diff --git a/.agents/skills/hf-bf16-gguf-conversion-jobs/SKILL.md b/.agents/skills/hf-bf16-gguf-conversion-jobs/SKILL.md new file mode 100644 index 0000000000..d8aa805100 --- /dev/null +++ b/.agents/skills/hf-bf16-gguf-conversion-jobs/SKILL.md @@ -0,0 +1,139 @@ +--- +name: hf-bf16-gguf-conversion-jobs +description: Use when converting Hugging Face SafeTensors checkpoints into split BF16 GGUF model repos with skippy-quantize on Hugging Face Jobs or a local machine, then publishing the artifact to Hugging Face. +metadata: + short-description: Convert HF checkpoints to BF16 GGUF repos +--- + +# HF BF16 GGUF Conversion Jobs + +Use this skill when the source artifact is a Hugging Face checkpoint repo and +the target artifact is a split BF16 GGUF model repo. The operational tool is +`skippy-quantize`; do not use `convert_hf_to_gguf.py`, `hf_to_gguf.py`, or a +wrapper that shells out to either script. Treat `hf_to_gguff.py` as the same +forbidden path if it appears in old notes or logs. + +## Preconditions + +- Confirm the source checkpoint repo, revision, tokenizer files, target repo, + output basename, expected split count, and desired split size before spending + HF Jobs credits. +- Build the standalone binary with `just skippy-quantize-standalone-release-build` + for local runs or in the job image/script for HF Jobs. +- Use `--output-type bf16` unless the experiment explicitly records a different + target precision. +- Prefer a split output with `--window-size 1` for first full-model runs. Raise + the window only after a smaller fixture proves the memory and I/O budget. +- Publish only complete windows, write per-window records, and resume from the + first missing target shard after cancellation. + +## Local Workflow + +Create a manifest: + +```bash +target/release/skippy-quantize init-convert \ + --source /path/to/checkpoint \ + --target /path/to/output-repo \ + --target-prefix BF16 \ + --output-basename -BF16 \ + --output-type bf16 \ + --expected-splits \ + --window-size 1 \ + --manifest /tmp/skippy-convert.json +``` + +Dry-run the next conversion window before spending I/O: + +```bash +target/release/skippy-quantize convert-job \ + --source /path/to/checkpoint \ + --target /path/to/output-repo \ + --target-prefix BF16 \ + --output-basename -BF16 \ + --output-type bf16 \ + --expected-splits \ + --window-size 1 \ + --manifest /tmp/skippy-convert.json \ + --max-memory 32G \ + --dry-run +``` + +Run until complete: + +```bash +target/release/skippy-quantize run-convert \ + --manifest /tmp/skippy-convert.json \ + --max-memory 32G \ + --split-max-size 50G \ + --stream-buffer-bytes 8388608 \ + --spool-dir /tmp/skippy-convert-output \ + --record-dir /tmp/skippy-convert-records \ + --json-event-file /tmp/skippy-convert-status.json \ + --json-event-interval-seconds 120 \ + --json-event-window 8 +``` + +Validate and publish: + +```bash +target/release/skippy-quantize verify-job \ + --manifest /tmp/skippy-convert.json \ + --json + +hf repo create / --type model --private +hf upload / /path/to/output-repo . --repo-type model +``` + +## HF Jobs Workflow + +Mount the source checkpoint and target model repo rather than downloading the +whole checkpoint into the job filesystem: + +```bash +hf jobs uv run \ + --namespace meshllm \ + --flavor cpu-upgrade \ + --timeout 3d \ + --secrets HF_TOKEN \ + --volume hf://models/:/mnt/checkpoint \ + --volume hf://models/:/mnt/target \ + --env SKIPPY_QUANTIZE_OUTPUT=json \ + --env PYTHONUNBUFFERED=1 \ + --detach \ + /path/to/skippy_convert_job.py \ + -- \ + --source /mnt/checkpoint \ + --target /mnt/target \ + --target-prefix BF16 \ + --output-basename -BF16 \ + --expected-splits \ + --split-max-size 50G \ + --max-memory 32G +``` + +The job script should only build or install `skippy-quantize`, create the +manifest if missing, run `run-convert`, verify the job, and upload sidecars. It +must not call the old Python converter. + +## Monitoring + +Use both HF Jobs status and `skippy-quantize` status: + +```bash +hf jobs inspect --namespace meshllm +hf jobs logs --namespace meshllm --tail 120 +target/release/skippy-quantize status --manifest /tmp/skippy-convert.json --json +``` + +For agents, prefer polling `/tmp/skippy-convert-status.json` over ingesting full +logs. Healthy snapshots show phase movement through `running`, `publishing`, +and `complete`, with only the last few high-level events retained. Stop and +diagnose if the same window restarts without a new published shard or memory +stays pinned near the hardware limit. + +## Record Keeping + +Record the job id, exact command, source revision, target repo commit, split +count, split size, memory budget, tokenizer notes, and follow-ups in the +experiment card or phase iteration card before promoting the artifact. diff --git a/.agents/skills/hf-bf16-gguf-conversion-jobs/agents/openai.yaml b/.agents/skills/hf-bf16-gguf-conversion-jobs/agents/openai.yaml new file mode 100644 index 0000000000..83d90d2d0d --- /dev/null +++ b/.agents/skills/hf-bf16-gguf-conversion-jobs/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "HF BF16 GGUF Conversion Jobs" + short_description: "Convert HF checkpoints to BF16 GGUF repos with skippy-quantize." + default_prompt: "Create or monitor a skippy-quantize BF16 GGUF conversion job." diff --git a/.agents/skills/hf-gguf-quant-jobs/SKILL.md b/.agents/skills/hf-gguf-quant-jobs/SKILL.md new file mode 100644 index 0000000000..39fabd3a14 --- /dev/null +++ b/.agents/skills/hf-gguf-quant-jobs/SKILL.md @@ -0,0 +1,202 @@ +--- +name: hf-gguf-quant-jobs +description: Use when creating, monitoring, validating, or documenting low-memory Hugging Face Jobs or local runs that quantize split BF16/FP16 GGUF model repos into custom quant GGUF repos with skippy-quantize. +--- + +# HF GGUF Quant Jobs + +Use this skill to turn an existing split BF16/FP16 GGUF model repo into a +quantized GGUF model repo without requiring the host to hold the full model in +memory or on local disk at once. The operational tool is `skippy-quantize`; do +not use `llama-quantize`, `llama-quantise`, or wrapper scripts that shell out to +those binaries. + +The supported pattern is: mount or point at the source BF16/FP16 GGUF repo, +quantize resumable split windows with `skippy-quantize`, publish completed +output shards to the target model repo, delete staged files immediately, and +resume from the first missing target shard after cancellation or failure. + +## Preconditions + +- Use a split BF16/FP16 GGUF repo as the source when possible. Do not re-read + SafeTensors for requants if a BF16 GGUF artifact already exists. +- Verify the source repo is complete before spending on quantization. Count all + expected split shards and refuse to run if any are missing. +- Use a tensor-type file for any custom recipe. Treat MTP tensors, output + tensors, precision-sensitive tensors, and latency-sensitive layer ranges as + explicit recipe inputs. +- Run jobs under the intended HF org and pass `HF_TOKEN` as a secret, not a + printed environment variable. +- Prefer mounted Hub repos over full `hf download` when the job only needs to + stream or stage one shard/window at a time. +- Build the standalone binary with `just skippy-quantize-standalone-release-build` + for local runs or in the job image/script for HF Jobs. + +## Workflow + +1. Identify the source BF16/FP16 GGUF repo, target quant repo, output prefix, + output basename, source prefix, quant type, tensor-type file, memory budget, + and split window size. +2. Preflight both Hub and mounted source paths with `skippy-quantize status`, + `next-window`, `validate-splits`, or a `quantize --preflight-only` run. Stop + if the source artifact is incomplete. +3. Write or upload a `quant-plan.json` with source repo/revision, target repo, + quant type, shard count, output prefix, tensor policy, and resume + settings. +4. Launch the job with `--window-size 1` for the first full model run unless a + smaller fixture proves a larger window is safe on the chosen hardware. +5. For each split window, stage only the required input shard, run + `skippy-quantize run-quant-window` or `run-quant`, publish finished shards, + then delete local staged input and output files. +6. Monitor for progress markers. A healthy job repeatedly emits staged source + copies, `quant_window`, publish completion, cleanup, and increasing split + progress. +7. Validate the target repo after completion by counting GGUF shards, checking + the first and last shard names, and confirming `quant-plan.json` plus the + tensor-type file are present. +8. Record the artifact in the experiment card and create an iteration card for + the run, including job id, command, environment, repo SHA, shard count, and + follow-up decisions. + +## Launch Template + +Create a quantization manifest: + +```bash +target/release/skippy-quantize init-quant \ + --source /mnt/source-gguf \ + --source-prefix \ + --target /mnt/target-quant \ + --target-prefix \ + --output-basename \ + --quant \ + --tensor-type-file /mnt/recipe/tensor-types.txt \ + --window-size 1 \ + --manifest /tmp/skippy-quantize.json +``` + +Dry-run the next quantization window before spending I/O: + +```bash +target/release/skippy-quantize quant-job \ + --source /mnt/source-gguf \ + --source-prefix \ + --target /mnt/target-quant \ + --target-prefix \ + --output-basename \ + --quant \ + --tensor-type-file /mnt/recipe/tensor-types.txt \ + --window-size 1 \ + --manifest /tmp/skippy-quantize.json \ + --backend llama-api \ + --max-memory 32G \ + --dry-run +``` + +Run until complete: + +```bash +target/release/skippy-quantize run-quant \ + --manifest /tmp/skippy-quantize.json \ + --backend llama-api \ + --max-memory 32G \ + --work-dir /tmp/skippy-quantize-work \ + --spool-dir /tmp/skippy-quantize-output \ + --record-dir /tmp/skippy-quantize-records \ + --json-event-file /tmp/skippy-quantize-status.json \ + --json-event-interval-seconds 120 \ + --json-event-window 8 +``` + +For HF Jobs, mount the BF16/FP16 source repo and target quant repo, then run the +same manifest and `run-quant` commands inside the job: + +```bash +hf jobs uv run \ + --namespace meshllm \ + --flavor cpu-upgrade \ + --timeout 3d \ + --secrets HF_TOKEN \ + --volume hf://models/:/mnt/source-gguf \ + --volume hf://models/:/mnt/target-quant \ + --env SKIPPY_QUANTIZE_OUTPUT=json \ + --env PYTHONUNBUFFERED=1 \ + --detach \ + /path/to/skippy_quant_job.py \ + -- \ + --source /mnt/source-gguf \ + --source-prefix \ + --target /mnt/target-quant \ + --target-prefix \ + --output-basename \ + --quant \ + --tensor-type-file /mnt/recipe/tensor-types.txt \ + --max-memory 32G +``` + +The job script should only build or install `skippy-quantize`, prepare the +manifest if missing, run `run-quant`, verify the job, and upload sidecars. + +## Monitoring + +Check status and logs: + +```bash +hf jobs inspect --namespace meshllm +hf jobs logs --namespace meshllm --tail 120 +``` + +For agents, prefer polling `/tmp/skippy-quantize-status.json` over ingesting +full logs. It is a periodically refreshed compact snapshot with the current +phase, current split window, and a bounded recent-event window. + +Useful healthy markers: + +- `Preflight QuantizeGguf with backend llama-api` +- `Source artifact is complete` +- `quant_window` +- `Published /mnt/target-quant/...` +- `Cleaned staged source` +- `split artifact ... 100.00%` + +Concerning markers: + +- repeated watchdog lines with no shard, tensor, upload, or cache-drop progress; +- cgroup memory pinned near the hardware limit; +- the same split window restarting repeatedly without new uploaded target files; +- fallback quant warnings for tensors that the recipe expected to preserve. + +If a job stalls, cancel it before changing code or hardware. The next run should +skip already published shards and resume at the first missing output shard. + +## Validation + +After completion, verify the target repo with an authenticated Hub API or CLI +check. Record at least: + +- target repo and commit SHA; +- privacy setting; +- total file count; +- GGUF shard count; +- first and last shard names; +- manifest/plan presence; +- tensor-type file presence. + +For local smoke tests, use a small split GGUF source first and verify: + +- `skippy-quantize verify-job --manifest --llama-load` succeeds; +- `skippy-quantize validate-splits --root --prefix ` succeeds; +- max RSS stays bounded compared with full-model size; +- `skippy-quantize status --manifest --json` reports completion. + +## Documentation Contract + +For Jianyang-style experiments, update both records: + +- the main experiment card with the promoted artifact; +- a phase iteration card with the job id, exact command, environment, + verification output, decision, and follow-ups. + +Keep post-experiment upstream notes separate from the run decision. The job can +be successful while the converter or quantizer patches still need extraction +into clean upstream PRs. diff --git a/.agents/skills/hf-gguf-quant-jobs/agents/openai.yaml b/.agents/skills/hf-gguf-quant-jobs/agents/openai.yaml new file mode 100644 index 0000000000..34004f14f2 --- /dev/null +++ b/.agents/skills/hf-gguf-quant-jobs/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "HF GGUF Quant Jobs" + short_description: "Run low-memory GGUF quantization jobs with skippy-quantize." + default_prompt: "Create or monitor a low-memory skippy-quantize GGUF quantization job." diff --git a/.agents/skills/hf-layer-package-jobs/SKILL.md b/.agents/skills/hf-layer-package-jobs/SKILL.md index b16ac9d990..2fad501e96 100644 --- a/.agents/skills/hf-layer-package-jobs/SKILL.md +++ b/.agents/skills/hf-layer-package-jobs/SKILL.md @@ -7,7 +7,11 @@ metadata: # HF Layer Package Jobs -Use this skill for the `models package` CLI, the `model-package` crate, and the daily Unsloth queue workflow. +Use this skill for the `models package` CLI, the `model-package` crate, and the +daily Unsloth queue workflow. This skill starts after a quantized GGUF artifact +exists. It does not quantize models; use `hf-gguf-quant-jobs` first or +`hf-quant-and-layer-package-jobs` when quantization and layer packaging should +run in one job. ## Workflow @@ -18,6 +22,59 @@ Use this skill for the `models package` CLI, the `model-package` crate, and the 5. The GitHub workflow should default to dry run. When confirmed, it should pass `--confirm`, submit at most the requested number of jobs, wait for every submitted HF Job, and fail if any job finishes unsuccessfully. 6. Prefer family-diverse candidate ordering after ranking by selected quant size, so one run does not consume the whole queue on a single model family. +## Commands + +Preview a package job: + +```bash +mesh-llm models package : --dry-run +``` + +Submit and follow: + +```bash +mesh-llm models package : --confirm --follow +``` + +Inspect jobs: + +```bash +mesh-llm models package --status +mesh-llm models package --logs +mesh-llm models package --list +``` + +For local package certification after the artifact exists: + +```bash +mesh-llm models certify --package-only --json +``` + +## Local Package Workflow + +When the quantized GGUF is already available on the local machine, build the +package locally with `skippy-model-package`, then publish the package directory +to a Hugging Face model repo: + +```bash +just build + +target/debug/skippy-model-package write-package \ + /: \ + --out-dir /tmp/-layers + +target/debug/skippy-model-package preflight \ + /tmp/-layers \ + --verify-sha256 + +hf repo create / --type model --private +hf upload / /tmp/-layers . --repo-type model +``` + +For local GGUF paths outside the Hugging Face cache, include explicit provenance +flags on `write-package`: `--model-id`, `--source-repo`, `--source-revision`, +and `--source-file`. + ## Validation Run Rust formatting and the focused package checks before committing: diff --git a/.agents/skills/hf-quant-and-layer-package-jobs/SKILL.md b/.agents/skills/hf-quant-and-layer-package-jobs/SKILL.md new file mode 100644 index 0000000000..7b62844930 --- /dev/null +++ b/.agents/skills/hf-quant-and-layer-package-jobs/SKILL.md @@ -0,0 +1,164 @@ +--- +name: hf-quant-and-layer-package-jobs +description: Use when running quantization of a BF16/FP16 GGUF repo and Skippy layer-package creation as one local or Hugging Face Jobs workflow, publishing both artifacts to Hugging Face. +metadata: + short-description: Quantize and package in one workflow +--- + +# HF Quant And Layer Package Jobs + +Use this skill when a workflow should produce both a quantized GGUF repo and a +Skippy layer package from an existing BF16/FP16 GGUF repo. The quantization +phase must use `skippy-quantize`; do not use `llama-quantize`, +`llama-quantise`, `convert_hf_to_gguf.py`, `hf_to_gguf.py`, or the misspelled +old notes form `hf_to_gguff.py`. + +## Preconditions + +- Source BF16/FP16 GGUF repo is complete and has a known selector/prefix. +- Target quant repo, quant selector, tensor-type file, output basename, expected + split count, and memory budget are known. +- Target layer-package repo is known or intentionally auto-derived by + `mesh-llm models package`. +- The layer package phase starts only after `skippy-quantize verify-job` + succeeds for the quantized artifact. + +## Local Workflow + +Quantize first: + +```bash +target/release/skippy-quantize init-quant \ + --source /mnt/bf16 \ + --source-prefix BF16 \ + --target /mnt/quant \ + --target-prefix \ + --output-basename - \ + --quant \ + --tensor-type-file /mnt/recipe/tensor-types.txt \ + --window-size 1 \ + --manifest /tmp/skippy-quantize.json + +target/release/skippy-quantize run-quant \ + --manifest /tmp/skippy-quantize.json \ + --backend llama-api \ + --max-memory 32G \ + --work-dir /tmp/skippy-quantize-work \ + --spool-dir /tmp/skippy-quantize-output \ + --record-dir /tmp/skippy-quantize-records \ + --json-event-file /tmp/skippy-quantize-status.json \ + --json-event-interval-seconds 120 \ + --json-event-window 8 + +target/release/skippy-quantize verify-job \ + --manifest /tmp/skippy-quantize.json \ + --llama-load +``` + +Before the real run, dry-run the same quant job and confirm it reports the +expected source, target, tensor recipe, backend, memory budget, and next window: + +```bash +target/release/skippy-quantize quant-job \ + --source /mnt/bf16 \ + --source-prefix BF16 \ + --target /mnt/quant \ + --target-prefix \ + --output-basename - \ + --quant \ + --tensor-type-file /mnt/recipe/tensor-types.txt \ + --window-size 1 \ + --manifest /tmp/skippy-quantize.json \ + --backend llama-api \ + --max-memory 32G \ + --dry-run +``` + +Publish the quant repo if the target is not already a mounted Hub repo: + +```bash +hf repo create / --type model --private +hf upload / /mnt/quant . --repo-type model +``` + +Package the published quant: + +```bash +mesh-llm models package /: --dry-run +mesh-llm models package /: --confirm --follow +``` + +Or package locally and publish: + +```bash +target/debug/skippy-model-package write-package \ + /: \ + --out-dir /tmp/-layers + +target/debug/skippy-model-package preflight \ + /tmp/-layers \ + --verify-sha256 + +hf repo create / --type model --private +hf upload / /tmp/-layers . --repo-type model +``` + +## HF Jobs Workflow + +When combining both phases in one HF Job, keep the quantized GGUF repo as the +durable boundary: + +1. Mount the BF16/FP16 source repo read-only. +2. Mount the target quant repo read/write. +3. Run `skippy-quantize init-quant` if the manifest is missing. +4. Run `skippy-quantize run-quant` until complete. +5. Run `skippy-quantize verify-job`; stop if it fails. +6. Submit or run the `mesh-llm models package :` package + phase. +7. Record both the quant repo commit and the layer-package repo commit. + +Template: + +```bash +hf jobs uv run \ + --namespace meshllm \ + --flavor cpu-upgrade \ + --timeout 4d \ + --secrets HF_TOKEN \ + --volume hf://models/:/mnt/bf16 \ + --volume hf://models/:/mnt/quant \ + --env SKIPPY_QUANTIZE_OUTPUT=json \ + --env PYTHONUNBUFFERED=1 \ + --detach \ + /path/to/skippy_quant_then_package_job.py \ + -- \ + --source /mnt/bf16 \ + --source-prefix BF16 \ + --target /mnt/quant \ + --target-prefix \ + --output-basename - \ + --quant \ + --tensor-type-file /mnt/recipe/tensor-types.txt \ + --package-ref /: \ + --max-memory 32G +``` + +## Resume Rules + +- If quant shards already exist, `skippy-quantize` resumes at the first missing + shard. +- If the quant repo verifies successfully, skip quantization and run or inspect + the package job. +- Do not delete a verified quant repo to force a clean package run. Package jobs + should consume the published quant artifact as the source of truth. + +## Validation + +Before promoting the combined run, record: + +- source BF16/FP16 repo revision; +- quant repo commit, quant selector, tensor recipe, split count, and verify + output; +- layer-package job id, target repo, target commit, and package certification; +- total HF job cost and whether the combined workflow saved time or only saved + operator steps. diff --git a/.agents/skills/hf-quant-and-layer-package-jobs/agents/openai.yaml b/.agents/skills/hf-quant-and-layer-package-jobs/agents/openai.yaml new file mode 100644 index 0000000000..64b9f9293d --- /dev/null +++ b/.agents/skills/hf-quant-and-layer-package-jobs/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "HF Quant And Layer Package Jobs" + short_description: "Run quantization and layer packaging as one workflow." + default_prompt: "Create or monitor a skippy-quantize quantization plus layer-package workflow." diff --git a/Cargo.lock b/Cargo.lock index 99d216f61c..bea2e65d01 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3616,6 +3616,13 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" +[[package]] +name = "llama-quant-ffi" +version = "0.68.0" +dependencies = [ + "libloading", +] + [[package]] name = "llama-spec-bench" version = "0.68.0" @@ -7409,6 +7416,7 @@ dependencies = [ "model-artifact", "model-hf", "model-ref", + "reqwest 0.12.28", "serde", "serde_json", "sha2 0.10.9", @@ -7472,6 +7480,19 @@ dependencies = [ "serde", ] +[[package]] +name = "skippy-quantize" +version = "0.68.0" +dependencies = [ + "anyhow", + "clap", + "libc", + "llama-quant-ffi", + "serde", + "serde_json", + "skippy-ffi", +] + [[package]] name = "skippy-runtime" version = "0.68.0" diff --git a/Cargo.toml b/Cargo.toml index d85880df07..4e34746bf9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -47,12 +47,14 @@ members = [ "crates/skippy-server", "crates/metrics-server", "crates/skippy-model-package", + "crates/skippy-quantize", "crates/model-package", "crates/skippy-correctness", "crates/llama-spec-bench", "crates/skippy-bench", "crates/skippy-prompt", "crates/mesh-mixture-of-agents", + "crates/llama-quant-ffi", "tools/xtask", ] default-members = [ diff --git a/Justfile b/Justfile index 3f729133cd..daaae7a301 100644 --- a/Justfile +++ b/Justfile @@ -205,6 +205,34 @@ metrics-server-build: skippy-wan-lab-build-bins: cargo build --release --locked -p skippy-server -p skippy-prompt -p metrics-server -p skippy-model-package +# Build the resumable GGUF conversion/quantization replacement CLI. +[unix] +skippy-quantize-build: + just with-lld cargo build -p skippy-quantize + +[windows] +skippy-quantize-build: + @just with-lld cargo build -p skippy-quantize + +# Build the release binary used in HF conversion/quantization job images. +[unix] +skippy-quantize-release-build: + just with-lld cargo build --release --locked -p skippy-quantize + +[windows] +skippy-quantize-release-build: + @just with-lld cargo build --release --locked -p skippy-quantize + +# Build skippy-quantize as a standalone quantization binary with the pinned +# llama.cpp quantization ABI linked into the executable. +[unix] +skippy-quantize-standalone-build backend="cpu": + LLAMA_STAGE_BACKEND="{{ backend }}" LLAMA_STAGE_LINK_MODE=static just with-lld cargo build -p skippy-quantize + +[unix] +skippy-quantize-standalone-release-build backend="cpu": + LLAMA_STAGE_BACKEND="{{ backend }}" LLAMA_STAGE_LINK_MODE=static just with-lld cargo build --release --locked -p skippy-quantize + # Generate a reproducible benchmark corpus for skippy bench tooling. bench-corpus tier="smoke" *ARGS="": scripts/generate-bench-corpus.py "{{ tier }}" {{ ARGS }} diff --git a/SPD_SKIPPY_PROJECT.md b/SPD_SKIPPY_PROJECT.md new file mode 100644 index 0000000000..1d0d5e500f --- /dev/null +++ b/SPD_SKIPPY_PROJECT.md @@ -0,0 +1,434 @@ +# GLM 4.7 SPD-on-MTP Project Handoff + +This branch captures the clean GLM 4.7 SPD-on-MTP experiment path. It starts +from the current Skippy native-MTP verifier work and ports only the compact SPD +training, export, manifest, and latency-model pieces needed to train a GLM 4.7 +sidecar. + +The working hypothesis is that native GLM MTP is the verifier/correctness +foundation, while a trained SPD sidecar may be the better drafting oracle for +`N > 1` proposals. The branch should therefore be judged by whether a freshly +trained GLM 4.7 sidecar can propose enough verified future tokens to improve +decode tok/s against the golden vanilla no-MTP baseline. + +It intentionally excludes private lab hosts, credentials, local IPs, and +machine-specific notes. Use it as a research/implementation handoff that another +engineer can reproduce from open models, open data, and the checked-in scripts. + +PR #860 is the compact GLM donor branch. PR #859 is treated as proof +archaeology only; do not transplant its broad live-serving/protocol work until +offline GLM sidecar quality justifies that next step. + +## Source Paper + +Paper: **Speculative Pipeline Decoding: Higher-Accuracy and Zero-Bubble +Speculation via Pipeline Parallelism** + +- arXiv: `https://arxiv.org/abs/2605.30852` +- Reference code: `https://github.com/yuyijiong/speculative_pipeline_decoding` + +The core idea is to combine pipeline parallel target-model execution with a +trained speculation module. The target model is partitioned into `n` pipeline +stages. While the target pipeline is processing one token per stage, the SPD +head consumes selected intermediate hidden states from the pipeline and proposes +future draft token(s). The target model still verifies the draft tokens, so with +verification enabled the output follows the base model's decoding path. + +## Why This Matters for Skippy + +Skippy already splits a model across staged runtimes. Ordinary split decoding is +sensitive to stage and network latency because each generated token must traverse +the full stage chain before the next target token is known. + +SPD is interesting because it can fill the pipeline and amortize that stage/hop +latency across accepted speculative work. The quality question is whether the +sidecar head accepts enough tokens. The engineering question is whether Skippy +can expose the required hidden-state taps and verify proposals without breaking +target-model equivalence. + +Headline background result: the pretrained `Qwen/Qwen3.5-4B` SPD head accepted +`1230 / 1536` draft flags on the local reference eval, with equivalent accept +length `2.4704` and token-weighted theoretical gain `163.39%`. Feeding that +same real trace into a four-stage Skippy latency model with `4ms` per stage +estimated `9.882x` SPD-vs-serial split throughput at `0ms` hop, `8.117x` at +`10ms` hop, and `7.752x` at `25ms` hop. + +That Qwen result is motivation, not the target artifact. This branch's target +artifact is a newly trained GLM 4.7 SPD sidecar exported as a Skippy-readable +manifest plus `spd-head.safetensors`. + +## What Works Today + +### 1. Real Small-Model Training Proof + +`evals/spd/hf_train_eval_qwen06.py` trains a real SPD head using the paper's +reference code. + +Recorded local proof: + +| Model | Training data | Head | Generated tokens | Accepted flags | Acceptance | Equivalent accept length | Theoretical gain | +| --- | --- | --- | ---: | ---: | ---: | ---: | ---: | +| `Qwen/Qwen3-0.6B` | `HuggingFaceH4/ultrachat_200k`, split `train_sft`, first 1024 rows | 4 spec layers, 2 stages | 1536 | 326 / 1536 | 0.5628 | 1.1257 | 12.67% | + +This proves the train/eval/export path. It is not the high-gain target. + +### 2. Strong Modest-Model Acceptance Signal + +The author-published `Qwen3.5-4B_s4_l4.pt` SPD head evaluates well with the +reference verifier. + +Recorded local proof: + +| Model | Head | Generated tokens | Accepted flags | Acceptance | Equivalent accept length | Theoretical gain | +| --- | --- | ---: | ---: | ---: | ---: | ---: | +| `Qwen/Qwen3.5-4B` | pretrained, 4 stages / 4 spec layers | 1536 | 1230 / 1536 | 0.6176 | 2.4704 | 163.39% | + +Per-dataset theoretical gains from the same run: + +| Dataset | Acceptance | Equivalent accept length | Theoretical gain | +| --- | ---: | ---: | ---: | +| MT-Bench | 0.4918 | 1.9673 | 98.42% | +| HumanEval | 0.8797 | 3.5189 | 254.18% | +| GSM8K | 0.5926 | 2.3704 | 137.58% | + +This is the main reason to keep pursuing SPD for Skippy. + +### 3. Trace-Based Skippy Latency Model + +`evals/spd/simulate_latency.py` consumes real per-sample SPD eval traces and +models split-stage latency. It does not invent acceptance. + +Recorded Qwen3.5-4B trace with four target stages at `4ms,4ms,4ms,4ms`: + +| Hop ms | Serial split tok/s | SPD pipeline tok/s | SPD vs serial split | Paper-like gain | P50 serial ms | P50 SPD ms | +| ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 0 | 62.50 | 617.61 | 9.882x | 2.470x | 1024.00 | 106.50 | +| 1 | 52.63 | 494.09 | 9.388x | 2.470x | 1216.00 | 133.12 | +| 5 | 32.26 | 274.49 | 8.509x | 2.470x | 1984.00 | 239.62 | +| 10 | 21.74 | 176.46 | 8.117x | 2.470x | 2944.00 | 372.75 | +| 25 | 10.99 | 85.19 | 7.752x | 2.470x | 5824.00 | 772.12 | + +The paper-like gain is from the SPD trace itself. The Skippy comparison models +ordinary split serving as requiring each generated token to traverse all +stages/hops before the next target token is known. + +### 4. Rust Serving Artifact Validation + +`crates/skippy-runtime/src/spd.rs` adds a manifest parser and validator for SPD +head artifacts: + +- schema: `skippy-spd-head/v1` +- checkpoint path, byte size, sha256 +- base model path/id +- checkpoint format/version +- hidden size +- vocab size +- draft vocab size and optional draft token ids +- number of target stages +- number of spec layers +- shallow hidden-layer tap indices +- optional safetensors serving checkpoint path, size, checksum, tensor count, + and dtype + +`evals/spd/export_spd_head.py` exports the reference `.pt` checkpoint into +`spd-head.safetensors` and updates the manifest with a `serving_checkpoint` +section. This is still validation only: Skippy can inspect the serving artifact +but does not run the SPD head yet. + +## What Does Not Work Yet + +- Skippy/Rust does not execute the SPD head. +- Skippy/Rust does not load tensor values into an executable SPD head yet. +- Skippy does not yet expose live hidden-state taps for SPD. +- No live Skippy request has used trained SPD proposals. +- No larger-than-4B head has been trained by us yet. + +## Correctness Contract + +SPD should be treated as a verified speculative path. + +For greedy decoding: + +1. SPD proposes a token. +2. The target model computes the verified logits for that position. +3. The token is accepted only if it equals the target argmax. +4. On rejection, Skippy rolls back speculative state and emits the target token. + +For sampling: + +1. SPD proposes from draft distribution `q`. +2. Target distribution `p` is computed by the base model. +3. Standard speculative rejection sampling accepts with the corrected + probability and otherwise samples from residual `max(0, p - q)`. + +Do not ship unverified/lossy SPD as the default path. Lossy SPD should only be a +separate explicit experiment because wrong accepted tokens change the future +context. + +## Practical Skippy Hosting Model + +Treat the SPD head as a sidecar artifact attached to a Skippy stage topology. +It should not be exposed as a separate OpenAI model and should not mutate the +base GGUF/layer-package weights. + +Recommended first implementation: + +- one SPD sidecar runtime per active Skippy topology/session group +- host it in one Skippy process first, likely coordinator or final-stage side +- other stages expose/send selected hidden-state taps +- SPD proposes token candidates +- normal Skippy stages verify every emitted token + +Distributed SPD execution across all stage nodes may become useful later, but it +is not the first proof path. + +## llama.cpp / Stage Runtime Dependencies + +The current proof branch does not require James's GLM/MTP work to reproduce the +Python SPD results or validate the Rust manifest. The live Skippy path will, +however, need additional staged-runtime/llama-side capability. + +Likely required: + +- hidden-state tap export from selected layers/stages, with token position, + dtype, shape, and stage ownership metadata +- enough sideband transport to return those taps to the SPD sidecar without + changing ordinary generation output +- verification support that can run proposed SPD tokens through the real target + stages and return the target-model decision +- rollback/session-trim support for rejected speculative tokens +- ABI version bumps and Rust `skippy-ffi` mirrors for any new staged-runtime + calls + +Adjacent work that may help: + +- native MTP verification work has similar concerns around speculative proposal, + target verification, sideband data, and rollback +- GLM/MTP branches may contain useful patterns for verifier plumbing, but they + are not SPD themselves and should not be merged wholesale just to start SPD +- package-declared draft speculation work may be useful later for advertising + optional SPD artifacts in model/layer packages + +First Skippy implementation should add the minimum SPD-specific stage-runtime +surface needed for Qwen3.5-4B parity: capture required hidden taps, run the SPD +head, and verify proposed tokens. Pull reusable verifier/rollback patterns from +MTP work only after confirming they apply cleanly to SPD. + +## Artifact Layout Target + +Layer packages should eventually support optional SPD artifacts: + +```text +package/ + manifest.json + parts/ + ... + spd/ + skippy-spd-head.json + spd-head.safetensors + draft-vocab.json +``` + +The manifest must bind the head to: + +- base model digest/path +- tokenizer/vocab identity +- hidden size +- split topology +- stage count +- layer taps +- draft vocab +- head tensor checksum + +## Reproduction Commands + +### Train Small Qwen3-0.6B Head + +```bash +python evals/spd/hf_train_eval_qwen06.py \ + --work-dir /tmp/skippy-spd-qwen06-proof \ + --model-name Qwen/Qwen3-0.6B \ + --dataset HuggingFaceH4/ultrachat_200k \ + --dataset-split train_sft \ + --train-rows 1024 \ + --eval-rows-per-set 8 \ + --num-stages 2 \ + --num-spec-layers 4 \ + --max-length 256 \ + --max-new-tokens 64 \ + --draft-top-k 4 \ + --device mps \ + --upload-repo '' +``` + +Use `--device cuda` on a CUDA host. + +### Evaluate Pretrained Qwen3.5-4B Head + +```bash +python evals/spd/hf_train_eval_qwen06.py \ + --work-dir /tmp/skippy-spd-qwen35-4b-pretrained-s4l4 \ + --model-name Qwen/Qwen3.5-4B \ + --spec-head-repo yuyijiong/speculative_pipeline_decoding \ + --spec-head-file Qwen3.5-4B_s4_l4.pt \ + --manifest-base-model-path Qwen/Qwen3.5-4B \ + --skip-train \ + --device mps \ + --eval-rows-per-set 8 \ + --max-new-tokens 64 \ + --draft-top-k 4 \ + --upload-repo '' +``` + +### Export Serving Checkpoint + +```bash +python evals/spd/export_spd_head.py \ + --checkpoint /tmp/skippy-spd-qwen35-4b-pretrained-s4l4/artifacts//train/speculation_head_final.pt \ + --manifest /tmp/skippy-spd-qwen35-4b-pretrained-s4l4/artifacts//train/skippy-spd-head.json \ + --base-model-path Qwen/Qwen3.5-4B +``` + +### Simulate Split Latency From Trace + +```bash +python evals/spd/simulate_latency.py \ + --raw /tmp/skippy-spd-qwen35-4b-pretrained-s4l4/artifacts//eval/raw/pipeline_eval__train__speculation_head_final__nt24__per_sample.jsonl \ + --stage-ms 4,4,4,4 \ + --hop-ms 0,1,5,10,25 +``` + +## Engineering Next Steps + +### Milestone 1: Tensor Export + +Goal: produce a Rust-serving artifact from the `.pt` checkpoint. + +Tasks: + +1. Add/export a `safetensors` writer for the SPD checkpoint. Done in + `evals/spd/export_spd_head.py`. +2. Preserve tensor names, shapes, dtype, draft vocab ids, and config. Done via + the safetensors file and `skippy-spd-head.json`. +3. Extend `skippy-spd-head.json` to reference the serving checkpoint. Done with + the optional `serving_checkpoint` section. +4. Add a small shape/checksum inspection command or test fixture. Done in + `skippy-runtime` tests. + +Exit criteria: + +- Qwen3.5-4B SPD head exports deterministically. +- Rust can validate the manifest and enumerate expected tensors. +- To validate a local exported head through Rust, set + `SKIPPY_SPD_MANIFEST=/tmp/.../train/skippy-spd-head.json` and run: + +```bash +cargo test -p skippy-runtime validates_external_manifest_when_skippy_spd_manifest_is_set +``` + +### Milestone 2: Rust Forward Pass Parity + +Goal: Rust computes the same draft candidates as Python for recorded inputs. + +Tasks: + +1. Record hidden-state tap fixtures from Python reference execution. +2. Implement the Qwen3.5-4B SPD head forward pass in Rust. +3. Compare Rust top-k draft candidates against Python top-k on the same hidden + states. +4. Add focused tests with small fixture tensors. + +Exit criteria: + +- Rust top-k proposals match Python within tolerance on recorded fixtures. +- No Skippy serving integration is required for this milestone. + +### Milestone 3: Skippy Hidden-State Taps + +Goal: Skippy can expose the hidden states the SPD head needs. + +Tasks: + +1. Identify the target layer taps from `skippy-spd-head.json`. +2. Add a hidden-state sideband/tap path in the staged runtime. +3. Validate dtype, shape, token position, and stage ownership. +4. Write a correctness test that compares tapped hidden states against a known + reference for a small prompt. + +Exit criteria: + +- Skippy can capture the required taps for a live prompt without changing + normal generation output. + +### Milestone 4: Live Verified SPD in Skippy + +Goal: Skippy uses SPD proposals during generation and verifies every token. + +Tasks: + +1. Wire SPD proposal generation into `skippy-server`. +2. Feed proposals into the existing target verification path. +3. Roll back speculative KV/session state on rejection. +4. Emit metrics for proposals, accepted tokens, rejected tokens, equivalent + accept length, and decode-loop steps. +5. Run ordinary split serving and SPD serving against the same prompts. + +Exit criteria: + +- Greedy outputs match ordinary target-model decoding. +- Acceptance and equivalent accept length are non-zero and close to reference + trace behavior. +- Latency improves under injected hop/stage delay. + +### Milestone 5: Larger-Model Training Proof + +Goal: prove the SPD head generation pipeline scales beyond the pretrained 4B +artifact. + +Recommended first target: + +- a larger Qwen-family model with architecture/tokenizer support close to the + reference implementation +- avoid custom huge MoE targets for the first scaling proof + +Tasks: + +1. Train with open conversation data mix. +2. Keep draft vocab capped at 32k or 50k. +3. Evaluate on the same MT-Bench/HumanEval/GSM8K prompt sets. +4. Record acceptance, equivalent accept length, and latency simulation. +5. Publish only artifact manifests, scripts, and aggregate metrics unless model + licensing allows the trained head to be shared. + +Exit criteria: + +- Larger-model head reaches useful equivalent accept length. +- Training process and artifact production are reproducible by another engineer. + +## Branch Scope + +This branch should stay focused on SPD proof and handoff material: + +- `SPD_SKIPPY_PROJECT.md` +- `evals/spd/` +- `crates/skippy-runtime/src/spd.rs` +- minimal module export from `skippy-runtime` + +Avoid mixing in unrelated MTP, GLM, packaging, branch-reconciliation, or private +lab automation work. Those can be inputs later, but the purpose of this branch +is to make the SPD path clear and reproducible. + +## Validation + +Run: + +```bash +python3 -m py_compile evals/spd/hf_train_eval_qwen06.py evals/spd/simulate_latency.py evals/spd/export_spd_head.py +cargo fmt --all -- --check +cargo test -p skippy-runtime spd +cargo clippy -p skippy-runtime --all-targets -- -D warnings +``` + +Before publishing or handing off, run the repo's normal secret scan and also +check the diff for private hostnames, private IPs, access tokens, credentials, +and absolute developer-machine paths. diff --git a/crates/llama-quant-ffi/Cargo.toml b/crates/llama-quant-ffi/Cargo.toml new file mode 100644 index 0000000000..97e30cda77 --- /dev/null +++ b/crates/llama-quant-ffi/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "llama-quant-ffi" +edition.workspace = true +license.workspace = true +version.workspace = true +build = "build.rs" +description = "Small Rust FFI surface for llama.cpp GGUF quantization" +repository = "https://github.com/Mesh-LLM/mesh-llm" +homepage = "https://github.com/Mesh-LLM/mesh-llm" + +[features] +default = [] +dynamic-runtime = ["dep:libloading"] + +[dependencies] +libloading = { version = "0.8", optional = true } diff --git a/crates/llama-quant-ffi/build.rs b/crates/llama-quant-ffi/build.rs new file mode 100644 index 0000000000..a0ee0ead01 --- /dev/null +++ b/crates/llama-quant-ffi/build.rs @@ -0,0 +1,412 @@ +fn main() { + print_rerun_envs(); + + if std::env::var_os("CARGO_FEATURE_DYNAMIC_RUNTIME").is_some() { + return; + } + + let link_mode = + std::env::var("LLAMA_STAGE_LINK_MODE").or_else(|_| std::env::var("SKIPPY_LLAMA_LINK_MODE")); + if link_mode.as_deref() == Ok("dynamic") { + link_dynamic_runtime(); + return; + } + + let workspace_root = workspace_root(); + let target = std::env::var("TARGET").unwrap_or_default(); + let backend = std::env::var("LLAMA_STAGE_BACKEND") + .or_else(|_| std::env::var("SKIPPY_LLAMA_BACKEND")) + .unwrap_or_else(|_| default_backend(&target).to_string()); + let build_dir = configured_build_dir(&workspace_root, &backend); + ensure_static_native_ready(&workspace_root, &build_dir, &target, &backend); + emit_static_link(&build_dir, &target); +} + +fn print_rerun_envs() { + for key in [ + "LLAMA_STAGE_BUILD_DIR", + "LLAMA_STAGE_LIB_DIR", + "LLAMA_STAGE_LINK_MODE", + "SKIPPY_LLAMA_BUILD_DIR", + "SKIPPY_LLAMA_LIB_DIR", + "SKIPPY_LLAMA_LINK_MODE", + "LLAMA_STAGE_BACKEND", + "SKIPPY_LLAMA_BACKEND", + "SKIPPY_LLAMA_AUTO_BUILD", + "MESH_LLM_AUTO_BUILD_LLAMA", + "CUDA_PATH", + "HIP_PATH", + "ROCM_PATH", + "LLVMInstallDir", + "VULKAN_SDK", + ] { + println!("cargo:rerun-if-env-changed={key}"); + } +} + +fn link_dynamic_runtime() { + if let Ok(lib_dir) = + std::env::var("LLAMA_STAGE_LIB_DIR").or_else(|_| std::env::var("SKIPPY_LLAMA_LIB_DIR")) + { + println!("cargo:rustc-link-search=native={lib_dir}"); + } + println!("cargo:rustc-link-lib=dylib=llama-common"); + println!("cargo:rustc-link-lib=dylib=llama"); +} + +fn workspace_root() -> std::path::PathBuf { + std::path::PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR")) + .join("../..") +} + +fn configured_build_dir(workspace_root: &std::path::Path, backend: &str) -> std::path::PathBuf { + std::env::var("LLAMA_STAGE_BUILD_DIR") + .or_else(|_| std::env::var("SKIPPY_LLAMA_BUILD_DIR")) + .map(std::path::PathBuf::from) + .map(|path| { + if path.is_absolute() { + path + } else { + workspace_root.join(path) + } + }) + .unwrap_or_else(|_| { + workspace_root.join(format!( + ".deps/llama-build/build-stage-abi-static-{backend}" + )) + }) +} + +fn default_backend(target: &str) -> &'static str { + if target.contains("apple") { + "metal" + } else { + "cpu" + } +} + +fn ensure_static_native_ready( + workspace_root: &std::path::Path, + build_dir: &std::path::Path, + target: &str, + backend: &str, +) { + if required_static_archives_exist(build_dir) { + return; + } + if !native_auto_build_enabled() { + panic!( + "patched llama.cpp quant archives are missing from {}; run `just llama-build`, set LLAMA_STAGE_BUILD_DIR, or enable SKIPPY_LLAMA_AUTO_BUILD=1", + build_dir.display() + ); + } + if target.contains("windows") { + panic!( + "patched llama.cpp quant archives are missing from {}; automatic native preparation is not supported for Windows from build.rs yet", + build_dir.display() + ); + } + + let prepare = workspace_root.join("scripts/prepare-llama.sh"); + let build = workspace_root.join("scripts/build-llama.sh"); + println!("cargo:rerun-if-changed={}", prepare.display()); + println!("cargo:rerun-if-changed={}", build.display()); + if !prepare.exists() || !build.exists() { + panic!( + "patched llama.cpp quant archives are missing from {}, and mesh-llm build scripts were not found under {}", + build_dir.display(), + workspace_root.display() + ); + } + + run_native_script( + workspace_root, + &prepare, + ["pinned"].as_slice(), + backend, + build_dir, + ); + run_native_script(workspace_root, &build, [].as_slice(), backend, build_dir); + + if !required_static_archives_exist(build_dir) { + panic!( + "patched llama.cpp quant build finished but required archives are still missing from {}", + build_dir.display() + ); + } +} + +fn native_auto_build_enabled() -> bool { + for key in ["SKIPPY_LLAMA_AUTO_BUILD", "MESH_LLM_AUTO_BUILD_LLAMA"] { + if let Ok(value) = std::env::var(key) { + return !matches!( + value.to_ascii_lowercase().as_str(), + "0" | "false" | "no" | "off" + ); + } + } + true +} + +fn run_native_script( + workspace_root: &std::path::Path, + script: &std::path::Path, + args: &[&str], + backend: &str, + build_dir: &std::path::Path, +) { + let mut command = std::process::Command::new("bash"); + command.current_dir(workspace_root).arg(script).args(args); + command.env("LLAMA_WORKDIR", workspace_root.join(".deps/llama.cpp")); + command.env("LLAMA_BUILD_DIR", build_dir); + command.env("LLAMA_STAGE_BUILD_DIR", build_dir); + command.env("LLAMA_STAGE_LINK_MODE", "static"); + command.env("LLAMA_STAGE_BACKEND", backend); + let status = command.status().unwrap_or_else(|error| { + panic!("failed to run {}: {error}", script.display()); + }); + if !status.success() { + panic!("{} failed with status {status}", script.display()); + } +} + +fn emit_static_link(build_dir: &std::path::Path, target: &str) { + for dir in static_search_dirs(build_dir) + .iter() + .filter(|dir| dir.exists()) + { + println!("cargo:rustc-link-search=native={}", dir.display()); + } + let cmake_cache = build_dir.join("CMakeCache.txt"); + if cmake_cache.exists() { + println!("cargo:rerun-if-changed={}", cmake_cache.display()); + } + emit_archive_reruns(build_dir); + + println!("cargo:rustc-link-lib=static=llama-common"); + println!("cargo:rustc-link-lib=static=llama-common-base"); + println!("cargo:rustc-link-lib=static=llama"); + println!("cargo:rustc-link-lib=static=ggml"); + let has_cuda = static_archive_exists( + build_dir, + "ggml/src/ggml-cuda/libggml-cuda.a", + "ggml/src/ggml-cuda/ggml-cuda.lib", + ); + if has_cuda { + println!("cargo:rustc-link-lib=static=ggml-cuda"); + } + let has_hip = static_archive_exists( + build_dir, + "ggml/src/ggml-hip/libggml-hip.a", + "ggml/src/ggml-hip/ggml-hip.lib", + ); + if has_hip { + println!("cargo:rustc-link-lib=static=ggml-hip"); + } + let has_vulkan = static_archive_exists( + build_dir, + "ggml/src/ggml-vulkan/libggml-vulkan.a", + "ggml/src/ggml-vulkan/ggml-vulkan.lib", + ); + if has_vulkan { + println!("cargo:rustc-link-lib=static=ggml-vulkan"); + } + println!("cargo:rustc-link-lib=static=ggml-cpu"); + if static_archive_exists( + build_dir, + "ggml/src/ggml-blas/libggml-blas.a", + "ggml/src/ggml-blas/ggml-blas.lib", + ) { + println!("cargo:rustc-link-lib=static=ggml-blas"); + } + if static_archive_exists( + build_dir, + "ggml/src/ggml-metal/libggml-metal.a", + "ggml/src/ggml-metal/ggml-metal.lib", + ) { + println!("cargo:rustc-link-lib=static=ggml-metal"); + } + println!("cargo:rustc-link-lib=static=ggml-base"); + emit_system_links( + build_dir, + &cmake_cache, + target, + has_cuda, + has_hip, + has_vulkan, + ); +} + +fn static_search_dirs(build_dir: &std::path::Path) -> [std::path::PathBuf; 9] { + [ + build_dir.join("common"), + build_dir.join("src"), + build_dir.join("ggml/src"), + build_dir.join("ggml/src/ggml-cpu"), + build_dir.join("ggml/src/ggml-blas"), + build_dir.join("ggml/src/ggml-cuda"), + build_dir.join("ggml/src/ggml-hip"), + build_dir.join("ggml/src/ggml-metal"), + build_dir.join("ggml/src/ggml-vulkan"), + ] +} + +fn emit_archive_reruns(build_dir: &std::path::Path) { + for (unix_archive, msvc_archive) in [ + ("src/libllama.a", "src/llama.lib"), + ("common/libllama-common.a", "common/llama-common.lib"), + ( + "common/libllama-common-base.a", + "common/llama-common-base.lib", + ), + ("ggml/src/libggml.a", "ggml/src/ggml.lib"), + ("ggml/src/libggml-base.a", "ggml/src/ggml-base.lib"), + ( + "ggml/src/ggml-cpu/libggml-cpu.a", + "ggml/src/ggml-cpu/ggml-cpu.lib", + ), + ("ggml/src/libggml-cpu.a", "ggml/src/ggml-cpu.lib"), + ( + "ggml/src/ggml-blas/libggml-blas.a", + "ggml/src/ggml-blas/ggml-blas.lib", + ), + ( + "ggml/src/ggml-cuda/libggml-cuda.a", + "ggml/src/ggml-cuda/ggml-cuda.lib", + ), + ( + "ggml/src/ggml-hip/libggml-hip.a", + "ggml/src/ggml-hip/ggml-hip.lib", + ), + ( + "ggml/src/ggml-metal/libggml-metal.a", + "ggml/src/ggml-metal/ggml-metal.lib", + ), + ( + "ggml/src/ggml-vulkan/libggml-vulkan.a", + "ggml/src/ggml-vulkan/ggml-vulkan.lib", + ), + ] { + for archive in [unix_archive, msvc_archive] + .iter() + .map(|path| build_dir.join(path)) + .filter(|archive| archive.exists()) + { + println!("cargo:rerun-if-changed={}", archive.display()); + } + } +} + +fn emit_system_links( + build_dir: &std::path::Path, + cmake_cache: &std::path::Path, + target: &str, + has_cuda: bool, + has_hip: bool, + has_vulkan: bool, +) { + if target.contains("apple") { + println!("cargo:rustc-link-lib=c++"); + println!("cargo:rustc-link-lib=framework=Accelerate"); + if static_archive_exists( + build_dir, + "ggml/src/ggml-metal/libggml-metal.a", + "ggml/src/ggml-metal/ggml-metal.lib", + ) { + println!("cargo:rustc-link-lib=framework=Foundation"); + println!("cargo:rustc-link-lib=framework=Metal"); + println!("cargo:rustc-link-lib=framework=MetalKit"); + } + } else if target.contains("linux") { + println!("cargo:rustc-link-lib=stdc++"); + println!("cargo:rustc-link-lib=dylib=m"); + println!("cargo:rustc-link-lib=dylib=dl"); + println!("cargo:rustc-link-lib=dylib=pthread"); + if has_cuda { + link_linux_cuda_libs(cmake_cache); + } + if has_hip { + link_linux_hip_libs(); + } + if has_vulkan { + println!("cargo:rustc-link-lib=dylib=vulkan"); + } + } +} + +fn required_static_archives_exist(build_dir: &std::path::Path) -> bool { + [ + &["src/libllama.a", "src/llama.lib"][..], + &["common/libllama-common.a", "common/llama-common.lib"], + &[ + "common/libllama-common-base.a", + "common/llama-common-base.lib", + ], + &["ggml/src/libggml.a", "ggml/src/ggml.lib"], + &["ggml/src/libggml-base.a", "ggml/src/ggml-base.lib"], + &[ + "ggml/src/libggml-cpu.a", + "ggml/src/ggml-cpu.lib", + "ggml/src/ggml-cpu/libggml-cpu.a", + "ggml/src/ggml-cpu/ggml-cpu.lib", + ], + ] + .iter() + .all(|candidates| { + candidates + .iter() + .any(|candidate| build_dir.join(candidate).exists()) + }) +} + +fn static_archive_exists( + build_dir: &std::path::Path, + unix_archive: &str, + msvc_archive: &str, +) -> bool { + build_dir.join(unix_archive).exists() || build_dir.join(msvc_archive).exists() +} + +fn link_linux_cuda_libs(cmake_cache: &std::path::Path) { + for (cache_key, lib) in [ + ("CUDA_cuda_driver_LIBRARY", "cuda"), + ("CUDA_cudart_LIBRARY", "cudart"), + ("CUDA_cublas_LIBRARY", "cublas"), + ("CUDA_cublasLt_LIBRARY", "cublasLt"), + ] { + link_linux_lib_from_cache(cmake_cache, cache_key, lib); + } +} + +fn link_linux_hip_libs() { + for search_path in ["/opt/rocm/lib", "/opt/rocm/hip/lib"] { + if std::path::Path::new(search_path).is_dir() { + println!("cargo:rustc-link-search=native={search_path}"); + } + } + for lib in ["amdhip64", "rocblas", "hipblas"] { + println!("cargo:rustc-link-lib=dylib={lib}"); + } +} + +fn link_linux_lib_from_cache(cmake_cache: &std::path::Path, cache_key: &str, lib: &str) { + if let Ok(cache) = std::fs::read_to_string(cmake_cache) + && let Some(path) = cmake_cache_value(&cache, cache_key) + { + let path = std::path::PathBuf::from(path); + if path.exists() + && let Some(parent) = path.parent() + { + println!("cargo:rustc-link-search=native={}", parent.display()); + } + } + println!("cargo:rustc-link-lib=dylib={lib}"); +} + +fn cmake_cache_value(cache: &str, key: &str) -> Option { + cache.lines().find_map(|line| { + let (lhs, rhs) = line.split_once('=')?; + let (name, _) = lhs.split_once(':')?; + (name == key).then(|| rhs.to_string()) + }) +} diff --git a/crates/llama-quant-ffi/src/lib.rs b/crates/llama-quant-ffi/src/lib.rs new file mode 100644 index 0000000000..aff20a6159 --- /dev/null +++ b/crates/llama-quant-ffi/src/lib.rs @@ -0,0 +1,352 @@ +use std::ffi::c_char; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(i32)] +pub enum LlamaFileType { + AllF32 = 0, + MostlyF16 = 1, + MostlyQ4_0 = 2, + MostlyQ4_1 = 3, + MostlyQ8_0 = 7, + MostlyQ5_0 = 8, + MostlyQ5_1 = 9, + MostlyQ2K = 10, + MostlyQ3KS = 11, + MostlyQ3KM = 12, + MostlyQ3KL = 13, + MostlyQ4KS = 14, + MostlyQ4KM = 15, + MostlyQ5KS = 16, + MostlyQ5KM = 17, + MostlyQ6K = 18, + MostlyIQ2XXS = 19, + MostlyIQ2XS = 20, + MostlyQ2KS = 21, + MostlyIQ3XS = 22, + MostlyIQ3XXS = 23, + MostlyIQ1S = 24, + MostlyIQ4NL = 25, + MostlyIQ3S = 26, + MostlyIQ3M = 27, + MostlyIQ2S = 28, + MostlyIQ2M = 29, + MostlyIQ4XS = 30, + MostlyIQ1M = 31, + MostlyBf16 = 32, + MostlyTQ1_0 = 36, + MostlyTQ2_0 = 37, + MostlyMxfp4Moe = 38, + MostlyNvfp4 = 39, + MostlyQ1_0 = 40, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(i32)] +pub enum GgmlType { + F32 = 0, + F16 = 1, + Q4_0 = 2, + Q4_1 = 3, + Q5_0 = 6, + Q5_1 = 7, + Q8_0 = 8, + Q8_1 = 9, + Q2K = 10, + Q3K = 11, + Q4K = 12, + Q5K = 13, + Q6K = 14, + Q8K = 15, + IQ2XXS = 16, + IQ2XS = 17, + IQ3XXS = 18, + IQ1S = 19, + IQ4NL = 20, + IQ3S = 21, + IQ2S = 22, + IQ4XS = 23, + I8 = 24, + I16 = 25, + I32 = 26, + I64 = 27, + F64 = 28, + IQ1M = 29, + Bf16 = 30, + TQ1_0 = 34, + TQ2_0 = 35, + Mxfp4 = 39, + Nvfp4 = 40, + Q1_0 = 41, + Count = 42, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(i32)] +pub enum LlamaModelKvOverrideType { + Int = 0, + Float = 1, + Bool = 2, + Str = 3, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub union LlamaModelKvOverrideValue { + pub val_i64: i64, + pub val_f64: f64, + pub val_bool: bool, + pub val_str: [c_char; 128], +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct LlamaModelKvOverride { + pub tag: LlamaModelKvOverrideType, + pub key: [c_char; 128], + pub value: LlamaModelKvOverrideValue, +} + +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct LlamaModelTensorOverride { + pub pattern: *const c_char, + pub tensor_type: GgmlType, +} + +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct LlamaModelImatrixData { + pub name: *const c_char, + pub data: *const f32, + pub size: usize, +} + +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct LlamaModelQuantizeParams { + pub nthread: i32, + pub ftype: LlamaFileType, + pub output_tensor_type: GgmlType, + pub token_embedding_type: GgmlType, + pub allow_requantize: bool, + pub quantize_output_tensor: bool, + pub only_copy: bool, + pub pure: bool, + pub keep_split: bool, + pub dry_run: bool, + pub imatrix: *const LlamaModelImatrixData, + pub kv_overrides: *const LlamaModelKvOverride, + pub tt_overrides: *const LlamaModelTensorOverride, + pub prune_layers: *const i32, +} + +#[derive(Debug)] +pub enum NativeRuntimeLoadError { + Load(String), + AlreadyLoaded, +} + +impl std::fmt::Display for NativeRuntimeLoadError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Load(message) => write!(f, "{message}"), + Self::AlreadyLoaded => write!(f, "native runtime library is already loaded"), + } + } +} + +impl std::error::Error for NativeRuntimeLoadError {} + +#[cfg(not(feature = "dynamic-runtime"))] +pub fn native_runtime_loaded() -> bool { + true +} + +#[cfg(not(feature = "dynamic-runtime"))] +/// No-op for statically linked builds. +/// +/// # Safety +/// +/// Static builds resolve the native ABI at process link/load time, so this +/// function does not dereference the supplied path or mutate loader state. +pub unsafe fn load_native_runtime_library( + _path: impl AsRef, +) -> Result<(), NativeRuntimeLoadError> { + Ok(()) +} + +#[cfg(not(feature = "dynamic-runtime"))] +/// No-op for statically linked builds. +/// +/// # Safety +/// +/// Static builds resolve the native ABI at process link/load time, so this +/// function does not dereference the supplied paths or mutate loader state. +pub unsafe fn load_native_runtime_libraries(_paths: I) -> Result<(), NativeRuntimeLoadError> +where + I: IntoIterator, + P: AsRef, +{ + Ok(()) +} + +#[cfg(feature = "dynamic-runtime")] +mod dynamic { + use super::*; + use libloading::Library; + use std::sync::OnceLock; + + static SYMBOLS: OnceLock = OnceLock::new(); + + pub fn native_runtime_loaded() -> bool { + SYMBOLS.get().is_some() + } + + /// Load a native llama.cpp runtime library and resolve quantization symbols. + /// + /// # Safety + /// + /// The caller must ensure the library belongs to the same pinned llama.cpp + /// build and exposes an ABI-compatible `llama_model_quantize` surface. + pub unsafe fn load_native_runtime_library( + path: impl AsRef, + ) -> Result<(), NativeRuntimeLoadError> { + let symbols = unsafe { Symbols::load_paths(&[path.as_ref()]) }?; + SYMBOLS + .set(symbols) + .map_err(|_| NativeRuntimeLoadError::AlreadyLoaded) + } + + /// Load native runtime libraries and resolve quantization symbols. + /// + /// Libraries are searched from last to first so dependencies can be passed + /// before the primary `libllama`/`llama.dll` library. + /// + /// # Safety + /// + /// The caller must ensure every library belongs to the same pinned llama.cpp + /// build and exposes an ABI-compatible quantization surface. + pub unsafe fn load_native_runtime_libraries( + paths: I, + ) -> Result<(), NativeRuntimeLoadError> + where + I: IntoIterator, + P: AsRef, + { + let collected = paths + .into_iter() + .map(|path| path.as_ref().to_path_buf()) + .collect::>(); + let symbols = unsafe { Symbols::load_paths(&collected) }?; + SYMBOLS + .set(symbols) + .map_err(|_| NativeRuntimeLoadError::AlreadyLoaded) + } + + fn symbols() -> &'static Symbols { + SYMBOLS + .get() + .expect("llama quant native runtime library has not been loaded") + } + + struct Symbols { + _libraries: Vec, + llama_model_quantize_default_params: unsafe extern "C" fn() -> LlamaModelQuantizeParams, + llama_model_quantize: unsafe extern "C" fn( + *const c_char, + *const c_char, + *const LlamaModelQuantizeParams, + ) -> u32, + } + + impl Symbols { + unsafe fn load_paths

(paths: &[P]) -> Result + where + P: AsRef, + { + if paths.is_empty() { + return Err(NativeRuntimeLoadError::Load( + "native runtime did not provide any libraries".to_string(), + )); + } + let mut libraries = Vec::with_capacity(paths.len()); + for path in paths { + libraries.push( + unsafe { Library::new(path.as_ref()) } + .map_err(|err| NativeRuntimeLoadError::Load(err.to_string()))?, + ); + } + let llama_model_quantize_default_params = lookup_symbol( + &libraries, + b"llama_model_quantize_default_params\0", + "llama_model_quantize_default_params", + )?; + let llama_model_quantize = lookup_symbol( + &libraries, + b"llama_model_quantize\0", + "llama_model_quantize", + )?; + Ok(Self { + _libraries: libraries, + llama_model_quantize_default_params, + llama_model_quantize, + }) + } + } + + fn lookup_symbol( + libraries: &[Library], + name: &[u8], + label: &str, + ) -> Result + where + Sym: Copy + 'static, + { + for library in libraries.iter().rev() { + if let Ok(symbol) = unsafe { library.get::(name) } { + return Ok(*symbol); + } + } + Err(NativeRuntimeLoadError::Load(format!( + "native runtime symbol not found: {label}" + ))) + } + + /// Return llama.cpp quantization default parameters. + /// + /// # Safety + /// + /// The loaded native runtime must expose an ABI-compatible implementation. + pub unsafe fn llama_model_quantize_default_params() -> LlamaModelQuantizeParams { + unsafe { (symbols().llama_model_quantize_default_params)() } + } + + /// Quantize a GGUF model through llama.cpp. + /// + /// # Safety + /// + /// `fname_inp`, `fname_out`, and `params` must be valid pointers matching + /// llama.cpp's `llama_model_quantize` contract for the loaded runtime. + pub unsafe fn llama_model_quantize( + fname_inp: *const c_char, + fname_out: *const c_char, + params: *const LlamaModelQuantizeParams, + ) -> u32 { + unsafe { (symbols().llama_model_quantize)(fname_inp, fname_out, params) } + } +} + +#[cfg(feature = "dynamic-runtime")] +pub use dynamic::*; + +#[cfg(not(feature = "dynamic-runtime"))] +#[allow(clippy::missing_safety_doc)] +unsafe extern "C" { + pub fn llama_model_quantize_default_params() -> LlamaModelQuantizeParams; + + pub fn llama_model_quantize( + fname_inp: *const c_char, + fname_out: *const c_char, + params: *const LlamaModelQuantizeParams, + ) -> u32; +} diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs index 0875ea9651..76ccde89c1 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/mod.rs @@ -32,9 +32,8 @@ use skippy_runtime::ModelInfo; use skippy_server::{ DEFAULT_EMBEDDED_MAX_TOKENS, EmbeddedOpenAiArgs, EmbeddedRuntimeOptions, EmbeddedRuntimeStatus, EmbeddedServerHandle, EmbeddedState, OpenAiGuardrailsConfig, OpenAiGuardrailsStatus, - OpenAiGuardrailsTarget, SkippyRuntimeHandle, binary_transport::PredictionReturnListener, - binary_transport::WireCondition, embedded_openai_backend, telemetry::Telemetry, - telemetry::TelemetryLevel, + OpenAiGuardrailsTarget, SkippyRuntimeHandle, binary_transport::WireCondition, + embedded_openai_backend, telemetry::Telemetry, telemetry::TelemetryLevel, }; pub use certification::{ @@ -347,7 +346,6 @@ pub(crate) struct SkippyModelHandle { started_at_unix_nanos: i64, status: Arc>, _materialized_pin: Option, - _prediction_return_listener: Option, } pub(crate) struct SkippyHttpHandle { @@ -449,7 +447,6 @@ impl SkippyModelHandle { reply_credit_limit: embedded_args.reply_credit_limit, downstream_connect_timeout_secs: embedded_args.downstream_connect_timeout_secs, downstream_wire_condition: WireCondition::new(0.0, None)?, - prediction_returns: None, telemetry, hook_policy, openai_guardrails: None, @@ -473,7 +470,6 @@ impl SkippyModelHandle { last_error: None, })), _materialized_pin: None, - _prediction_return_listener: None, }) } @@ -544,7 +540,6 @@ impl SkippyModelHandle { reply_credit_limit: embedded_args.reply_credit_limit, downstream_connect_timeout_secs: embedded_args.downstream_connect_timeout_secs, downstream_wire_condition: WireCondition::new(0.0, None)?, - prediction_returns: None, telemetry, hook_policy, openai_guardrails: None, @@ -568,7 +563,6 @@ impl SkippyModelHandle { last_error: None, })), _materialized_pin: None, - _prediction_return_listener: None, }) } @@ -680,16 +674,6 @@ impl SkippyModelHandle { runtime_config.clone(), telemetry.level, ); - let prediction_return_listener = if runtime_config.downstream.is_some() { - Some(PredictionReturnListener::start( - runtime_config.bind_addr.parse()?, - )?) - } else { - None - }; - let prediction_returns = prediction_return_listener - .as_ref() - .map(PredictionReturnListener::hub); let binding = embedded_openai_backend(EmbeddedOpenAiArgs { bind_addr: "127.0.0.1:0" .parse() @@ -715,7 +699,6 @@ impl SkippyModelHandle { reply_credit_limit: embedded_args.reply_credit_limit, downstream_connect_timeout_secs: embedded_args.downstream_connect_timeout_secs, downstream_wire_condition: WireCondition::new(0.0, None)?, - prediction_returns, telemetry, hook_policy, openai_guardrails: None, @@ -739,7 +722,6 @@ impl SkippyModelHandle { last_error: None, })), _materialized_pin: materialized_pin, - _prediction_return_listener: prediction_return_listener, }) } @@ -802,16 +784,6 @@ impl SkippyModelHandle { runtime_config.clone(), telemetry.level, ); - let prediction_return_listener = if runtime_config.downstream.is_some() { - Some(PredictionReturnListener::start( - runtime_config.bind_addr.parse()?, - )?) - } else { - None - }; - let prediction_returns = prediction_return_listener - .as_ref() - .map(PredictionReturnListener::hub); let binding = embedded_openai_backend(EmbeddedOpenAiArgs { bind_addr: "127.0.0.1:0" .parse() @@ -837,7 +809,6 @@ impl SkippyModelHandle { reply_credit_limit: embedded_args.reply_credit_limit, downstream_connect_timeout_secs: embedded_args.downstream_connect_timeout_secs, downstream_wire_condition: WireCondition::new(0.0, None)?, - prediction_returns, telemetry, hook_policy, openai_guardrails: None, @@ -861,7 +832,6 @@ impl SkippyModelHandle { last_error: None, })), _materialized_pin: materialized_pin, - _prediction_return_listener: prediction_return_listener, }) } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs index d963f2fae3..6f44170249 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs @@ -412,7 +412,6 @@ impl ResolvedEmbeddedOpenAiArgs { 0.0, None, ) .expect("static downstream wire condition should construct"), - prediction_returns: None, telemetry, hook_policy, openai_guardrails: None, diff --git a/crates/model-artifact/src/gguf.rs b/crates/model-artifact/src/gguf.rs index 18ceadf8e0..e2b2dfd454 100644 --- a/crates/model-artifact/src/gguf.rs +++ b/crates/model-artifact/src/gguf.rs @@ -239,6 +239,7 @@ pub struct GgufCompactMeta { pub rope_freq_base: f32, pub expert_count: u32, pub expert_used_count: u32, + pub nextn_predict_layers: u32, } impl GgufCompactMeta { @@ -489,21 +490,28 @@ pub fn scan_gguf_compact_meta(path: &Path) -> Option { if let Ok(Some(v)) = read_gguf_value_as_u32(&mut f, vtype) { meta.expert_used_count = v; } + } else if key.ends_with(".nextn_predict_layers") { + if let Ok(Some(v)) = read_gguf_value_as_u32(&mut f, vtype) { + meta.nextn_predict_layers = v; + } } else { skip_gguf_value(&mut f, vtype).ok()?; } } - if meta.key_length == 0 - && meta.head_count > 0 - && let Some(key_length) = meta.embedding_size.checked_div(meta.head_count) - { + let derived_key_length = (meta.key_length == 0 && meta.head_count > 0) + .then(|| meta.embedding_size.checked_div(meta.head_count)) + .flatten(); + if let Some(key_length) = derived_key_length { meta.key_length = key_length; } - if meta.value_length == 0 - && let Some(effective_kv) = meta.effective_kv_head_count() - && let Some(value_length) = meta.embedding_size.checked_div(effective_kv) - { + let derived_value_length = (meta.value_length == 0) + .then(|| { + meta.effective_kv_head_count() + .and_then(|effective_kv| meta.embedding_size.checked_div(effective_kv)) + }) + .flatten(); + if let Some(value_length) = derived_value_length { meta.value_length = value_length; } @@ -802,6 +810,25 @@ mod tests { let _ = std::fs::remove_file(path); } + #[test] + fn scan_gguf_compact_meta_preserves_nextn_predict_layers() { + let mut bytes = Vec::new(); + bytes.extend_from_slice(b"GGUF"); + bytes.extend_from_slice(&2u32.to_le_bytes()); + bytes.extend_from_slice(&0i64.to_le_bytes()); + bytes.extend_from_slice(&2i64.to_le_bytes()); + push_gguf_string(&mut bytes, "general.architecture"); + bytes.extend_from_slice(&(GgufType::String as u32).to_le_bytes()); + push_gguf_string(&mut bytes, "deepseek2"); + push_u32_kv(&mut bytes, "deepseek2.nextn_predict_layers", 1); + + let path = write_bytes("model-artifact-gguf-nextn", &bytes); + let meta = scan_gguf_compact_meta(&path).expect("should parse GGUF"); + assert_eq!(meta.architecture, "deepseek2"); + assert_eq!(meta.nextn_predict_layers, 1); + let _ = std::fs::remove_file(path); + } + #[test] fn kv_cache_quant_prices_key_and_value_types_independently() { let meta = GgufCompactMeta { diff --git a/crates/skippy-bench/src/cli.rs b/crates/skippy-bench/src/cli.rs index 400355db00..eb970e53ea 100644 --- a/crates/skippy-bench/src/cli.rs +++ b/crates/skippy-bench/src/cli.rs @@ -20,6 +20,8 @@ pub enum CommandKind { LocalSplitBinary(LocalSplitBinaryArgs), LocalSplitCompare(LocalSplitCompareArgs), LocalSplitChainBinary(LocalSplitChainBinaryArgs), + #[command(name = "verify-span-local")] + VerifySpanLocal(VerifySpanLocalArgs), #[command(name = "chat-corpus")] ChatCorpus(ChatCorpusArgs), #[command(name = "token-lengths")] @@ -88,6 +90,39 @@ pub struct TokenLengthsArgs { pub summary_json: Option, } +#[derive(Parser)] +pub struct VerifySpanLocalArgs { + #[arg(long)] + pub model_path: PathBuf, + #[arg(long, default_value_t = 48)] + pub layer_end: u32, + #[arg(long)] + pub split_layer: Option, + #[arg(long, default_value_t = 4096)] + pub ctx_size: u32, + #[arg(long, default_value_t = -1, allow_hyphen_values = true)] + pub n_gpu_layers: i32, + #[arg(long, default_value = "f16")] + pub cache_type_k: String, + #[arg(long, default_value = "f16")] + pub cache_type_v: String, + #[arg(long)] + pub n_batch: Option, + #[arg(long)] + pub n_ubatch: Option, + #[arg(long, default_value_t = 64)] + pub iterations: usize, + #[arg(long, default_value_t = 8)] + pub warmup: usize, + #[arg( + long, + default_value = "Write a Rust function that parses a list of integers and returns the median." + )] + pub prompt: String, + #[arg(long)] + pub output: Option, +} + #[derive(Parser)] pub struct ChatCorpusArgs { #[arg(long, default_value = "http://127.0.0.1:9337/v1")] @@ -415,6 +450,8 @@ pub struct LocalSplitChainBinaryArgs { #[cfg(test)] mod tests { + use std::path::PathBuf; + use clap::Parser; use super::{Cli, CommandKind, FocusedRuntimeScenario}; @@ -471,4 +508,53 @@ mod tests { assert_eq!(args.run.max_new_tokens, None); } + + #[test] + fn parses_verify_span_local_command() { + let cli = Cli::try_parse_from([ + "skippy-bench", + "verify-span-local", + "--model-path", + "/tmp/model.gguf", + "--layer-end", + "48", + "--iterations", + "3", + "--warmup", + "1", + "--n-gpu-layers", + "-1", + ]) + .unwrap(); + + let CommandKind::VerifySpanLocal(args) = cli.command else { + panic!("expected verify-span-local subcommand"); + }; + + assert_eq!(args.model_path, PathBuf::from("/tmp/model.gguf")); + assert_eq!(args.layer_end, 48); + assert_eq!(args.split_layer, None); + assert_eq!(args.iterations, 3); + assert_eq!(args.warmup, 1); + assert_eq!(args.n_gpu_layers, -1); + } + + #[test] + fn parses_verify_span_local_split_layer() { + let cli = Cli::try_parse_from([ + "skippy-bench", + "verify-span-local", + "--model-path", + "/tmp/model.gguf", + "--split-layer", + "24", + ]) + .unwrap(); + + let CommandKind::VerifySpanLocal(args) = cli.command else { + panic!("expected verify-span-local subcommand"); + }; + + assert_eq!(args.split_layer, Some(24)); + } } diff --git a/crates/skippy-bench/src/direct_return.rs b/crates/skippy-bench/src/direct_return.rs deleted file mode 100644 index eb6c8f579f..0000000000 --- a/crates/skippy-bench/src/direct_return.rs +++ /dev/null @@ -1,160 +0,0 @@ -use std::{ - collections::HashMap, - io, - net::{SocketAddr, TcpListener, TcpStream}, - sync::{Arc, Mutex, mpsc}, - thread, - time::Duration, -}; - -use anyhow::{Context, Result, anyhow, bail}; -use skippy_protocol::binary::{StageReply, WireMessageKind, WireReplyKind, recv_reply, send_ready}; - -#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] -struct DirectReturnKey { - request_id: u64, - session_id: u64, -} - -type DirectReturnResult = Result; -type DirectReturnSender = mpsc::Sender; -type DirectReturnWaiters = Arc>>; - -pub(crate) struct BenchDirectReturnServer { - local_addr: SocketAddr, - waiters: DirectReturnWaiters, -} - -impl BenchDirectReturnServer { - pub(crate) fn start(bind_addr: &str) -> Result { - let listener = TcpListener::bind(bind_addr) - .with_context(|| format!("bind benchmark direct-return listener {bind_addr}"))?; - let local_addr = listener - .local_addr() - .context("read benchmark direct-return listener address")?; - let waiters = Arc::new(Mutex::new(HashMap::new())); - let thread_waiters = waiters.clone(); - thread::spawn(move || { - for stream in listener.incoming() { - match stream { - Ok(stream) => { - let waiters = thread_waiters.clone(); - thread::spawn(move || { - if let Err(error) = - handle_bench_direct_return_connection(waiters, stream) - { - eprintln!("benchmark direct-return connection failed: {error:#}"); - } - }); - } - Err(error) if error.kind() == io::ErrorKind::Interrupted => continue, - Err(error) => { - eprintln!("benchmark direct-return listener failed: {error}"); - break; - } - } - } - }); - Ok(Self { - local_addr, - waiters, - }) - } - - pub(crate) fn endpoint(&self) -> String { - self.local_addr.to_string() - } - - pub(crate) fn register( - &self, - request_id: u64, - session_id: u64, - ) -> Result { - let key = DirectReturnKey { - request_id, - session_id, - }; - let (sender, receiver) = mpsc::channel(); - self.waiters - .lock() - .map_err(|_| anyhow!("benchmark direct-return hub lock poisoned"))? - .insert(key, sender); - Ok(BenchDirectReturnReceiver { - key, - waiters: self.waiters.clone(), - receiver, - }) - } -} - -pub(crate) struct BenchDirectReturnReceiver { - key: DirectReturnKey, - waiters: DirectReturnWaiters, - receiver: mpsc::Receiver, -} - -impl BenchDirectReturnReceiver { - pub(crate) fn recv_expected(&self, expected: WireReplyKind) -> Result { - let reply = self - .receiver - .recv_timeout(Duration::from_secs(300)) - .context("timed out waiting for benchmark direct prediction return")? - .map_err(|error| anyhow!(error))?; - if reply.kind != expected { - bail!( - "expected {expected:?} direct prediction return, got {:?}", - reply.kind - ); - } - Ok(reply) - } -} - -impl Drop for BenchDirectReturnReceiver { - fn drop(&mut self) { - if let Ok(mut waiters) = self.waiters.lock() { - waiters.remove(&self.key); - } - } -} - -fn handle_bench_direct_return_connection( - waiters: DirectReturnWaiters, - mut stream: TcpStream, -) -> Result<()> { - send_ready(&mut stream).context("send benchmark direct-return ready")?; - let open = skippy_protocol::binary::read_stage_message(&mut stream, 0) - .context("read benchmark direct-return open")?; - if open.kind != WireMessageKind::PredictionReturnOpen { - bail!("expected prediction-return-open message"); - } - let key = DirectReturnKey { - request_id: open.request_id, - session_id: open.session_id, - }; - let sender = waiters - .lock() - .map_err(|_| anyhow!("benchmark direct-return hub lock poisoned"))? - .get(&key) - .cloned() - .ok_or_else(|| { - anyhow!( - "no benchmark direct-return waiter for request {}", - key.request_id - ) - })?; - loop { - match recv_reply(&mut stream) { - Ok(reply) => { - if sender.send(Ok(reply)).is_err() { - return Ok(()); - } - } - Err(error) if error.kind() == io::ErrorKind::UnexpectedEof => return Ok(()), - Err(error) => { - let _ = sender.send(Err(error.to_string())); - return Err(error).context("read benchmark direct prediction return"); - } - } - } -} diff --git a/crates/skippy-bench/src/distributed.rs b/crates/skippy-bench/src/distributed.rs index 1138b8a307..796752122b 100644 --- a/crates/skippy-bench/src/distributed.rs +++ b/crates/skippy-bench/src/distributed.rs @@ -35,8 +35,6 @@ use crate::{ support::{ChildGuard, parse_wire_dtype, retry}, }; -use crate::direct_return::BenchDirectReturnServer; - struct DistributedRunOutcome { run_id: String, topology_id: String, @@ -187,7 +185,6 @@ struct DeploymentPlan { work_dir: PathBuf, metrics_http: String, metrics_otlp_grpc: String, - driver_return_bind_addr: String, driver_return_endpoint: String, stages: Vec, execute_remote: bool, @@ -1075,7 +1072,6 @@ fn build_deployment_plan( work_dir: args.work_dir.clone(), metrics_http, metrics_otlp_grpc: metrics_otlp, - driver_return_bind_addr: driver_return_bind_addr(args), driver_return_endpoint: driver_return_endpoint(args, &stages)?, stages, execute_remote: args.execute_remote, @@ -1215,10 +1211,6 @@ fn parse_load_mode(load_mode: &str) -> Result { } } -fn driver_return_bind_addr(args: &RunArgs) -> String { - format!("0.0.0.0:{}", driver_return_port(args)) -} - fn driver_return_endpoint(args: &RunArgs, stages: &[StageAssignment]) -> Result { let first = stages.first().context("deployment plan has no stages")?; let endpoint = first @@ -1584,8 +1576,6 @@ fn run_remote_prompt_driver(args: &RunArgs, plan: &DeploymentPlan) -> Result Result() as f64 / results.len() as f64 } +fn ensure_reply_kind( + reply: &skippy_protocol::binary::StageReply, + expected: WireReplyKind, +) -> Result<()> { + if reply.kind != expected { + bail!("expected {expected:?} reply, got {:?}", reply.kind); + } + Ok(()) +} + fn run_remote_prompt_case( args: &RunArgs, first: &StageAssignment, @@ -1704,7 +1697,6 @@ fn run_remote_prompt_case( prompt_case: &PromptCase, token_ids: Vec, prompt_index: usize, - direct_returns: &BenchDirectReturnServer, ) -> Result { if token_ids.is_empty() { bail!("prompt produced no tokens"); @@ -1721,7 +1713,6 @@ fn run_remote_prompt_case( let wire_started = Instant::now(); let request_id = 10_000_u64 + prompt_index as u64; let session_id = 20_000_u64 + prompt_index as u64; - let direct_return = direct_returns.register(request_id, session_id)?; send_generation_config( &mut stream, wire_dtype, @@ -1791,11 +1782,10 @@ fn run_remote_prompt_case( write_stage_message(&mut stream, &message, wire_dtype).with_context(|| { format!("send remote decode step {decode_step} for prompt {prompt_index}") })?; - let reply = direct_return - .recv_expected(WireReplyKind::PredictedToken) - .with_context(|| { - format!("receive direct decode step {decode_step} reply for prompt {prompt_index}") - })?; + let reply = recv_reply(&mut stream).with_context(|| { + format!("receive decode step {decode_step} reply for prompt {prompt_index}") + })?; + ensure_reply_kind(&reply, WireReplyKind::PredictedToken)?; if decode_step == 0 { ttft_ms = wire_started.elapsed().as_millis(); } @@ -3349,7 +3339,6 @@ mod tests { work_dir: PathBuf::from("/tmp/work"), metrics_http: "http://127.0.0.1:18080".to_string(), metrics_otlp_grpc: "http://coordinator.local:14317".to_string(), - driver_return_bind_addr: "0.0.0.0:20031".to_string(), driver_return_endpoint: "host.local:20031".to_string(), stages: Vec::new(), execute_remote: true, diff --git a/crates/skippy-bench/src/local_split.rs b/crates/skippy-bench/src/local_split.rs index 3f304eff29..97a13d7ac6 100644 --- a/crates/skippy-bench/src/local_split.rs +++ b/crates/skippy-bench/src/local_split.rs @@ -21,7 +21,6 @@ use crate::{ LocalSplitBinaryArgs, LocalSplitChainBinaryArgs, LocalSplitCompareArgs, LocalSplitInprocessArgs, }, - direct_return::BenchDirectReturnServer, model_identity::model_identity_for_path, support::{ ChildGuard, activation_width, connect_ready, generate_run_id, parse_wire_dtype, @@ -56,6 +55,16 @@ struct BinaryChainResult { layer_end: u32, } +fn ensure_reply_kind( + reply: &skippy_protocol::binary::StageReply, + expected: WireReplyKind, +) -> Result<()> { + if reply.kind != expected { + bail!("expected {expected:?} reply, got {:?}", reply.kind); + } + Ok(()) +} + pub fn local_split_binary(args: LocalSplitBinaryArgs) -> Result<()> { let result = run_binary_split(BinarySplitConfig { stage_server_bin: args.stage_server_bin, @@ -308,7 +317,6 @@ fn run_binary_split(args: BinarySplitConfig) -> Result { bail!("stage 0 produced an empty activation frame"); } let activation_width = activation_width(&boundary)?; - let direct_returns = BenchDirectReturnServer::start("127.0.0.1:0")?; let run_id = generate_run_id(); let config_path = temp_config_path_for(&run_id, "stage-1"); @@ -341,7 +349,7 @@ fn run_binary_split(args: BinarySplitConfig) -> Result { LocalSplitTopologyStage { stage_id: "stage-0", stage_index: 0, - endpoint: format!("tcp://{}", direct_returns.endpoint()), + endpoint: "driver".to_string(), layer_start: 0, layer_end: args.split_layer, }, @@ -388,7 +396,6 @@ fn run_binary_split(args: BinarySplitConfig) -> Result { .context("stage 1 binary server did not become ready")?; let request_id = 1; let session_id = 1; - let direct_return = direct_returns.register(request_id, session_id)?; send_generation_config(&mut stream, wire_dtype, request_id, session_id, 1) .context("send binary split generation config")?; let mut state = StageStateHeader::new(WireMessageKind::DecodeEmbd, wire_dtype); @@ -421,9 +428,8 @@ fn run_binary_split(args: BinarySplitConfig) -> Result { raw_bytes: Vec::new(), }; write_stage_message(&mut stream, &message, wire_dtype).context("send binary decode")?; - let reply = direct_return - .recv_expected(WireReplyKind::PredictedToken) - .context("receive direct binary reply")?; + let reply = recv_reply(&mut stream).context("receive binary split prediction reply")?; + ensure_reply_kind(&reply, WireReplyKind::PredictedToken)?; write_stage_message(&mut stream, &StageWireMessage::stop(wire_dtype), wire_dtype) .context("send binary stop")?; @@ -495,7 +501,6 @@ fn run_binary_chain(args: LocalSplitChainBinaryArgs) -> Result Result Result Result Result<()> { @@ -29,6 +30,7 @@ fn main() -> Result<()> { CommandKind::LocalSplitBinary(args) => local_split_binary(args), CommandKind::LocalSplitCompare(args) => local_split_compare(args), CommandKind::LocalSplitChainBinary(args) => local_split_chain_binary(args), + CommandKind::VerifySpanLocal(args) => verify_span_local(args), CommandKind::ChatCorpus(args) => chat_corpus(args), CommandKind::TokenLengths(args) => token_lengths(args), CommandKind::FocusedRuntime(args) => focused_runtime(args), diff --git a/crates/skippy-bench/src/verify_span_local.rs b/crates/skippy-bench/src/verify_span_local.rs new file mode 100644 index 0000000000..24e9a39915 --- /dev/null +++ b/crates/skippy-bench/src/verify_span_local.rs @@ -0,0 +1,888 @@ +use std::{ + fs, + path::PathBuf, + time::{Duration, Instant}, +}; + +use anyhow::{Context, Result, bail}; +use serde::Serialize; +use skippy_runtime::{ + FlashAttentionType, RuntimeConfig, RuntimeLoadMode, SamplingConfig, StageModel, StageSession, + parse_cache_type, +}; + +use crate::cli::VerifySpanLocalArgs; + +#[derive(Debug, Serialize)] +struct TimingStats { + count: usize, + total_us: u128, + avg_us: f64, + min_us: u128, + p50_us: u128, + p95_us: u128, + max_us: u128, +} + +#[derive(Debug, Serialize)] +struct TimingShape { + first_half: TimingStats, + second_half: TimingStats, + second_half_avg_vs_first_half_avg: f64, + first_sample_us: u128, + last_sample_us: u128, + samples_us: Vec, +} + +#[derive(Debug, Serialize)] +struct VerifySpanLocalReport { + mode: &'static str, + model_path: PathBuf, + layer_end: u32, + split_layer: Option, + ctx_size: u32, + n_gpu_layers: i32, + n_batch: Option, + n_ubatch: Option, + cache_type_k: String, + cache_type_v: String, + prompt_token_count: usize, + verify_tokens: Vec, + warmup: usize, + iterations: usize, + batched_width2: TimingStats, + serial_two_decode_mtp_n1: TimingStats, + split_inprocess_width2: Option, + batched_avg_vs_serial_avg: f64, + batched_token_per_sec: f64, + serial_token_per_sec: f64, + first_batched_prediction: Vec, + first_serial_prediction: Vec, +} + +#[derive(Debug, Serialize)] +struct SplitInprocessReport { + split_layer: u32, + boundary_payload_bytes: usize, + serial_boundary_payload_bytes: usize, + total: TimingStats, + stage0: TimingStats, + stage1: TimingStats, + serial_total: TimingStats, + serial_stage0: TimingStats, + serial_stage1: TimingStats, + total_token_per_sec: f64, + serial_total_token_per_sec: f64, + total_avg_vs_full_batched_avg: f64, + total_avg_vs_serial_total_avg: f64, + diagnostics: SplitTimingDiagnostics, + first_prediction: Vec, + first_serial_prediction: Vec, +} + +#[derive(Debug, Serialize)] +struct SplitTimingDiagnostics { + batched_total: TimingShape, + batched_stage0: TimingShape, + batched_stage1: TimingShape, + serial_total: TimingShape, + serial_stage0: TimingShape, + serial_stage1: TimingShape, +} + +pub fn verify_span_local(args: VerifySpanLocalArgs) -> Result<()> { + validate_args(&args)?; + let output = args.output.clone(); + let full = run_full_model_samples(&args)?; + let split = match args.split_layer { + Some(split_layer) => Some(run_split_inprocess_samples( + &args, + split_layer, + &full.tokens, + &full.verify_tokens, + full.samples.batched_avg_us()?, + )?), + None => None, + }; + let report = build_report(args, full, split)?; + let encoded = serde_json::to_vec_pretty(&report)?; + + if let Some(path) = output { + fs::write(&path, &encoded) + .with_context(|| format!("failed to write {}", path.display()))?; + } + println!("{}", String::from_utf8(encoded)?); + Ok(()) +} + +fn validate_args(args: &VerifySpanLocalArgs) -> Result<()> { + if args.layer_end == 0 { + bail!("layer_end must be greater than zero"); + } + if args.iterations == 0 { + bail!("iterations must be greater than zero"); + } + if let Some(split_layer) = args.split_layer + && (split_layer == 0 || split_layer >= args.layer_end) + { + bail!("split_layer must be greater than zero and less than layer_end"); + } + Ok(()) +} + +fn full_runtime_config(args: &VerifySpanLocalArgs) -> Result { + Ok(RuntimeConfig { + stage_index: 0, + layer_start: 0, + layer_end: args.layer_end, + ctx_size: args.ctx_size, + lane_count: 1, + n_batch: args.n_batch, + n_ubatch: args.n_ubatch, + n_threads: None, + n_threads_batch: None, + n_gpu_layers: args.n_gpu_layers, + selected_backend_device: None, + cache_type_k: parse_cache_type(&args.cache_type_k)?, + cache_type_v: parse_cache_type(&args.cache_type_v)?, + flash_attn_type: FlashAttentionType::Auto, + load_mode: RuntimeLoadMode::RuntimeSlice, + projector_path: None, + include_embeddings: true, + include_output: true, + filter_tensors_on_load: false, + }) +} + +struct FullModelSamples { + tokens: Vec, + verify_tokens: Vec, + samples: SampleSet, +} + +fn run_full_model_samples(args: &VerifySpanLocalArgs) -> Result { + let config = full_runtime_config(args)?; + let model = StageModel::open(&args.model_path, &config) + .with_context(|| format!("failed to open {}", args.model_path.display()))?; + let tokens = model + .tokenize(&args.prompt, true) + .context("failed to tokenize prompt")?; + if tokens.is_empty() { + bail!("prompt produced no tokens"); + } + + let mut session = model.create_session().context("failed to create session")?; + session + .prefill_chunked(&tokens) + .context("failed to prefill prompt")?; + let base_token_count = session.token_count(); + let verify_tokens = choose_verify_tokens( + &mut session, + base_token_count, + &tokens, + &args.prompt, + &config, + )?; + let samples = run_samples( + &mut session, + base_token_count, + &verify_tokens, + args.warmup, + args.iterations, + )?; + Ok(FullModelSamples { + tokens, + verify_tokens, + samples, + }) +} + +fn choose_verify_tokens( + session: &mut StageSession, + base_token_count: u64, + prompt_tokens: &[i32], + prompt: &str, + config: &RuntimeConfig, +) -> Result> { + session + .trim_session(base_token_count) + .context("failed to trim session before choosing verify tokens")?; + let current = *prompt_tokens + .first() + .context("prompt produced no token for verify-token seed")?; + let (_after_current, native_mtp, _frame) = session + .decode_step_frame_sampled_mtp_n1(current, Some(&SamplingConfig::default()), None, 0) + .with_context(|| { + format!( + "failed to get native MTP draft from {} after prompt {:?}", + model_description(config), + prompt + ) + })?; + let Some(draft) = native_mtp else { + bail!("model did not produce a native MTP n=1 draft token"); + }; + session + .trim_session(base_token_count) + .context("failed to trim session after choosing verify tokens")?; + Ok(vec![current, draft.token_id]) +} + +fn run_samples( + session: &mut StageSession, + base_token_count: u64, + verify_tokens: &[i32], + warmup: usize, + iterations: usize, +) -> Result { + let total = warmup + .checked_add(iterations) + .context("sample count overflow")?; + let mut batched = Vec::with_capacity(iterations); + let mut serial = Vec::with_capacity(iterations); + let mut first_batched_prediction = None; + let mut first_serial_prediction = None; + + { + let mut targets = SampleRecordTargets { + batched: &mut batched, + serial: &mut serial, + first_batched_prediction: &mut first_batched_prediction, + first_serial_prediction: &mut first_serial_prediction, + }; + for index in 0..total { + let record = index >= warmup; + if index.is_multiple_of(2) { + measure_batched_then_serial( + session, + base_token_count, + verify_tokens, + record, + &mut targets, + )?; + } else { + measure_serial_then_batched( + session, + base_token_count, + verify_tokens, + record, + &mut targets, + )?; + } + } + } + + Ok(SampleSet { + batched, + serial, + first_batched_prediction: first_batched_prediction.unwrap_or_default(), + first_serial_prediction: first_serial_prediction.unwrap_or_default(), + }) +} + +fn measure_batched_then_serial( + session: &mut StageSession, + base_token_count: u64, + verify_tokens: &[i32], + record: bool, + targets: &mut SampleRecordTargets<'_>, +) -> Result<()> { + let (batched_duration, batched_prediction) = + measure_batched(session, base_token_count, verify_tokens)?; + let (serial_duration, serial_prediction) = + measure_serial(session, base_token_count, verify_tokens)?; + record_sample( + record, + (batched_duration, batched_prediction), + (serial_duration, serial_prediction), + targets, + ); + Ok(()) +} + +fn measure_serial_then_batched( + session: &mut StageSession, + base_token_count: u64, + verify_tokens: &[i32], + record: bool, + targets: &mut SampleRecordTargets<'_>, +) -> Result<()> { + let (serial_duration, serial_prediction) = + measure_serial(session, base_token_count, verify_tokens)?; + let (batched_duration, batched_prediction) = + measure_batched(session, base_token_count, verify_tokens)?; + record_sample( + record, + (batched_duration, batched_prediction), + (serial_duration, serial_prediction), + targets, + ); + Ok(()) +} + +struct SampleRecordTargets<'a> { + batched: &'a mut Vec, + serial: &'a mut Vec, + first_batched_prediction: &'a mut Option>, + first_serial_prediction: &'a mut Option>, +} + +fn record_sample( + record: bool, + batched_sample: (Duration, Vec), + serial_sample: (Duration, Vec), + targets: &mut SampleRecordTargets<'_>, +) { + if !record { + return; + } + targets.batched.push(batched_sample.0); + targets.serial.push(serial_sample.0); + targets + .first_batched_prediction + .get_or_insert(batched_sample.1); + targets + .first_serial_prediction + .get_or_insert(serial_sample.1); +} + +fn measure_batched( + session: &mut StageSession, + base_token_count: u64, + verify_tokens: &[i32], +) -> Result<(Duration, Vec)> { + session + .trim_session(base_token_count) + .context("failed to trim session before batched verify")?; + let start = Instant::now(); + let prediction = session + .verify_tokens_frame_sampled(verify_tokens, Some(&SamplingConfig::default()), None, 0) + .context("batched width-2 VerifySpan failed")? + .0; + Ok((start.elapsed(), prediction)) +} + +fn measure_serial( + session: &mut StageSession, + base_token_count: u64, + verify_tokens: &[i32], +) -> Result<(Duration, Vec)> { + session + .trim_session(base_token_count) + .context("failed to trim session before serial verify")?; + let start = Instant::now(); + let prediction = serial_decode_mtp_n1(session, verify_tokens)?; + Ok((start.elapsed(), prediction)) +} + +fn run_split_inprocess_samples( + args: &VerifySpanLocalArgs, + split_layer: u32, + tokens: &[i32], + verify_tokens: &[i32], + full_batched_avg_us: f64, +) -> Result { + let (stage0_config, stage1_config) = split_runtime_configs(args, split_layer)?; + let stage0 = StageModel::open(&args.model_path, &stage0_config) + .context("failed to open in-process split stage 0")?; + let stage1 = StageModel::open(&args.model_path, &stage1_config) + .context("failed to open in-process split stage 1")?; + let mut session0 = stage0 + .create_session() + .context("failed to create in-process split stage 0 session")?; + let mut session1 = stage1 + .create_session() + .context("failed to create in-process split stage 1 session")?; + prefill_split_sessions(&mut session0, &mut session1, tokens)?; + let base0 = session0.token_count(); + let base1 = session1.token_count(); + let samples = run_split_samples( + &mut session0, + &mut session1, + base0, + base1, + verify_tokens, + args.warmup, + args.iterations, + )?; + split_report(split_layer, samples, full_batched_avg_us) +} + +fn split_runtime_configs( + args: &VerifySpanLocalArgs, + split_layer: u32, +) -> Result<(RuntimeConfig, RuntimeConfig)> { + let cache_type_k = parse_cache_type(&args.cache_type_k)?; + let cache_type_v = parse_cache_type(&args.cache_type_v)?; + let stage0 = RuntimeConfig { + stage_index: 0, + layer_start: 0, + layer_end: split_layer, + ctx_size: args.ctx_size, + lane_count: 1, + n_batch: args.n_batch, + n_ubatch: args.n_ubatch, + n_threads: None, + n_threads_batch: None, + n_gpu_layers: args.n_gpu_layers, + selected_backend_device: None, + cache_type_k, + cache_type_v, + flash_attn_type: FlashAttentionType::Auto, + load_mode: RuntimeLoadMode::RuntimeSlice, + projector_path: None, + include_embeddings: true, + include_output: false, + filter_tensors_on_load: true, + }; + let stage1 = RuntimeConfig { + stage_index: 1, + layer_start: split_layer, + layer_end: args.layer_end, + ctx_size: args.ctx_size, + lane_count: 1, + n_batch: args.n_batch, + n_ubatch: args.n_ubatch, + n_threads: None, + n_threads_batch: None, + n_gpu_layers: args.n_gpu_layers, + selected_backend_device: None, + cache_type_k, + cache_type_v, + flash_attn_type: FlashAttentionType::Auto, + load_mode: RuntimeLoadMode::RuntimeSlice, + projector_path: None, + include_embeddings: false, + include_output: true, + filter_tensors_on_load: true, + }; + Ok((stage0, stage1)) +} + +fn prefill_split_sessions( + session0: &mut StageSession, + session1: &mut StageSession, + tokens: &[i32], +) -> Result<()> { + let (_stage0_prediction, boundary) = session0 + .prefill_chunk_frame_sampled(tokens, Some(&SamplingConfig::default()), None, 0) + .context("in-process split stage 0 failed to prefill")?; + if boundary.payload.is_empty() { + bail!("in-process split stage 0 produced an empty prefill activation frame"); + } + session1 + .prefill_chunk_frame_sampled(tokens, Some(&SamplingConfig::default()), Some(&boundary), 0) + .context("in-process split stage 1 failed to prefill")?; + Ok(()) +} + +fn run_split_samples( + session0: &mut StageSession, + session1: &mut StageSession, + base0: u64, + base1: u64, + verify_tokens: &[i32], + warmup: usize, + iterations: usize, +) -> Result { + let total = warmup + .checked_add(iterations) + .context("split sample count overflow")?; + let mut total_samples = Vec::with_capacity(iterations); + let mut stage0_samples = Vec::with_capacity(iterations); + let mut stage1_samples = Vec::with_capacity(iterations); + let mut serial_total_samples = Vec::with_capacity(iterations); + let mut serial_stage0_samples = Vec::with_capacity(iterations); + let mut serial_stage1_samples = Vec::with_capacity(iterations); + let mut boundary_payload_bytes = 0usize; + let mut serial_boundary_payload_bytes = 0usize; + let mut first_prediction = None; + let mut first_serial_prediction = None; + + for index in 0..total { + let (batched, serial) = if index.is_multiple_of(2) { + ( + measure_split_batched(session0, session1, base0, base1, verify_tokens)?, + measure_split_serial(session0, session1, base0, base1, verify_tokens)?, + ) + } else { + let serial = measure_split_serial(session0, session1, base0, base1, verify_tokens)?; + let batched = measure_split_batched(session0, session1, base0, base1, verify_tokens)?; + (batched, serial) + }; + if index >= warmup { + total_samples.push(batched.total); + stage0_samples.push(batched.stage0); + stage1_samples.push(batched.stage1); + serial_total_samples.push(serial.total); + serial_stage0_samples.push(serial.stage0); + serial_stage1_samples.push(serial.stage1); + boundary_payload_bytes = batched.boundary_payload_bytes; + serial_boundary_payload_bytes = serial.boundary_payload_bytes; + first_prediction.get_or_insert(batched.prediction); + first_serial_prediction.get_or_insert(serial.prediction); + } + } + + Ok(SplitSampleSet { + total: total_samples, + stage0: stage0_samples, + stage1: stage1_samples, + serial_total: serial_total_samples, + serial_stage0: serial_stage0_samples, + serial_stage1: serial_stage1_samples, + boundary_payload_bytes, + serial_boundary_payload_bytes, + first_prediction: first_prediction.unwrap_or_default(), + first_serial_prediction: first_serial_prediction.unwrap_or_default(), + }) +} + +fn measure_split_batched( + session0: &mut StageSession, + session1: &mut StageSession, + base0: u64, + base1: u64, + verify_tokens: &[i32], +) -> Result { + session0 + .trim_session(base0) + .context("failed to trim split stage 0 before verify")?; + session1 + .trim_session(base1) + .context("failed to trim split stage 1 before verify")?; + + let total_start = Instant::now(); + let stage0_start = Instant::now(); + let (_stage0_prediction, boundary) = session0 + .verify_tokens_frame_sampled(verify_tokens, Some(&SamplingConfig::default()), None, 0) + .context("in-process split stage 0 VerifySpan failed")?; + let stage0 = stage0_start.elapsed(); + let boundary_payload_bytes = boundary.payload.len(); + if boundary_payload_bytes == 0 { + bail!("in-process split stage 0 produced an empty VerifySpan activation frame"); + } + + let stage1_start = Instant::now(); + let prediction = session1 + .verify_tokens_frame_sampled( + verify_tokens, + Some(&SamplingConfig::default()), + Some(&boundary), + 0, + ) + .context("in-process split stage 1 VerifySpan failed")? + .0; + let stage1 = stage1_start.elapsed(); + Ok(SplitSample { + total: total_start.elapsed(), + stage0, + stage1, + boundary_payload_bytes, + prediction, + }) +} + +fn measure_split_serial( + session0: &mut StageSession, + session1: &mut StageSession, + base0: u64, + base1: u64, + verify_tokens: &[i32], +) -> Result { + session0 + .trim_session(base0) + .context("failed to trim split stage 0 before serial verify")?; + session1 + .trim_session(base1) + .context("failed to trim split stage 1 before serial verify")?; + + let total_start = Instant::now(); + let mut stage0_total = Duration::ZERO; + let mut stage1_total = Duration::ZERO; + let mut boundary_payload_bytes = 0usize; + let mut prediction = Vec::with_capacity(verify_tokens.len() + 3); + let mut last_draft = None; + + for token_id in verify_tokens { + let stage0_start = Instant::now(); + let (_stage0_prediction, _stage0_draft, boundary) = session0 + .decode_step_frame_sampled_mtp_n1(*token_id, Some(&SamplingConfig::default()), None, 0) + .context("in-process split stage 0 serial decode failed")?; + stage0_total += stage0_start.elapsed(); + if boundary.payload.is_empty() { + bail!("in-process split stage 0 produced an empty serial activation frame"); + } + boundary_payload_bytes += boundary.payload.len(); + + let stage1_start = Instant::now(); + let (predicted, native_mtp, _output) = session1 + .decode_step_frame_sampled_mtp_n1( + *token_id, + Some(&SamplingConfig::default()), + Some(&boundary), + 0, + ) + .context("in-process split stage 1 serial decode failed")?; + stage1_total += stage1_start.elapsed(); + if predicted >= 0 { + prediction.push(predicted); + } + last_draft = native_mtp; + } + + if let Some(draft) = last_draft { + prediction.push(draft.token_id); + prediction.push(i32::try_from(draft.proposal_compute_us.max(0)).unwrap_or(i32::MAX)); + } + + Ok(SplitSample { + total: total_start.elapsed(), + stage0: stage0_total, + stage1: stage1_total, + boundary_payload_bytes, + prediction, + }) +} + +fn serial_decode_mtp_n1(session: &mut StageSession, verify_tokens: &[i32]) -> Result> { + let mut predicted_tokens = Vec::with_capacity(verify_tokens.len() + 3); + let mut last_draft = None; + for token_id in verify_tokens { + let (predicted, native_mtp, _frame) = session + .decode_step_frame_sampled_mtp_n1(*token_id, Some(&SamplingConfig::default()), None, 0) + .context("serial native MTP n=1 decode failed")?; + if predicted >= 0 { + predicted_tokens.push(predicted); + } + last_draft = native_mtp; + } + if let Some(draft) = last_draft { + predicted_tokens.push(draft.token_id); + predicted_tokens.push(i32::try_from(draft.proposal_compute_us.max(0)).unwrap_or(i32::MAX)); + } + Ok(predicted_tokens) +} + +#[derive(Debug)] +struct SplitSample { + total: Duration, + stage0: Duration, + stage1: Duration, + boundary_payload_bytes: usize, + prediction: Vec, +} + +#[derive(Debug)] +struct SplitSampleSet { + total: Vec, + stage0: Vec, + stage1: Vec, + serial_total: Vec, + serial_stage0: Vec, + serial_stage1: Vec, + boundary_payload_bytes: usize, + serial_boundary_payload_bytes: usize, + first_prediction: Vec, + first_serial_prediction: Vec, +} + +fn split_report( + split_layer: u32, + samples: SplitSampleSet, + full_batched_avg_us: f64, +) -> Result { + let total = timing_stats(&samples.total)?; + let serial_total = timing_stats(&samples.serial_total)?; + let total_avg = total.avg_us; + let serial_total_avg = serial_total.avg_us; + Ok(SplitInprocessReport { + split_layer, + boundary_payload_bytes: samples.boundary_payload_bytes, + serial_boundary_payload_bytes: samples.serial_boundary_payload_bytes, + stage0: timing_stats(&samples.stage0)?, + stage1: timing_stats(&samples.stage1)?, + serial_stage0: timing_stats(&samples.serial_stage0)?, + serial_stage1: timing_stats(&samples.serial_stage1)?, + total, + serial_total, + total_token_per_sec: verified_tokens_per_sec(total_avg, 2), + serial_total_token_per_sec: verified_tokens_per_sec(serial_total_avg, 2), + total_avg_vs_full_batched_avg: total_avg / full_batched_avg_us, + total_avg_vs_serial_total_avg: total_avg / serial_total_avg, + diagnostics: split_timing_diagnostics(&samples)?, + first_prediction: samples.first_prediction, + first_serial_prediction: samples.first_serial_prediction, + }) +} + +fn split_timing_diagnostics(samples: &SplitSampleSet) -> Result { + Ok(SplitTimingDiagnostics { + batched_total: timing_shape(&samples.total)?, + batched_stage0: timing_shape(&samples.stage0)?, + batched_stage1: timing_shape(&samples.stage1)?, + serial_total: timing_shape(&samples.serial_total)?, + serial_stage0: timing_shape(&samples.serial_stage0)?, + serial_stage1: timing_shape(&samples.serial_stage1)?, + }) +} + +#[derive(Debug)] +struct SampleSet { + batched: Vec, + serial: Vec, + first_batched_prediction: Vec, + first_serial_prediction: Vec, +} + +impl SampleSet { + fn batched_avg_us(&self) -> Result { + Ok(timing_stats(&self.batched)?.avg_us) + } +} + +fn build_report( + args: VerifySpanLocalArgs, + full: FullModelSamples, + split_inprocess_width2: Option, +) -> Result { + let batched_width2 = timing_stats(&full.samples.batched)?; + let serial_two_decode_mtp_n1 = timing_stats(&full.samples.serial)?; + let batched_avg = batched_width2.avg_us; + let serial_avg = serial_two_decode_mtp_n1.avg_us; + Ok(VerifySpanLocalReport { + mode: "verify-span-local", + model_path: args.model_path, + layer_end: args.layer_end, + split_layer: args.split_layer, + ctx_size: args.ctx_size, + n_gpu_layers: args.n_gpu_layers, + n_batch: args.n_batch, + n_ubatch: args.n_ubatch, + cache_type_k: args.cache_type_k, + cache_type_v: args.cache_type_v, + prompt_token_count: full.tokens.len(), + verify_tokens: full.verify_tokens, + warmup: args.warmup, + iterations: args.iterations, + batched_width2, + serial_two_decode_mtp_n1, + split_inprocess_width2, + batched_avg_vs_serial_avg: batched_avg / serial_avg, + batched_token_per_sec: verified_tokens_per_sec(batched_avg, 2), + serial_token_per_sec: verified_tokens_per_sec(serial_avg, 2), + first_batched_prediction: full.samples.first_batched_prediction, + first_serial_prediction: full.samples.first_serial_prediction, + }) +} + +fn timing_stats(samples: &[Duration]) -> Result { + if samples.is_empty() { + bail!("cannot summarize empty timing samples"); + } + let mut micros = samples.iter().map(Duration::as_micros).collect::>(); + micros.sort_unstable(); + let total_us = micros.iter().sum::(); + let avg_us = total_us as f64 / micros.len() as f64; + Ok(TimingStats { + count: micros.len(), + total_us, + avg_us, + min_us: micros[0], + p50_us: percentile(µs, 0.50), + p95_us: percentile(µs, 0.95), + max_us: *micros.last().context("missing max timing sample")?, + }) +} + +fn timing_shape(samples: &[Duration]) -> Result { + if samples.is_empty() { + bail!("cannot summarize empty timing shape"); + } + let split_at = (samples.len() / 2).max(1); + let (first_half, second_half) = samples.split_at(split_at); + let second_half = if second_half.is_empty() { + first_half + } else { + second_half + }; + let first_half_stats = timing_stats(first_half)?; + let second_half_stats = timing_stats(second_half)?; + let samples_us = samples.iter().map(Duration::as_micros).collect::>(); + Ok(TimingShape { + second_half_avg_vs_first_half_avg: second_half_stats.avg_us / first_half_stats.avg_us, + first_half: first_half_stats, + second_half: second_half_stats, + first_sample_us: samples_us[0], + last_sample_us: *samples_us.last().context("missing timing shape sample")?, + samples_us, + }) +} + +fn percentile(sorted_micros: &[u128], percentile: f64) -> u128 { + let last_index = sorted_micros.len().saturating_sub(1); + let index = (last_index as f64 * percentile).round() as usize; + sorted_micros[index.min(last_index)] +} + +fn verified_tokens_per_sec(avg_us: f64, token_count: usize) -> f64 { + token_count as f64 * 1_000_000.0 / avg_us +} + +fn model_description(config: &RuntimeConfig) -> String { + format!( + "layers={}..{} ctx={} n_gpu_layers={}", + config.layer_start, config.layer_end, config.ctx_size, config.n_gpu_layers + ) +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use super::{percentile, timing_shape, timing_stats, verified_tokens_per_sec}; + + #[test] + fn timing_stats_sorts_and_summarizes_microseconds() { + let stats = timing_stats(&[ + Duration::from_micros(30), + Duration::from_micros(10), + Duration::from_micros(20), + ]) + .unwrap(); + + assert_eq!(stats.count, 3); + assert_eq!(stats.total_us, 60); + assert_eq!(stats.min_us, 10); + assert_eq!(stats.p50_us, 20); + assert_eq!(stats.max_us, 30); + } + + #[test] + fn percentile_clamps_to_last_sample() { + assert_eq!(percentile(&[10, 20, 30], 1.0), 30); + } + + #[test] + fn token_rate_uses_verified_token_count() { + assert_eq!(verified_tokens_per_sec(20_000.0, 2), 100.0); + } + + #[test] + fn timing_shape_reports_half_drift_in_sample_order() { + let shape = timing_shape(&[ + Duration::from_micros(10), + Duration::from_micros(20), + Duration::from_micros(30), + Duration::from_micros(50), + ]) + .unwrap(); + + assert_eq!(shape.first_sample_us, 10); + assert_eq!(shape.last_sample_us, 50); + assert_eq!(shape.samples_us, vec![10, 20, 30, 50]); + assert_eq!(shape.first_half.avg_us, 15.0); + assert_eq!(shape.second_half.avg_us, 40.0); + assert_eq!(shape.second_half_avg_vs_first_half_avg, 40.0 / 15.0); + } +} diff --git a/crates/skippy-correctness/Cargo.toml b/crates/skippy-correctness/Cargo.toml index 56f0a2a5d4..7152af38d7 100644 --- a/crates/skippy-correctness/Cargo.toml +++ b/crates/skippy-correctness/Cargo.toml @@ -12,6 +12,7 @@ skippy-runtime = { path = "../skippy-runtime" } model-artifact = { path = "../model-artifact" } model-hf = { path = "../model-hf" } model-ref = { path = "../model-ref" } +reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] } serde.workspace = true serde_json.workspace = true sha2 = "0.10" diff --git a/crates/skippy-correctness/README.md b/crates/skippy-correctness/README.md index 5bc3be9ddb..0351b675de 100644 --- a/crates/skippy-correctness/README.md +++ b/crates/skippy-correctness/README.md @@ -83,6 +83,36 @@ skippy-correctness state-handoff \ --cache-hit-repeats 3 \ --n-gpu-layers=-1 \ --report-out reports/state-handoff.json + +skippy-correctness native-mtp-open-ai-ab \ + --model model.gguf \ + --model-id org/repo:Q4_K_M \ + --prompt 'Write a Python function named add_one that returns x + 1.' \ + --layer-end 48 \ + --split-layer 24 \ + --ctx-size 128 \ + --n-batch 128 \ + --n-ubatch 128 \ + --n-gpu-layers 999 \ + --activation-width 2048 \ + --activation-wire-dtype f16 \ + --max-tokens 12 \ + --report-out reports/native-mtp-openai-ab.json + +skippy-correctness native-mtp-open-ai-ab \ + --model /Volumes/External/models/huggingface/.../model.gguf \ + --model-id org/repo:Q4_K_M \ + --stage1-model /Volumes/models/huggingface/.../model.gguf \ + --stage1-ssh-host micstudio \ + --stage1-remote-workdir /Users/micn/src/mesh-llm-codex \ + --stage1-remote-stage-server-bin target/debug/skippy-server \ + --openai-bind-addr 192.168.0.5:19170 \ + --stage0-bind-addr 192.168.0.5:19171 \ + --stage0-endpoint-addr 192.168.0.5:19171 \ + --stage1-bind-addr 192.168.0.10:19172 \ + --stage1-endpoint-addr 192.168.0.10:19172 \ + --layer-end 48 \ + --split-layer 24 ``` All commands emit JSON, optionally write the same JSON with `--report-out`, and @@ -178,6 +208,26 @@ same activation/cache contracts without requiring a monolithic full GGUF. `--cache-hit-repeats` to repeatedly attach the exported state and decode the same continuation, producing a recompute-vs-cache-hit speedup estimate. Use `--allow-mismatch` only for diagnostic payloads such as recurrent-only. +- `native-mtp-open-ai-ab` launches a real two-stage `skippy-server + serve-binary` split three times through the embedded OpenAI frontend. The + baseline run sets `SKIPPY_NATIVE_MTP_ENABLED=0`, the n=1 run sets + `SKIPPY_NATIVE_MTP_BATCHED_VERIFY=0`, and the batched run uses the default + transactional verification path. The command fails unless all responses are + HTTP 200, baseline/n=1/batched output content is byte-identical, native MTP + metrics are observed for n=1 and batched runs, and the batched run emits + `stage.openai_native_mtp_verify` events. +- For lab split config generation, use `--stage0-endpoint-addr` and + `--stage1-endpoint-addr` when the address a peer should dial differs from the + local bind. Use `--stage0-model` and `--stage1-model` when the same model is + mounted at different paths on each host. +- For harness-owned lab splits, pass `--stage1-ssh-host`, + `--stage1-remote-workdir`, and `--stage1-remote-stage-server-bin`. The harness + copies the generated stage-1 config/topology to the remote host, launches + `serve-binary` there for each baseline/n=1/batched case, waits for the binary + endpoint, and collects the remote log back into the report directory. +- For manually managed lab splits, pass `--external-stage1`. The harness writes + the stage configs and waits for the configured stage-1 endpoint, but it does + not start or stop the remote process. - Requires a built `skippy-server` binary for binary transport checks. - Uses the same llama-backed runtime ABI as the server. - The default build statically links llama from diff --git a/crates/skippy-correctness/src/cli.rs b/crates/skippy-correctness/src/cli.rs index a50a64dde5..459d9528ec 100644 --- a/crates/skippy-correctness/src/cli.rs +++ b/crates/skippy-correctness/src/cli.rs @@ -17,6 +17,7 @@ pub enum CommandKind { SplitScan(SplitScanArgs), DtypeMatrix(DtypeMatrixArgs), StateHandoff(StateHandoffArgs), + NativeMtpOpenAiAb(Box), } #[derive(Args, Clone)] @@ -76,6 +77,15 @@ pub struct ServerArgs { pub max_inflight: usize, } +#[derive(Args, Clone, Copy)] +pub struct NativeMtpArgs { + #[arg( + long, + help = "Fail the correctness run unless the final stage returns a native MTP draft sideband" + )] + pub require_native_mtp_draft: bool, +} + #[derive(Args)] pub struct OutputArgs { #[arg(long)] @@ -89,6 +99,8 @@ pub struct SingleStepArgs { #[command(flatten)] pub server: ServerArgs, #[command(flatten)] + pub native_mtp: NativeMtpArgs, + #[command(flatten)] pub output: OutputArgs, #[arg(long, default_value_t = 15)] pub split_layer: u32, @@ -107,6 +119,8 @@ pub struct ChainArgs { #[command(flatten)] pub server: ServerArgs, #[command(flatten)] + pub native_mtp: NativeMtpArgs, + #[command(flatten)] pub output: OutputArgs, #[arg(long, default_value = "10,20")] pub splits: String, @@ -127,6 +141,8 @@ pub struct SplitScanArgs { #[command(flatten)] pub server: ServerArgs, #[command(flatten)] + pub native_mtp: NativeMtpArgs, + #[command(flatten)] pub output: OutputArgs, #[arg(long, default_value = "1..30")] pub splits: String, @@ -145,6 +161,8 @@ pub struct DtypeMatrixArgs { #[command(flatten)] pub server: ServerArgs, #[command(flatten)] + pub native_mtp: NativeMtpArgs, + #[command(flatten)] pub output: OutputArgs, #[arg(long, default_value_t = 15)] pub split_layer: u32, @@ -200,6 +218,58 @@ pub struct StateHandoffArgs { pub allow_mismatch: bool, } +#[derive(Args)] +pub struct NativeMtpOpenAiAbArgs { + #[command(flatten)] + pub runtime: RuntimeArgs, + #[command(flatten)] + pub server: ServerArgs, + #[command(flatten)] + pub output: OutputArgs, + #[arg(long, default_value_t = 24)] + pub split_layer: u32, + #[arg(long, default_value = "127.0.0.1:19170")] + pub openai_bind_addr: SocketAddr, + #[arg(long, default_value = "127.0.0.1:19171")] + pub stage0_bind_addr: SocketAddr, + #[arg(long)] + pub stage0_endpoint_addr: Option, + #[arg(long, default_value = "127.0.0.1:19172")] + pub stage1_bind_addr: SocketAddr, + #[arg(long)] + pub stage1_endpoint_addr: Option, + #[arg(long)] + pub stage0_model: Option, + #[arg(long)] + pub stage1_model: Option, + #[arg(long)] + pub case_root: Option, + #[arg(long)] + pub external_stage1: bool, + #[arg(long)] + pub stage1_ssh_host: Option, + #[arg(long)] + pub stage1_remote_stage_server_bin: Option, + #[arg(long, default_value = "/tmp/skippy-native-mtp-openai-ab")] + pub stage1_remote_root: String, + #[arg(long)] + pub stage1_remote_workdir: Option, + #[arg(long, default_value_t = 10)] + pub batched_port_offset: u16, + #[arg(long, default_value_t = 2048)] + pub activation_width: i32, + #[arg(long, default_value = "f16")] + pub activation_wire_dtype: String, + #[arg(long, default_value_t = 12)] + pub max_tokens: u32, + #[arg(long, default_value_t = 60)] + pub request_timeout_secs: u64, + #[arg(long)] + pub allow_missing_batched_events: bool, + #[arg(long)] + pub allow_mismatch: bool, +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] #[value(rename_all = "kebab-case")] pub enum StatePayloadKind { diff --git a/crates/skippy-correctness/src/direct_return.rs b/crates/skippy-correctness/src/direct_return.rs deleted file mode 100644 index c2398de787..0000000000 --- a/crates/skippy-correctness/src/direct_return.rs +++ /dev/null @@ -1,160 +0,0 @@ -use std::{ - collections::HashMap, - io, - net::{SocketAddr, TcpListener, TcpStream}, - sync::{Arc, Mutex, mpsc}, - thread, - time::Duration, -}; - -use anyhow::{Context, Result, anyhow, bail}; -use skippy_protocol::binary::{StageReply, WireMessageKind, WireReplyKind, recv_reply, send_ready}; - -#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] -struct DirectReturnKey { - request_id: u64, - session_id: u64, -} - -type DirectReturnResult = Result; -type DirectReturnSender = mpsc::Sender; -type DirectReturnWaiters = Arc>>; - -pub(crate) struct CorrectnessDirectReturnServer { - local_addr: SocketAddr, - waiters: DirectReturnWaiters, -} - -impl CorrectnessDirectReturnServer { - pub(crate) fn start(bind_addr: &str) -> Result { - let listener = TcpListener::bind(bind_addr) - .with_context(|| format!("bind correctness direct-return listener {bind_addr}"))?; - let local_addr = listener - .local_addr() - .context("read correctness direct-return listener address")?; - let waiters = Arc::new(Mutex::new(HashMap::new())); - let thread_waiters = waiters.clone(); - thread::spawn(move || { - for stream in listener.incoming() { - match stream { - Ok(stream) => { - let waiters = thread_waiters.clone(); - thread::spawn(move || { - if let Err(error) = - handle_correctness_direct_return_connection(waiters, stream) - { - eprintln!("correctness direct-return connection failed: {error:#}"); - } - }); - } - Err(error) if error.kind() == io::ErrorKind::Interrupted => continue, - Err(error) => { - eprintln!("correctness direct-return listener failed: {error}"); - break; - } - } - } - }); - Ok(Self { - local_addr, - waiters, - }) - } - - pub(crate) fn endpoint(&self) -> String { - self.local_addr.to_string() - } - - pub(crate) fn register( - &self, - request_id: u64, - session_id: u64, - ) -> Result { - let key = DirectReturnKey { - request_id, - session_id, - }; - let (sender, receiver) = mpsc::channel(); - self.waiters - .lock() - .map_err(|_| anyhow!("correctness direct-return hub lock poisoned"))? - .insert(key, sender); - Ok(CorrectnessDirectReturnReceiver { - key, - waiters: self.waiters.clone(), - receiver, - }) - } -} - -pub(crate) struct CorrectnessDirectReturnReceiver { - key: DirectReturnKey, - waiters: DirectReturnWaiters, - receiver: mpsc::Receiver, -} - -impl CorrectnessDirectReturnReceiver { - pub(crate) fn recv_expected(&self, expected: WireReplyKind) -> Result { - let reply = self - .receiver - .recv_timeout(Duration::from_secs(300)) - .context("timed out waiting for correctness direct prediction return")? - .map_err(|error| anyhow!(error))?; - if reply.kind != expected { - bail!( - "expected {expected:?} direct prediction return, got {:?}", - reply.kind - ); - } - Ok(reply) - } -} - -impl Drop for CorrectnessDirectReturnReceiver { - fn drop(&mut self) { - if let Ok(mut waiters) = self.waiters.lock() { - waiters.remove(&self.key); - } - } -} - -fn handle_correctness_direct_return_connection( - waiters: DirectReturnWaiters, - mut stream: TcpStream, -) -> Result<()> { - send_ready(&mut stream).context("send correctness direct-return ready")?; - let open = skippy_protocol::binary::read_stage_message(&mut stream, 0) - .context("read correctness direct-return open")?; - if open.kind != WireMessageKind::PredictionReturnOpen { - bail!("expected prediction-return-open message"); - } - let key = DirectReturnKey { - request_id: open.request_id, - session_id: open.session_id, - }; - let sender = waiters - .lock() - .map_err(|_| anyhow!("correctness direct-return hub lock poisoned"))? - .get(&key) - .cloned() - .ok_or_else(|| { - anyhow!( - "no correctness direct-return waiter for request {}", - key.request_id - ) - })?; - loop { - match recv_reply(&mut stream) { - Ok(reply) => { - if sender.send(Ok(reply)).is_err() { - return Ok(()); - } - } - Err(error) if error.kind() == io::ErrorKind::UnexpectedEof => return Ok(()), - Err(error) => { - let _ = sender.send(Err(error.to_string())); - return Err(error).context("read correctness direct prediction return"); - } - } - } -} diff --git a/crates/skippy-correctness/src/main.rs b/crates/skippy-correctness/src/main.rs index d6e0bb089c..efccd869d8 100644 --- a/crates/skippy-correctness/src/main.rs +++ b/crates/skippy-correctness/src/main.rs @@ -1,5 +1,5 @@ mod cli; -mod direct_return; +mod native_mtp_openai; mod report; mod runner; mod support; @@ -9,6 +9,7 @@ use clap::Parser; use crate::{ cli::{Cli, CommandKind}, + native_mtp_openai::native_mtp_openai_ab, runner::{chain, dtype_matrix, single_step, split_scan, state_handoff}, }; @@ -19,5 +20,6 @@ fn main() -> Result<()> { CommandKind::SplitScan(args) => split_scan(args), CommandKind::DtypeMatrix(args) => dtype_matrix(args), CommandKind::StateHandoff(args) => state_handoff(args), + CommandKind::NativeMtpOpenAiAb(args) => native_mtp_openai_ab(*args), } } diff --git a/crates/skippy-correctness/src/native_mtp_openai/metrics.rs b/crates/skippy-correctness/src/native_mtp_openai/metrics.rs new file mode 100644 index 0000000000..2a89fc8032 --- /dev/null +++ b/crates/skippy-correctness/src/native_mtp_openai/metrics.rs @@ -0,0 +1,293 @@ +use std::{fs, path::Path}; + +use anyhow::{Context, Result}; +use serde_json::Value; + +use crate::report::NativeMtpOpenAiMetricsReport; + +pub(super) fn read_metrics( + stage0_log: &Path, + stage1_log: &Path, +) -> Result { + let mut metrics = NativeMtpOpenAiMetricsReport::default(); + for path in [stage0_log, stage1_log] { + let text = fs::read_to_string(path) + .with_context(|| format!("failed to read {}", path.display()))?; + metrics.fatal_error_events += count_fatal_log_lines(&text); + for line in text.lines() { + let Ok(event) = serde_json::from_str::(line) else { + continue; + }; + apply_telemetry_event(&mut metrics, &event); + } + } + Ok(metrics) +} + +fn apply_telemetry_event(metrics: &mut NativeMtpOpenAiMetricsReport, event: &Value) { + let Some(name) = event.get("event").and_then(Value::as_str) else { + return; + }; + let attrs = event.get("attributes").unwrap_or(&Value::Null); + match name { + "stage.openai_decode_token" => { + metrics.decode_token_events += 1; + apply_native_mtp_verification(metrics, attrs); + } + "stage.openai_native_mtp_verify" => { + metrics.batched_verify_events += 1; + match apply_native_mtp_verification(metrics, attrs) { + Some("accepted") => metrics.batched_accepted_events += 1, + Some("rejected") => metrics.batched_rejected_events += 1, + _ => {} + } + } + "stage.openai_decode" | "stage.openai_generation_summary" => { + apply_generation_summary(metrics, attrs); + } + _ => {} + } +} + +fn apply_generation_summary(metrics: &mut NativeMtpOpenAiMetricsReport, attrs: &Value) { + metrics.native_mtp_enabled |= attr_bool(attrs, "llama_stage.native_mtp.enabled"); + if let Some(drafted) = attr_u64(attrs, "llama_stage.native_mtp.drafted") { + metrics.drafted_tokens = drafted; + } + if let Some(accepted) = attr_u64(attrs, "llama_stage.native_mtp.accepted") { + metrics.accepted_tokens = accepted; + } + if let Some(rejected) = attr_u64(attrs, "llama_stage.native_mtp.rejected") { + metrics.rejected_tokens = rejected; + } + if let Some(pending) = attr_u64(attrs, "llama_stage.native_mtp.pending") { + metrics.pending_tokens = pending; + } + if let Some(verifications) = attr_u64(attrs, "llama_stage.native_mtp.verifications") { + metrics.verification_count = verifications; + } + if let Some(proposal_compute_us) = attr_i64(attrs, "llama_stage.native_mtp.proposal_compute_us") + { + metrics.proposal_compute_us = proposal_compute_us; + } + if let Some(verification_compute_us) = + attr_i64(attrs, "llama_stage.native_mtp.verification_compute_us") + { + metrics.verification_compute_us = verification_compute_us; + } + if let Some(accept_rate) = attrs + .get("llama_stage.native_mtp.accept_rate") + .and_then(Value::as_f64) + { + metrics.accept_rate = accept_rate; + } + if let Some(value) = attr_f64( + attrs, + "llama_stage.native_mtp.batched.accepted_verify_elapsed_ms", + ) { + metrics.batched_accepted_verify_elapsed_ms = value; + } + if let Some(value) = attr_f64( + attrs, + "llama_stage.native_mtp.batched.accepted_verify_elapsed_avg_ms", + ) { + metrics.batched_accepted_verify_avg_ms = value; + } + if let Some(value) = attr_f64( + attrs, + "llama_stage.native_mtp.batched.rejected_verify_elapsed_ms", + ) { + metrics.batched_rejected_verify_elapsed_ms = value; + } + if let Some(value) = attr_f64( + attrs, + "llama_stage.native_mtp.batched.rejected_verify_elapsed_avg_ms", + ) { + metrics.batched_rejected_verify_avg_ms = value; + } + if let Some(value) = attr_u64(attrs, "llama_stage.native_mtp.batched.consumed_positions") { + metrics.batched_consumed_positions = value; + } + if let Some(value) = attr_u64(attrs, "llama_stage.native_mtp.batched.committed_positions") { + metrics.batched_committed_positions = value; + } + if let Some(value) = attr_u64(attrs, "llama_stage.native_mtp.batched.trim_count") { + metrics.batched_trim_count = value; + } + if let Some(value) = attr_f64(attrs, "llama_stage.native_mtp.batched.trim_elapsed_ms") { + metrics.batched_trim_elapsed_ms = value; + } + if let Some(value) = attr_f64(attrs, "llama_stage.native_mtp.batched.trim_local_ms") { + metrics.batched_trim_local_ms = value; + } + if let Some(value) = attr_f64( + attrs, + "llama_stage.native_mtp.batched.trim_downstream_write_ms", + ) { + metrics.batched_trim_downstream_write_ms = value; + } + if let Some(value) = attr_f64( + attrs, + "llama_stage.native_mtp.batched.trim_downstream_wait_ms", + ) { + metrics.batched_trim_downstream_wait_ms = value; + } +} + +fn apply_native_mtp_verification<'a>( + metrics: &mut NativeMtpOpenAiMetricsReport, + attrs: &'a Value, +) -> Option<&'a str> { + let verification = attrs + .get("llama_stage.native_mtp.verification") + .and_then(Value::as_str)?; + match verification { + "accepted" => { + metrics.native_mtp_enabled = true; + metrics.accepted_tokens += 1; + metrics.verification_count += 1; + } + "rejected" => { + metrics.native_mtp_enabled = true; + metrics.rejected_tokens += 1; + metrics.verification_count += 1; + } + _ => {} + } + Some(verification) +} + +fn count_fatal_log_lines(text: &str) -> u64 { + text.lines() + .filter(|line| { + line.contains("panicked") + || line.contains("service_unavailable") + || line.contains("llama_decode failed for MTP sidecar sync") + }) + .count() as u64 +} + +fn attr_bool(attrs: &Value, key: &str) -> bool { + attrs.get(key).and_then(Value::as_bool).unwrap_or(false) +} + +fn attr_u64(attrs: &Value, key: &str) -> Option { + attrs.get(key).and_then(Value::as_u64) +} + +fn attr_i64(attrs: &Value, key: &str) -> Option { + attrs.get(key).and_then(Value::as_i64) +} + +fn attr_f64(attrs: &Value, key: &str) -> Option { + attrs.get(key).and_then(Value::as_f64) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn telemetry_parser_extracts_native_mtp_counts() { + let mut metrics = NativeMtpOpenAiMetricsReport::default(); + apply_telemetry_event( + &mut metrics, + &json!({ + "event": "stage.openai_native_mtp_verify", + "attributes": { + "llama_stage.native_mtp.verification": "accepted", + } + }), + ); + apply_telemetry_event( + &mut metrics, + &json!({ + "event": "stage.openai_decode_token", + "attributes": {} + }), + ); + apply_telemetry_event( + &mut metrics, + &json!({ + "event": "stage.openai_decode", + "attributes": { + "llama_stage.native_mtp.enabled": true, + "llama_stage.native_mtp.drafted": 4, + "llama_stage.native_mtp.accepted": 3, + "llama_stage.native_mtp.rejected": 1, + "llama_stage.native_mtp.pending": 0, + "llama_stage.native_mtp.verifications": 4, + "llama_stage.native_mtp.accept_rate": 0.75, + "llama_stage.native_mtp.proposal_compute_us": 11, + "llama_stage.native_mtp.verification_compute_us": 22, + "llama_stage.native_mtp.batched.accepted_verify_elapsed_ms": 30.0, + "llama_stage.native_mtp.batched.accepted_verify_elapsed_avg_ms": 10.0, + "llama_stage.native_mtp.batched.rejected_verify_elapsed_ms": 12.0, + "llama_stage.native_mtp.batched.rejected_verify_elapsed_avg_ms": 12.0, + "llama_stage.native_mtp.batched.consumed_positions": 8, + "llama_stage.native_mtp.batched.committed_positions": 7, + "llama_stage.native_mtp.batched.trim_count": 1, + "llama_stage.native_mtp.batched.trim_elapsed_ms": 4.0, + "llama_stage.native_mtp.batched.trim_local_ms": 1.0, + "llama_stage.native_mtp.batched.trim_downstream_write_ms": 0.5, + "llama_stage.native_mtp.batched.trim_downstream_wait_ms": 2.5, + } + }), + ); + + assert!(metrics.native_mtp_enabled); + assert_eq!(metrics.drafted_tokens, 4); + assert_eq!(metrics.accepted_tokens, 3); + assert_eq!(metrics.rejected_tokens, 1); + assert_eq!(metrics.verification_count, 4); + assert_eq!(metrics.accept_rate, 0.75); + assert_eq!(metrics.proposal_compute_us, 11); + assert_eq!(metrics.verification_compute_us, 22); + assert_eq!(metrics.batched_accepted_verify_elapsed_ms, 30.0); + assert_eq!(metrics.batched_accepted_verify_avg_ms, 10.0); + assert_eq!(metrics.batched_rejected_verify_elapsed_ms, 12.0); + assert_eq!(metrics.batched_rejected_verify_avg_ms, 12.0); + assert_eq!(metrics.batched_consumed_positions, 8); + assert_eq!(metrics.batched_committed_positions, 7); + assert_eq!(metrics.batched_trim_count, 1); + assert_eq!(metrics.batched_trim_elapsed_ms, 4.0); + assert_eq!(metrics.batched_trim_local_ms, 1.0); + assert_eq!(metrics.batched_trim_downstream_write_ms, 0.5); + assert_eq!(metrics.batched_trim_downstream_wait_ms, 2.5); + assert_eq!(metrics.batched_verify_events, 1); + assert_eq!(metrics.batched_accepted_events, 1); + assert_eq!(metrics.decode_token_events, 1); + } + + #[test] + fn telemetry_parser_counts_batched_rejections() { + let mut metrics = NativeMtpOpenAiMetricsReport::default(); + + apply_telemetry_event( + &mut metrics, + &json!({ + "event": "stage.openai_native_mtp_verify", + "attributes": { + "llama_stage.native_mtp.verification": "rejected", + } + }), + ); + + assert!(metrics.native_mtp_enabled); + assert_eq!(metrics.rejected_tokens, 1); + assert_eq!(metrics.verification_count, 1); + assert_eq!(metrics.batched_verify_events, 1); + assert_eq!(metrics.batched_rejected_events, 1); + } + + #[test] + fn fatal_counter_ignores_connection_retry_noise() { + let text = "\ +downstream connect retry: error=Connection refused (os error 61) +llama_decode failed for MTP sidecar sync +"; + assert_eq!(count_fatal_log_lines(text), 1); + } +} diff --git a/crates/skippy-correctness/src/native_mtp_openai/mod.rs b/crates/skippy-correctness/src/native_mtp_openai/mod.rs new file mode 100644 index 0000000000..519251546b --- /dev/null +++ b/crates/skippy-correctness/src/native_mtp_openai/mod.rs @@ -0,0 +1,416 @@ +mod metrics; +mod remote; +mod reporting; +mod stage_process; + +use std::{ + fs, + net::SocketAddr, + path::{Path, PathBuf}, + time::Duration, +}; + +use anyhow::{Context, Result, bail}; +use reqwest::blocking::Client; +use serde_json::{Value, json}; + +use crate::{ + cli::{NativeMtpOpenAiAbArgs, StageLoadMode}, + report::{NativeMtpOpenAiAbReport, NativeMtpOpenAiCaseReport}, + support::generate_run_id, +}; + +use metrics::read_metrics; +use reporting::{emit_report, status}; +use stage_process::{ + case_addr, spawn_stage, start_stage1, wait_openai_ready, wait_stage1_ready, write_stage_config, +}; + +pub(super) const BATCHED_VERIFY_ENV: &str = "SKIPPY_NATIVE_MTP_BATCHED_VERIFY"; +pub(super) const NATIVE_MTP_ENABLED_ENV: &str = "SKIPPY_NATIVE_MTP_ENABLED"; + +struct OpenAiCaseConfig { + case: &'static str, + native_mtp_enabled: bool, + batched_verify_enabled: bool, + run_id: String, + root: PathBuf, + openai_bind_addr: SocketAddr, + stage0_bind_addr: SocketAddr, + stage0_endpoint_addr: SocketAddr, + stage1_bind_addr: SocketAddr, + stage1_endpoint_addr: SocketAddr, +} + +struct OpenAiStageConfig<'a> { + run_id: &'a str, + model_id: &'a str, + model_path: &'a Path, + stage_id: &'a str, + stage_index: u32, + layer_start: u32, + layer_end: u32, + bind_addr: SocketAddr, + upstream: Option, + downstream: Option, +} + +pub fn native_mtp_openai_ab(args: NativeMtpOpenAiAbArgs) -> Result<()> { + if args.runtime.stage_load_mode != StageLoadMode::RuntimeSlice { + bail!("native-mtp-open-ai-ab currently supports --stage-load-mode runtime-slice only"); + } + if args.split_layer == 0 || args.split_layer >= args.runtime.layer_end { + bail!("split_layer must be greater than zero and less than layer_end"); + } + if args.external_stage1 && args.stage1_ssh_host.is_some() { + bail!("--external-stage1 cannot be combined with --stage1-ssh-host"); + } + if args.stage1_ssh_host.is_some() && args.stage1_remote_stage_server_bin.is_none() { + bail!("--stage1-remote-stage-server-bin is required with --stage1-ssh-host"); + } + + let model_id = args.runtime.model_id.clone().unwrap_or_else(|| { + args.runtime + .model + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("local-model") + .to_string() + }); + let root = args + .case_root + .clone() + .unwrap_or_else(|| std::env::temp_dir().join(generate_run_id())); + fs::create_dir_all(&root).with_context(|| format!("failed to create {}", root.display()))?; + let client = Client::builder() + .timeout(Duration::from_secs(args.request_timeout_secs.max(1))) + .build() + .context("failed to build HTTP client")?; + + let baseline = run_openai_case( + &args, + &client, + &model_id, + OpenAiCaseConfig { + case: "baseline", + native_mtp_enabled: false, + batched_verify_enabled: false, + run_id: format!("{}-baseline", generate_run_id()), + root: root.join("baseline"), + openai_bind_addr: case_addr(args.openai_bind_addr, args.batched_port_offset, 0)?, + stage0_bind_addr: case_addr(args.stage0_bind_addr, args.batched_port_offset, 0)?, + stage0_endpoint_addr: case_addr( + args.stage0_endpoint_addr.unwrap_or(args.stage0_bind_addr), + args.batched_port_offset, + 0, + )?, + stage1_bind_addr: case_addr(args.stage1_bind_addr, args.batched_port_offset, 0)?, + stage1_endpoint_addr: case_addr( + args.stage1_endpoint_addr.unwrap_or(args.stage1_bind_addr), + args.batched_port_offset, + 0, + )?, + }, + )?; + let n1 = run_openai_case( + &args, + &client, + &model_id, + OpenAiCaseConfig { + case: "n1", + native_mtp_enabled: true, + batched_verify_enabled: false, + run_id: format!("{}-n1", generate_run_id()), + root: root.join("n1"), + openai_bind_addr: case_addr(args.openai_bind_addr, args.batched_port_offset, 1)?, + stage0_bind_addr: case_addr(args.stage0_bind_addr, args.batched_port_offset, 1)?, + stage0_endpoint_addr: case_addr( + args.stage0_endpoint_addr.unwrap_or(args.stage0_bind_addr), + args.batched_port_offset, + 1, + )?, + stage1_bind_addr: case_addr(args.stage1_bind_addr, args.batched_port_offset, 1)?, + stage1_endpoint_addr: case_addr( + args.stage1_endpoint_addr.unwrap_or(args.stage1_bind_addr), + args.batched_port_offset, + 1, + )?, + }, + )?; + let batched = run_openai_case( + &args, + &client, + &model_id, + OpenAiCaseConfig { + case: "batched", + native_mtp_enabled: true, + batched_verify_enabled: true, + run_id: format!("{}-batched", generate_run_id()), + root: root.join("batched"), + openai_bind_addr: case_addr(args.openai_bind_addr, args.batched_port_offset, 2)?, + stage0_bind_addr: case_addr(args.stage0_bind_addr, args.batched_port_offset, 2)?, + stage0_endpoint_addr: case_addr( + args.stage0_endpoint_addr.unwrap_or(args.stage0_bind_addr), + args.batched_port_offset, + 2, + )?, + stage1_bind_addr: case_addr(args.stage1_bind_addr, args.batched_port_offset, 2)?, + stage1_endpoint_addr: case_addr( + args.stage1_endpoint_addr.unwrap_or(args.stage1_bind_addr), + args.batched_port_offset, + 2, + )?, + }, + )?; + + let exact_content_match = baseline.content == n1.content && baseline.content == batched.content; + let batched_events_present = batched.metrics.batched_verify_events > 0; + let require_batched_events = !args.allow_missing_batched_events; + let matches = baseline.http_status == 200 + && n1.http_status == 200 + && batched.http_status == 200 + && exact_content_match + && !baseline.metrics.native_mtp_enabled + && n1.metrics.native_mtp_enabled + && batched.metrics.native_mtp_enabled + && baseline.metrics.fatal_error_events == 0 + && n1.metrics.fatal_error_events == 0 + && batched.metrics.fatal_error_events == 0 + && (!require_batched_events || batched_events_present); + + let report = NativeMtpOpenAiAbReport { + mode: "native-mtp-open-ai-ab", + status: status(matches), + model_id, + model_path: args.runtime.model.display().to_string(), + prompt: args.runtime.prompt, + max_tokens: args.max_tokens, + split_layer: args.split_layer, + layer_end: args.runtime.layer_end, + activation_width: args.activation_width, + activation_wire_dtype: args.activation_wire_dtype, + exact_content_match, + batched_events_required: require_batched_events, + batched_events_present, + matches, + baseline, + n1, + batched, + }; + emit_report(&report, args.output.report_out.as_deref())?; + if !report.matches && !args.allow_mismatch { + bail!("native MTP OpenAI n=1 and batched verification did not match"); + } + Ok(()) +} + +fn run_openai_case( + args: &NativeMtpOpenAiAbArgs, + client: &Client, + model_id: &str, + case: OpenAiCaseConfig, +) -> Result { + fs::create_dir_all(&case.root) + .with_context(|| format!("failed to create {}", case.root.display()))?; + let stage0_config_path = case.root.join("stage0.json"); + let stage1_config_path = case.root.join("stage1.json"); + let topology_path = case.root.join("topology.json"); + let stage0_log = case.root.join("stage0.log"); + let stage1_log = case.root.join("stage1.log"); + let stage0_model_path = args.stage0_model.as_deref().unwrap_or(&args.runtime.model); + let stage1_model_path = args.stage1_model.as_deref().unwrap_or(&args.runtime.model); + + write_stage_config( + &stage0_config_path, + &stage_config_json( + args, + OpenAiStageConfig { + run_id: &case.run_id, + model_id, + model_path: stage0_model_path, + stage_id: "stage-0", + stage_index: 0, + layer_start: 0, + layer_end: args.split_layer, + bind_addr: case.stage0_bind_addr, + upstream: None, + downstream: Some(json!({ + "stage_id": "stage-1", + "stage_index": 1, + "endpoint": format!("tcp://{}", case.stage1_endpoint_addr), + })), + }, + ), + )?; + write_stage_config( + &stage1_config_path, + &stage_config_json( + args, + OpenAiStageConfig { + run_id: &case.run_id, + model_id, + model_path: stage1_model_path, + stage_id: "stage-1", + stage_index: 1, + layer_start: args.split_layer, + layer_end: args.runtime.layer_end, + bind_addr: case.stage1_bind_addr, + upstream: Some(json!({ + "stage_id": "stage-0", + "stage_index": 0, + "endpoint": format!("tcp://{}", case.stage0_endpoint_addr), + })), + downstream: None, + }, + ), + )?; + write_stage_config( + &topology_path, + &json!({ + "topology_id": "native-mtp-open-ai-ab", + "model_id": model_id, + "stages": [ + { + "stage_id": "stage-0", + "stage_index": 0, + "host": "localhost", + "endpoint": format!("tcp://{}", case.stage0_endpoint_addr), + "layer_start": 0, + "layer_end": args.split_layer, + "load_mode": "runtime-slice", + }, + { + "stage_id": "stage-1", + "stage_index": 1, + "host": "localhost", + "endpoint": format!("tcp://{}", case.stage1_endpoint_addr), + "layer_start": args.split_layer, + "layer_end": args.runtime.layer_end, + "load_mode": "runtime-slice", + }, + ], + }), + )?; + + let (stage1, stage1_launch) = start_stage1( + args, + &case.run_id, + &stage1_config_path, + &topology_path, + &stage1_log, + case.native_mtp_enabled, + case.batched_verify_enabled, + )?; + wait_stage1_ready( + &stage1, + case.stage1_endpoint_addr, + args.server.startup_timeout_secs, + ) + .context("stage 1 binary server did not become ready")?; + let stage0 = spawn_stage( + args, + &stage0_config_path, + Some(case.openai_bind_addr), + &topology_path, + &stage0_log, + case.native_mtp_enabled, + case.batched_verify_enabled, + )?; + wait_openai_ready( + client, + case.openai_bind_addr, + args.server.startup_timeout_secs, + ) + .context("stage 0 OpenAI server did not become ready")?; + + let response = client + .post(format!( + "http://{}/v1/chat/completions", + case.openai_bind_addr + )) + .json(&json!({ + "model": model_id, + "messages": [ + { + "role": "user", + "content": args.runtime.prompt, + }, + ], + "temperature": 0, + "max_tokens": args.max_tokens, + })) + .send() + .context("failed to send OpenAI chat completion request")?; + let http_status = response.status().as_u16(); + let body: Value = response.json().context("failed to parse OpenAI response")?; + let content = body + .pointer("/choices/0/message/content") + .and_then(Value::as_str) + .unwrap_or_default() + .to_string(); + let completion_tokens = body + .pointer("/usage/completion_tokens") + .and_then(Value::as_u64); + + drop(stage0); + stage1.stop_and_collect(&stage1_log)?; + + let metrics = read_metrics(&stage0_log, &stage1_log)?; + Ok(NativeMtpOpenAiCaseReport { + case: case.case, + native_mtp_enabled: case.native_mtp_enabled, + batched_verify_enabled: case.batched_verify_enabled, + http_status, + content, + completion_tokens, + openai_bind_addr: case.openai_bind_addr.to_string(), + stage0_bind_addr: case.stage0_bind_addr.to_string(), + stage0_endpoint_addr: case.stage0_endpoint_addr.to_string(), + stage1_bind_addr: case.stage1_bind_addr.to_string(), + stage1_endpoint_addr: case.stage1_endpoint_addr.to_string(), + stage0_config: stage0_config_path.display().to_string(), + stage1_config: stage1_config_path.display().to_string(), + topology_config: topology_path.display().to_string(), + stage0_log: stage0_log.display().to_string(), + stage1_log: stage1_log.display().to_string(), + stage1_launch_mode: stage1_launch.launch_mode.to_string(), + stage1_remote_config: stage1_launch.remote_config, + stage1_remote_topology: stage1_launch.remote_topology, + stage1_remote_log: stage1_launch.remote_log, + metrics, + }) +} + +fn stage_config_json(args: &NativeMtpOpenAiAbArgs, stage: OpenAiStageConfig<'_>) -> Value { + json!({ + "run_id": stage.run_id, + "topology_id": "native-mtp-open-ai-ab", + "model_id": stage.model_id, + "model_path": stage.model_path, + "stage_id": stage.stage_id, + "stage_index": stage.stage_index, + "layer_start": stage.layer_start, + "layer_end": stage.layer_end, + "ctx_size": args.runtime.ctx_size, + "lane_count": 1, + "n_batch": args.runtime.n_batch, + "n_ubatch": args.runtime.n_ubatch, + "n_gpu_layers": args.runtime.n_gpu_layers, + "cache_type_k": "f16", + "cache_type_v": "f16", + "flash_attn_type": protocol_flash_attn(args.runtime.flash_attn), + "filter_tensors_on_load": true, + "load_mode": "runtime-slice", + "bind_addr": stage.bind_addr, + "upstream": stage.upstream, + "downstream": stage.downstream, + }) +} + +fn protocol_flash_attn(value: crate::cli::FlashAttentionArg) -> &'static str { + match value { + crate::cli::FlashAttentionArg::Auto => "auto", + crate::cli::FlashAttentionArg::Disabled => "disabled", + crate::cli::FlashAttentionArg::Enabled => "enabled", + } +} diff --git a/crates/skippy-correctness/src/native_mtp_openai/remote.rs b/crates/skippy-correctness/src/native_mtp_openai/remote.rs new file mode 100644 index 0000000000..293cfc509c --- /dev/null +++ b/crates/skippy-correctness/src/native_mtp_openai/remote.rs @@ -0,0 +1,290 @@ +use std::{fs, path::Path, process::Command, thread, time::Duration}; + +use anyhow::{Context, Result, bail}; + +use crate::cli::NativeMtpOpenAiAbArgs; + +use super::stage_process::Stage1LaunchReport; + +pub(super) struct RemoteStageGuard { + host: String, + pid: Option, + remote_log: String, +} + +pub(super) fn spawn_remote_stage1( + args: &NativeMtpOpenAiAbArgs, + host: &str, + run_id: &str, + config_path: &Path, + topology_path: &Path, + native_mtp_enabled: bool, + batched_verify_enabled: bool, +) -> Result<(RemoteStageGuard, Stage1LaunchReport)> { + let remote_dir = format!( + "{}/{}", + args.stage1_remote_root.trim_end_matches('/'), + run_id + ); + let remote_config = format!("{remote_dir}/stage1.json"); + let remote_topology = format!("{remote_dir}/topology.json"); + let remote_log = format!("{remote_dir}/stage1.log"); + let remote_pid = format!("{remote_dir}/stage1.pid"); + let remote_bin = args + .stage1_remote_stage_server_bin + .as_deref() + .context("missing stage 1 remote binary")?; + + ssh_success(host, &format!("mkdir -p {}", shell_quote(&remote_dir))) + .with_context(|| format!("create remote stage directory on {host}"))?; + scp_to(host, config_path, &remote_config).with_context(|| { + format!( + "copy stage 1 config {} to {host}:{remote_config}", + config_path.display() + ) + })?; + scp_to(host, topology_path, &remote_topology).with_context(|| { + format!( + "copy topology {} to {host}:{remote_topology}", + topology_path.display() + ) + })?; + + let command = remote_stage_command(RemoteStageCommand { + workdir: args.stage1_remote_workdir.as_deref(), + remote_bin, + remote_config: &remote_config, + remote_topology: &remote_topology, + remote_log: &remote_log, + remote_pid: &remote_pid, + activation_width: args.activation_width, + activation_wire_dtype: &args.activation_wire_dtype, + native_mtp_enabled, + batched_verify_enabled, + }); + let output = Command::new("ssh") + .arg(host) + .arg(command) + .output() + .with_context(|| format!("start remote stage 1 on {host}"))?; + if !output.status.success() { + bail!( + "remote stage 1 start on {host} failed with status {}: {}", + output.status, + String::from_utf8_lossy(&output.stderr) + ); + } + let stdout = String::from_utf8_lossy(&output.stdout); + let pid = stdout + .lines() + .rev() + .find_map(|line| line.trim().parse::().ok()) + .with_context(|| format!("remote stage 1 on {host} did not print a pid"))?; + + Ok(( + RemoteStageGuard { + host: host.to_string(), + pid: Some(pid), + remote_log: remote_log.clone(), + }, + Stage1LaunchReport { + launch_mode: "ssh", + remote_config: Some(remote_config), + remote_topology: Some(remote_topology), + remote_log: Some(remote_log), + }, + )) +} + +struct RemoteStageCommand<'a> { + workdir: Option<&'a str>, + remote_bin: &'a str, + remote_config: &'a str, + remote_topology: &'a str, + remote_log: &'a str, + remote_pid: &'a str, + activation_width: i32, + activation_wire_dtype: &'a str, + native_mtp_enabled: bool, + batched_verify_enabled: bool, +} + +fn remote_stage_command(args: RemoteStageCommand<'_>) -> String { + let mut command = String::new(); + if let Some(workdir) = args.workdir { + command.push_str("cd "); + command.push_str(&shell_quote(workdir)); + command.push_str(" && "); + } + command.push_str("SKIPPY_TELEMETRY_STDERR=1 "); + command.push_str("SKIPPY_NATIVE_MTP_ENABLED="); + command.push_str(if args.native_mtp_enabled { "1 " } else { "0 " }); + if !args.batched_verify_enabled { + command.push_str("SKIPPY_NATIVE_MTP_BATCHED_VERIFY=0 "); + } + command.push_str("nohup "); + command.push_str(&shell_quote(args.remote_bin)); + command.push_str(" serve-binary --config "); + command.push_str(&shell_quote(args.remote_config)); + command.push_str(" --topology "); + command.push_str(&shell_quote(args.remote_topology)); + command.push_str(" --activation-width "); + command.push_str(&args.activation_width.to_string()); + command.push_str(" --activation-wire-dtype "); + command.push_str(&shell_quote(args.activation_wire_dtype)); + command.push_str(" --telemetry-level debug > "); + command.push_str(&shell_quote(args.remote_log)); + command.push_str(" 2>&1 < /dev/null & pid=$!; echo $pid > "); + command.push_str(&shell_quote(args.remote_pid)); + command.push_str("; echo $pid"); + command +} + +impl RemoteStageGuard { + pub(super) fn wait_ready(&self, timeout_secs: u64) -> Result<()> { + let Some(pid) = self.pid else { + bail!("remote stage 1 has no pid"); + }; + let attempts = timeout_secs.saturating_mul(2).max(1); + let log = shell_quote(&self.remote_log); + let mut last_stderr = String::new(); + for _ in 0..attempts { + let command = format!( + "if ! kill -0 {pid} 2>/dev/null; then echo dead; exit 2; fi; \ + grep -q 'skippy-server listening:' {log}" + ); + let output = Command::new("ssh") + .arg(&self.host) + .arg(&command) + .output() + .with_context(|| format!("check remote stage 1 readiness on {}", self.host))?; + if output.status.success() { + return Ok(()); + } + if output.status.code() == Some(2) { + bail!("remote stage 1 on {} exited before readiness", self.host); + } + last_stderr = String::from_utf8_lossy(&output.stderr).trim().to_string(); + thread::sleep(Duration::from_millis(500)); + } + bail!( + "timed out waiting for remote stage 1 on {} to log readiness{}", + self.host, + if last_stderr.is_empty() { + String::new() + } else { + format!(": {last_stderr}") + } + ) + } + + pub(super) fn stop_and_collect(&mut self, local_log: &Path) -> Result<()> { + self.terminate(); + match scp_from(&self.host, &self.remote_log, local_log) { + Ok(()) => Ok(()), + Err(error) => { + let note = format!("failed to collect remote stage 1 log: {error:#}\n"); + fs::write(local_log, note) + .with_context(|| format!("failed to write {}", local_log.display())) + } + } + } + + fn terminate(&mut self) { + let Some(pid) = self.pid.take() else { + return; + }; + let command = format!("kill {pid} 2>/dev/null || true; sleep 0.2"); + let _ = Command::new("ssh").arg(&self.host).arg(command).status(); + } +} + +impl Drop for RemoteStageGuard { + fn drop(&mut self) { + self.terminate(); + } +} + +fn ssh_success(host: &str, remote_command: &str) -> Result<()> { + let status = Command::new("ssh") + .arg(host) + .arg(remote_command) + .status() + .with_context(|| format!("run ssh command on {host}"))?; + if !status.success() { + bail!("ssh command on {host} failed with status {status}"); + } + Ok(()) +} + +fn scp_to(host: &str, local_path: &Path, remote_path: &str) -> Result<()> { + let status = Command::new("scp") + .arg(local_path) + .arg(format!("{host}:{remote_path}")) + .status() + .with_context(|| format!("copy {} to {host}:{remote_path}", local_path.display()))?; + if !status.success() { + bail!("scp to {host}:{remote_path} failed with status {status}"); + } + Ok(()) +} + +fn scp_from(host: &str, remote_path: &str, local_path: &Path) -> Result<()> { + let status = Command::new("scp") + .arg(format!("{host}:{remote_path}")) + .arg(local_path) + .status() + .with_context(|| format!("copy {host}:{remote_path} to {}", local_path.display()))?; + if !status.success() { + bail!("scp from {host}:{remote_path} failed with status {status}"); + } + Ok(()) +} + +fn shell_quote(value: &str) -> String { + format!("'{}'", value.replace('\'', "'\"'\"'")) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn remote_stage_command_quotes_paths_and_sets_mtp_env() { + let command = remote_stage_command(RemoteStageCommand { + workdir: Some("/tmp/work dir"), + remote_bin: "target/debug/skippy-server", + remote_config: "/tmp/run/stage'1.json", + remote_topology: "/tmp/run/topology.json", + remote_log: "/tmp/run/stage1.log", + remote_pid: "/tmp/run/stage1.pid", + activation_width: 2048, + activation_wire_dtype: "f16", + native_mtp_enabled: true, + batched_verify_enabled: false, + }); + + assert!(command.contains("cd '/tmp/work dir' && ")); + assert!(command.contains("SKIPPY_NATIVE_MTP_ENABLED=1")); + assert!(command.contains("SKIPPY_NATIVE_MTP_BATCHED_VERIFY=0")); + assert!(command.contains("'target/debug/skippy-server' serve-binary")); + assert!(command.contains("'/tmp/run/stage'\"'\"'1.json'")); + + let batched_command = remote_stage_command(RemoteStageCommand { + batched_verify_enabled: true, + ..RemoteStageCommand { + workdir: None, + remote_bin: "/bin/skippy-server", + remote_config: "/tmp/stage1.json", + remote_topology: "/tmp/topology.json", + remote_log: "/tmp/stage1.log", + remote_pid: "/tmp/stage1.pid", + activation_width: 2048, + activation_wire_dtype: "f16", + native_mtp_enabled: true, + batched_verify_enabled: false, + } + }); + assert!(!batched_command.contains("SKIPPY_NATIVE_MTP_BATCHED_VERIFY")); + } +} diff --git a/crates/skippy-correctness/src/native_mtp_openai/reporting.rs b/crates/skippy-correctness/src/native_mtp_openai/reporting.rs new file mode 100644 index 0000000000..c27347b888 --- /dev/null +++ b/crates/skippy-correctness/src/native_mtp_openai/reporting.rs @@ -0,0 +1,31 @@ +use std::path::Path; + +use anyhow::{Context, Result}; + +pub(super) fn emit_report( + report: &T, + report_out: Option<&Path>, +) -> Result<()> { + let json = serde_json::to_vec_pretty(report)?; + if let Some(path) = report_out { + std::fs::write(path, &json) + .with_context(|| format!("failed to write {}", path.display()))?; + } + println!("{}", String::from_utf8(json)?); + Ok(()) +} + +pub(super) fn status(matches: bool) -> &'static str { + if matches { "pass" } else { "fail" } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn status_names_match_report_contract() { + assert_eq!(status(true), "pass"); + assert_eq!(status(false), "fail"); + } +} diff --git a/crates/skippy-correctness/src/native_mtp_openai/stage_process.rs b/crates/skippy-correctness/src/native_mtp_openai/stage_process.rs new file mode 100644 index 0000000000..3da6486234 --- /dev/null +++ b/crates/skippy-correctness/src/native_mtp_openai/stage_process.rs @@ -0,0 +1,236 @@ +use std::{ + fs::{self, File}, + net::SocketAddr, + path::Path, + process::{Command, Stdio}, + thread, + time::Duration, +}; + +use anyhow::{Context, Result, bail}; +use reqwest::blocking::Client; +use serde_json::Value; + +use crate::{ + cli::NativeMtpOpenAiAbArgs, + support::{ChildGuard, connect_ready}, +}; + +use super::{ + BATCHED_VERIFY_ENV, NATIVE_MTP_ENABLED_ENV, + remote::{RemoteStageGuard, spawn_remote_stage1}, +}; + +pub(super) struct Stage1LaunchReport { + pub(super) launch_mode: &'static str, + pub(super) remote_config: Option, + pub(super) remote_topology: Option, + pub(super) remote_log: Option, +} + +pub(super) enum StageHandle { + Local(ChildGuard), + Remote(RemoteStageGuard), + External, +} + +impl StageHandle { + pub(super) fn stop_and_collect(self, stage1_log: &Path) -> Result<()> { + match self { + StageHandle::Local(guard) => { + drop(guard); + Ok(()) + } + StageHandle::External => Ok(()), + StageHandle::Remote(mut guard) => guard.stop_and_collect(stage1_log), + } + } +} + +pub(super) fn start_stage1( + args: &NativeMtpOpenAiAbArgs, + run_id: &str, + config_path: &Path, + topology_path: &Path, + log_path: &Path, + native_mtp_enabled: bool, + batched_verify_enabled: bool, +) -> Result<(StageHandle, Stage1LaunchReport)> { + if args.external_stage1 { + fs::write(log_path, "stage 1 was externally managed by the caller\n") + .with_context(|| format!("failed to create {}", log_path.display()))?; + return Ok(( + StageHandle::External, + Stage1LaunchReport { + launch_mode: "external", + remote_config: None, + remote_topology: None, + remote_log: None, + }, + )); + } + + if let Some(host) = args.stage1_ssh_host.as_deref() { + let (remote, report) = spawn_remote_stage1( + args, + host, + run_id, + config_path, + topology_path, + native_mtp_enabled, + batched_verify_enabled, + )?; + return Ok((StageHandle::Remote(remote), report)); + } + + let local = spawn_stage( + args, + config_path, + None, + topology_path, + log_path, + native_mtp_enabled, + batched_verify_enabled, + )?; + Ok(( + StageHandle::Local(local), + Stage1LaunchReport { + launch_mode: "local", + remote_config: None, + remote_topology: None, + remote_log: None, + }, + )) +} + +pub(super) fn spawn_stage( + args: &NativeMtpOpenAiAbArgs, + config_path: &Path, + openai_bind_addr: Option, + topology_path: &Path, + log_path: &Path, + native_mtp_enabled: bool, + batched_verify_enabled: bool, +) -> Result { + let log = File::create(log_path) + .with_context(|| format!("failed to create {}", log_path.display()))?; + let mut command = Command::new(&args.server.stage_server_bin); + command.args([ + "serve-binary", + "--config", + config_path + .to_str() + .context("stage config path is not valid UTF-8")?, + "--topology", + topology_path + .to_str() + .context("topology path is not valid UTF-8")?, + "--activation-width", + &args.activation_width.to_string(), + "--activation-wire-dtype", + &args.activation_wire_dtype, + "--telemetry-level", + "debug", + ]); + if let Some(openai_bind_addr) = openai_bind_addr { + command.args(["--openai-bind-addr", &openai_bind_addr.to_string()]); + } + command.env("SKIPPY_TELEMETRY_STDERR", "1"); + if native_mtp_enabled { + command.env(NATIVE_MTP_ENABLED_ENV, "1"); + } else { + command.env(NATIVE_MTP_ENABLED_ENV, "0"); + } + if batched_verify_enabled { + command.env_remove(BATCHED_VERIFY_ENV); + } else { + command.env(BATCHED_VERIFY_ENV, "0"); + } + command.stdout(Stdio::from(log.try_clone()?)); + command.stderr(Stdio::from(log)); + ChildGuard::spawn(command) +} + +pub(super) fn wait_stage1_ready( + stage1: &StageHandle, + addr: SocketAddr, + timeout_secs: u64, +) -> Result<()> { + match stage1 { + StageHandle::Local(_) | StageHandle::External => { + drop(connect_ready(addr, timeout_secs)?); + Ok(()) + } + StageHandle::Remote(guard) => guard.wait_ready(timeout_secs), + } +} + +pub(super) fn wait_openai_ready( + client: &Client, + addr: SocketAddr, + timeout_secs: u64, +) -> Result<()> { + let attempts = timeout_secs.saturating_mul(4).max(1); + let url = format!("http://{addr}/v1/models"); + let mut last_error = None; + for _ in 0..attempts { + match client.get(&url).send() { + Ok(response) if response.status().is_success() => return Ok(()), + Ok(response) => last_error = Some(format!("HTTP {}", response.status())), + Err(error) => last_error = Some(error.to_string()), + } + thread::sleep(Duration::from_millis(250)); + } + bail!( + "timed out waiting for {url}: {}", + last_error.unwrap_or_else(|| "no attempts made".to_string()) + ); +} + +pub(super) fn case_addr(addr: SocketAddr, port_offset: u16, case_index: u16) -> Result { + let offset = port_offset + .checked_mul(case_index) + .context("case port offset exceeds u16")?; + offset_port(addr, offset) +} + +fn offset_port(mut addr: SocketAddr, offset: u16) -> Result { + let port = addr + .port() + .checked_add(offset) + .context("batched port offset exceeds u16")?; + addr.set_port(port); + Ok(addr) +} + +pub(super) fn write_stage_config(path: &Path, value: &Value) -> Result<()> { + fs::write(path, serde_json::to_vec_pretty(value)?) + .with_context(|| format!("failed to write {}", path.display())) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn case_addr_applies_per_case_port_offset() { + let addr: SocketAddr = "127.0.0.1:9000".parse().unwrap(); + + let shifted = case_addr(addr, 10, 2).unwrap(); + + assert_eq!(shifted.to_string(), "127.0.0.1:9020"); + } + + #[test] + fn case_addr_rejects_port_overflow() { + let addr: SocketAddr = "127.0.0.1:65530".parse().unwrap(); + + let error = case_addr(addr, 10, 1).unwrap_err(); + + assert!( + error + .to_string() + .contains("batched port offset exceeds u16") + ); + } +} diff --git a/crates/skippy-correctness/src/report.rs b/crates/skippy-correctness/src/report.rs index 9887eabf49..e65fc921e1 100644 --- a/crates/skippy-correctness/src/report.rs +++ b/crates/skippy-correctness/src/report.rs @@ -6,6 +6,8 @@ pub use model_artifact::ModelIdentity; pub struct BaselineReport { pub token_id: i32, pub predicted_token: i32, + #[serde(skip_serializing_if = "Option::is_none")] + pub second_predicted_token: Option, } #[derive(Debug, Serialize)] @@ -22,17 +24,57 @@ pub struct BoundaryReport { pub struct SplitReport { pub token_id: i32, pub predicted_token: i32, + #[serde(skip_serializing_if = "Option::is_none")] + pub second_predicted_token: Option, + pub native_mtp: NativeMtpSidebandReport, + #[serde(skip_serializing_if = "Option::is_none")] + pub native_mtp_n1: Option, pub activation_width: i32, pub wire_dtype: String, pub boundary: BoundaryReport, } +#[derive(Debug, Clone, Serialize)] +pub struct NativeMtpSidebandReport { + pub sideband_present: bool, + pub predicted_token_count: usize, + pub authoritative_matches_reply: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub authoritative_token: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub draft_token: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub proposal_compute_us: Option, +} + +#[derive(Debug, Clone, Serialize)] +pub struct NativeMtpN1VerificationReport { + pub drafted_tokens: u64, + pub accepted_tokens: u64, + pub rejected_tokens: u64, + pub pending_tokens: u64, + pub verification_count: u64, + pub accept_rate: f64, + pub byte_identical: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub draft_token: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub second_target_token: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub second_baseline_token: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub proposal_compute_us: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub verification_compute_us: Option, +} + #[derive(Debug, Serialize)] pub struct SingleStepReport { pub mode: &'static str, pub status: &'static str, pub model_identity: ModelIdentity, pub matches: bool, + pub native_mtp_draft_required: bool, pub baseline: BaselineReport, pub split: SplitReport, pub stage_models: Vec, @@ -55,9 +97,15 @@ pub struct ChainReport { pub status: &'static str, pub model_identity: ModelIdentity, pub matches: bool, + pub native_mtp_draft_required: bool, pub baseline: BaselineReport, pub token_id: i32, pub predicted_token: i32, + #[serde(skip_serializing_if = "Option::is_none")] + pub second_predicted_token: Option, + pub native_mtp: NativeMtpSidebandReport, + #[serde(skip_serializing_if = "Option::is_none")] + pub native_mtp_n1: Option, pub activation_width: i32, pub wire_dtype: String, pub stages: Vec, @@ -199,3 +247,78 @@ pub struct StatePayloadBlockDigestReport { pub bytes: usize, pub sha256: String, } + +#[derive(Debug, Serialize)] +pub struct NativeMtpOpenAiAbReport { + pub mode: &'static str, + pub status: &'static str, + pub model_id: String, + pub model_path: String, + pub prompt: String, + pub max_tokens: u32, + pub split_layer: u32, + pub layer_end: u32, + pub activation_width: i32, + pub activation_wire_dtype: String, + pub exact_content_match: bool, + pub batched_events_required: bool, + pub batched_events_present: bool, + pub matches: bool, + pub baseline: NativeMtpOpenAiCaseReport, + pub n1: NativeMtpOpenAiCaseReport, + pub batched: NativeMtpOpenAiCaseReport, +} + +#[derive(Debug, Serialize)] +pub struct NativeMtpOpenAiCaseReport { + pub case: &'static str, + pub native_mtp_enabled: bool, + pub batched_verify_enabled: bool, + pub http_status: u16, + pub content: String, + pub completion_tokens: Option, + pub openai_bind_addr: String, + pub stage0_bind_addr: String, + pub stage0_endpoint_addr: String, + pub stage1_bind_addr: String, + pub stage1_endpoint_addr: String, + pub stage0_config: String, + pub stage1_config: String, + pub topology_config: String, + pub stage0_log: String, + pub stage1_log: String, + pub stage1_launch_mode: String, + pub stage1_remote_config: Option, + pub stage1_remote_topology: Option, + pub stage1_remote_log: Option, + pub metrics: NativeMtpOpenAiMetricsReport, +} + +#[derive(Debug, Default, Serialize)] +pub struct NativeMtpOpenAiMetricsReport { + pub native_mtp_enabled: bool, + pub drafted_tokens: u64, + pub accepted_tokens: u64, + pub rejected_tokens: u64, + pub pending_tokens: u64, + pub verification_count: u64, + pub accept_rate: f64, + pub proposal_compute_us: i64, + pub verification_compute_us: i64, + pub decode_token_events: u64, + pub batched_verify_events: u64, + pub batched_accepted_events: u64, + pub batched_rejected_events: u64, + pub batched_accepted_verify_elapsed_ms: f64, + pub batched_accepted_verify_avg_ms: f64, + pub batched_rejected_verify_elapsed_ms: f64, + pub batched_rejected_verify_avg_ms: f64, + pub batched_consumed_positions: u64, + pub batched_committed_positions: u64, + pub batched_trim_count: u64, + pub batched_trim_elapsed_ms: f64, + pub batched_trim_local_ms: f64, + pub batched_trim_downstream_write_ms: f64, + pub batched_trim_downstream_wait_ms: f64, + pub fatal_error_events: u64, +} diff --git a/crates/skippy-correctness/src/runner.rs b/crates/skippy-correctness/src/runner.rs index a4a08952f2..14a297e1f7 100644 --- a/crates/skippy-correctness/src/runner.rs +++ b/crates/skippy-correctness/src/runner.rs @@ -8,7 +8,7 @@ use std::{ }; use anyhow::{Context, Result, bail}; -use model_artifact::ModelIdentity; +use model_artifact::{ModelIdentity, gguf::scan_gguf_compact_meta}; use model_hf::HfModelRepository; use model_ref::ModelRef; use serde::Deserialize; @@ -16,7 +16,7 @@ use serde::Serialize; use serde_json::json; use sha2::{Digest, Sha256}; use skippy_protocol::binary::{ - StageStateHeader, StageWireMessage, WireMessageKind, WireReplyKind, + StageReply, StageStateHeader, StageWireMessage, WireMessageKind, WireReplyKind, activation_state_flags_from_frame_flags, read_stage_message, recv_reply, state_flags, write_stage_message, }; @@ -28,15 +28,14 @@ use skippy_runtime::{ use crate::{ cli::{ - ChainArgs, DtypeMatrixArgs, FlashAttentionArg, RuntimeArgs, ServerArgs, SingleStepArgs, - SplitScanArgs, StageLoadMode, StateHandoffArgs, StatePayloadKind, + ChainArgs, DtypeMatrixArgs, FlashAttentionArg, NativeMtpArgs, RuntimeArgs, ServerArgs, + SingleStepArgs, SplitScanArgs, StageLoadMode, StateHandoffArgs, StatePayloadKind, }, - direct_return::CorrectnessDirectReturnServer, report::{ BaselineReport, BoundaryReport, ChainReport, ChainStageReport, DtypeMatrixReport, - PackagePartReport, PackageStageReport, SingleStepReport, SplitReport, SplitScanReport, - StageModelReport, StateHandoffReport, StatePayloadBlockDigestReport, - StatePayloadDigestReport, + NativeMtpN1VerificationReport, NativeMtpSidebandReport, PackagePartReport, + PackageStageReport, SingleStepReport, SplitReport, SplitScanReport, StageModelReport, + StateHandoffReport, StatePayloadBlockDigestReport, StatePayloadDigestReport, }, support::{ ChildGuard, activation_width, connect_ready, generate_run_id, parse_wire_dtype, @@ -47,6 +46,7 @@ use crate::{ struct FullModelResult { token_id: i32, predicted_token: i32, + second_predicted_token: Option, } struct BinarySplitConfig { @@ -68,11 +68,20 @@ struct BinarySplitConfig { startup_timeout_secs: u64, max_inflight: usize, model_identity: ModelIdentity, + native_mtp_verification: bool, +} + +#[derive(Clone, Copy)] +struct NativeMtpRequirement { + require_draft: bool, } struct BinarySplitResult { token_id: i32, predicted_token: i32, + second_predicted_token: Option, + native_mtp: NativeMtpSidebandReport, + native_mtp_verification_compute_us: Option, activation_width: i32, wire_dtype: String, boundary_producer_stage_index: i32, @@ -84,6 +93,37 @@ struct BinarySplitResult { stage_models: Vec, } +#[derive(Clone, Debug, Default, PartialEq, Eq)] +struct NativeMtpArtifactSummary { + nextn_predict_layers: u32, + has_eh_proj: bool, + has_enorm: bool, + has_hnorm: bool, +} + +impl NativeMtpArtifactSummary { + fn supports_native_mtp(&self) -> bool { + self.nextn_predict_layers > 0 && self.has_eh_proj && self.has_enorm && self.has_hnorm + } + + fn missing_reasons(&self) -> Vec<&'static str> { + let mut missing = Vec::new(); + if self.nextn_predict_layers == 0 { + missing.push("*.nextn_predict_layers > 0"); + } + if !self.has_eh_proj { + missing.push("*.nextn.eh_proj tensor"); + } + if !self.has_enorm { + missing.push("*.nextn.enorm tensor"); + } + if !self.has_hnorm { + missing.push("*.nextn.hnorm tensor"); + } + missing + } +} + struct BinaryChainConfig { stage_server_bin: PathBuf, model: PathBuf, @@ -105,11 +145,15 @@ struct BinaryChainConfig { startup_timeout_secs: u64, max_inflight: usize, model_identity: ModelIdentity, + native_mtp_verification: bool, } struct BinaryChainResult { token_id: i32, predicted_token: i32, + second_predicted_token: Option, + native_mtp: NativeMtpSidebandReport, + native_mtp_verification_compute_us: Option, activation_width: i32, wire_dtype: String, stage0_wire_payload_bytes: usize, @@ -291,6 +335,8 @@ impl LocalStatePayload { } pub fn single_step(args: SingleStepArgs) -> Result<()> { + let native_mtp = native_mtp_requirement(args.native_mtp); + ensure_native_mtp_artifact_if_required(&args.runtime, native_mtp)?; let model_identity = runtime_model_identity(&args.runtime)?; let baseline = run_full_model_decode(&args.runtime)?; let report = run_single_step_with_baseline( @@ -302,6 +348,7 @@ pub fn single_step(args: SingleStepArgs) -> Result<()> { split_layer: args.split_layer, stage1_bind_addr: args.stage1_bind_addr, activation_wire_dtype: args.activation_wire_dtype, + native_mtp, }, )?; emit_report(&report, args.output.report_out.as_deref())?; @@ -310,6 +357,8 @@ pub fn single_step(args: SingleStepArgs) -> Result<()> { } pub fn chain(args: ChainArgs) -> Result<()> { + let native_mtp_requirement = native_mtp_requirement(args.native_mtp); + ensure_native_mtp_artifact_if_required(&args.runtime, native_mtp_requirement)?; let splits = parse_chain_splits(&args.splits)?; let model_identity = runtime_model_identity(&args.runtime)?; let baseline = run_full_model_decode(&args.runtime)?; @@ -334,16 +383,31 @@ pub fn chain(args: ChainArgs) -> Result<()> { startup_timeout_secs: args.server.startup_timeout_secs, max_inflight: args.server.max_inflight, model_identity: model_identity.clone(), + native_mtp_verification: native_mtp_requirement.require_draft, })?; - let matches = baseline.predicted_token == chain.predicted_token; + let native_mtp = chain.native_mtp.clone(); + let native_mtp_n1 = native_mtp_n1_verification_report( + native_mtp_requirement.require_draft, + &native_mtp, + chain.second_predicted_token, + baseline.second_predicted_token, + chain.native_mtp_verification_compute_us, + ); + let matches = baseline.predicted_token == chain.predicted_token + && native_mtp_satisfies_requirement(&native_mtp, native_mtp_requirement) + && native_mtp_n1_satisfies_requirement(&native_mtp_n1, native_mtp_requirement); let report = ChainReport { mode: "chain", status: status(matches), model_identity, matches, + native_mtp_draft_required: native_mtp_requirement.require_draft, baseline: baseline_report(baseline), token_id: chain.token_id, predicted_token: chain.predicted_token, + second_predicted_token: chain.second_predicted_token, + native_mtp, + native_mtp_n1, activation_width: chain.activation_width, wire_dtype: chain.wire_dtype, stages: vec![ @@ -383,6 +447,8 @@ pub fn chain(args: ChainArgs) -> Result<()> { } pub fn split_scan(args: SplitScanArgs) -> Result<()> { + let native_mtp = native_mtp_requirement(args.native_mtp); + ensure_native_mtp_artifact_if_required(&args.runtime, native_mtp)?; let splits = parse_split_list(&args.splits)?; let model_identity = runtime_model_identity(&args.runtime)?; let baseline = run_full_model_decode(&args.runtime)?; @@ -401,11 +467,13 @@ pub fn split_scan(args: SplitScanArgs) -> Result<()> { FullModelResult { token_id: baseline.token_id, predicted_token: baseline.predicted_token, + second_predicted_token: baseline.second_predicted_token, }, SingleStepCase { split_layer, stage1_bind_addr: args.stage1_bind_addr, activation_wire_dtype: args.activation_wire_dtype.clone(), + native_mtp, }, )?); } @@ -425,6 +493,8 @@ pub fn split_scan(args: SplitScanArgs) -> Result<()> { } pub fn dtype_matrix(args: DtypeMatrixArgs) -> Result<()> { + let native_mtp = native_mtp_requirement(args.native_mtp); + ensure_native_mtp_artifact_if_required(&args.runtime, native_mtp)?; let dtypes = parse_csv(&args.dtypes)?; let model_identity = runtime_model_identity(&args.runtime)?; let baseline = run_full_model_decode(&args.runtime)?; @@ -437,11 +507,13 @@ pub fn dtype_matrix(args: DtypeMatrixArgs) -> Result<()> { FullModelResult { token_id: baseline.token_id, predicted_token: baseline.predicted_token, + second_predicted_token: baseline.second_predicted_token, }, SingleStepCase { split_layer: args.split_layer, stage1_bind_addr: args.stage1_bind_addr, activation_wire_dtype: dtype, + native_mtp, }, )?); } @@ -572,6 +644,7 @@ struct SingleStepCase { split_layer: u32, stage1_bind_addr: SocketAddr, activation_wire_dtype: String, + native_mtp: NativeMtpRequirement, } fn run_single_step_with_baseline( @@ -600,16 +673,27 @@ fn run_single_step_with_baseline( startup_timeout_secs: server.startup_timeout_secs, max_inflight: server.max_inflight, model_identity: model_identity.clone(), + native_mtp_verification: case.native_mtp.require_draft, })?; - let matches = baseline.predicted_token == split.predicted_token; + let native_mtp_n1 = native_mtp_n1_verification_report( + case.native_mtp.require_draft, + &split.native_mtp, + split.second_predicted_token, + baseline.second_predicted_token, + split.native_mtp_verification_compute_us, + ); + let matches = baseline.predicted_token == split.predicted_token + && native_mtp_satisfies_requirement(&split.native_mtp, case.native_mtp) + && native_mtp_n1_satisfies_requirement(&native_mtp_n1, case.native_mtp); let stage_models = split.stage_models.clone(); Ok(SingleStepReport { mode: "single-step", status: status(matches), model_identity: model_identity.clone(), matches, + native_mtp_draft_required: case.native_mtp.require_draft, baseline: baseline_report(baseline), - split: split_report(split), + split: split_report(split, native_mtp_n1), stage_models, }) } @@ -648,12 +732,66 @@ fn run_full_model_decode(args: &RuntimeArgs) -> Result { .decode_step_frame(token_id, None, 0) .context("full model failed to decode")? .0; + let second_predicted_token = session + .decode_step_frame(predicted_token, None, 0) + .context("full model failed to decode second token")? + .0; Ok(FullModelResult { token_id, predicted_token, + second_predicted_token: Some(second_predicted_token), }) } +struct BinaryDecodeMessageArgs<'a> { + wire_dtype: skippy_protocol::binary::WireActivationDType, + token_id: i32, + decode_step: i32, + source_stage_index: i32, + boundary: &'a ActivationFrame, + activation_width: i32, + request_id: u64, + session_id: u64, +} + +fn binary_decode_message(args: BinaryDecodeMessageArgs<'_>) -> Result { + let mut state = StageStateHeader::new(WireMessageKind::DecodeEmbd, args.wire_dtype); + state.prompt_token_count = 0; + state.decode_step = args.decode_step; + state.current_token = args.token_id; + state.source_stage_index = args.source_stage_index; + state.flags |= activation_state_flags(args.boundary); + let activation = skippy_protocol::binary::encode_f32_activation_payload_with_state_flags( + args.wire_dtype, + 1, + args.activation_width, + &args.boundary.payload, + activation_state_flags(args.boundary), + ) + .context("failed to encode boundary activation for wire")?; + Ok(StageWireMessage { + kind: WireMessageKind::DecodeEmbd, + pos_start: args.decode_step, + token_count: 1, + state, + request_id: args.request_id, + session_id: args.session_id, + sampling: None, + chat_sampling_metadata: None, + tokens: vec![args.token_id], + positions: vec![args.decode_step], + activation, + raw_bytes: Vec::new(), + }) +} + +fn ensure_reply_kind(reply: &StageReply, expected: WireReplyKind) -> Result<()> { + if reply.kind != expected { + bail!("expected {expected:?} reply, got {:?}", reply.kind); + } + Ok(()) +} + fn run_binary_split(args: BinarySplitConfig) -> Result { if args.split_layer == 0 || args.split_layer >= args.layer_end { bail!("split_layer must be greater than zero and less than layer_end"); @@ -729,7 +867,6 @@ fn run_binary_split(args: BinarySplitConfig) -> Result { } let activation_width = activation_width(&boundary)?; - let direct_returns = CorrectnessDirectReturnServer::start("127.0.0.1:0")?; let run_id = generate_run_id(); let model_id = args.model_identity.model_id.clone(); let config_path = temp_config_path_for(&run_id, "stage-1"); @@ -759,7 +896,7 @@ fn run_binary_split(args: BinarySplitConfig) -> Result { "upstream": { "stage_id": "stage-0", "stage_index": 0, - "endpoint": format!("tcp://{}", direct_returns.endpoint()) + "endpoint": "driver" }, "downstream": null }); @@ -770,7 +907,7 @@ fn run_binary_split(args: BinarySplitConfig) -> Result { CorrectnessTopologyStage { stage_id: "stage-0", stage_index: 0, - endpoint: format!("tcp://{}", direct_returns.endpoint()), + endpoint: "driver".to_string(), layer_start: 0, layer_end: args.split_layer, load_mode: protocol_load_mode(args.stage_load_mode), @@ -815,47 +952,59 @@ fn run_binary_split(args: BinarySplitConfig) -> Result { .context("stage 1 binary server did not become ready")?; let request_id = 1; let session_id = 1; - let direct_return = direct_returns.register(request_id, session_id)?; send_generation_config(&mut stream, wire_dtype, request_id, session_id, 1) .context("send binary generation config")?; - let mut state = StageStateHeader::new(WireMessageKind::DecodeEmbd, wire_dtype); - state.prompt_token_count = 0; - state.decode_step = 0; - state.current_token = token_id; - state.source_stage_index = 0; - state.flags |= activation_state_flags(&boundary); - let activation = skippy_protocol::binary::encode_f32_activation_payload_with_state_flags( + let message = binary_decode_message(BinaryDecodeMessageArgs { wire_dtype, - 1, + token_id, + decode_step: 0, + source_stage_index: 0, + boundary: &boundary, activation_width, - &boundary.payload, - activation_state_flags(&boundary), - ) - .context("failed to encode boundary activation for wire")?; - let message = StageWireMessage { - kind: WireMessageKind::DecodeEmbd, - pos_start: 0, - token_count: 1, - state, request_id, session_id, - sampling: None, - chat_sampling_metadata: None, - tokens: vec![token_id], - positions: vec![0], - activation, - raw_bytes: Vec::new(), - }; + })?; write_stage_message(&mut stream, &message, wire_dtype).context("send binary decode")?; - let reply = direct_return - .recv_expected(WireReplyKind::PredictedToken) - .context("receive direct binary reply")?; + let reply = recv_reply(&mut stream).context("receive binary prediction reply")?; + ensure_reply_kind(&reply, WireReplyKind::PredictedToken)?; + let native_mtp = native_mtp_sideband_report(&reply); + let (second_predicted_token, native_mtp_verification_compute_us) = + if args.native_mtp_verification { + let verification_timer = Instant::now(); + let (_boundary_prediction, second_boundary) = session0 + .decode_step_frame(reply.predicted, None, 0) + .context("stage 0 failed to produce second activation frame")?; + let second_message = binary_decode_message(BinaryDecodeMessageArgs { + wire_dtype, + token_id: reply.predicted, + decode_step: 1, + source_stage_index: 0, + boundary: &second_boundary, + activation_width, + request_id, + session_id, + })?; + write_stage_message(&mut stream, &second_message, wire_dtype) + .context("send second binary decode")?; + let second_reply = + recv_reply(&mut stream).context("receive second binary prediction reply")?; + ensure_reply_kind(&second_reply, WireReplyKind::PredictedToken)?; + ( + Some(second_reply.predicted), + Some(elapsed_us(verification_timer)), + ) + } else { + (None, None) + }; write_stage_message(&mut stream, &StageWireMessage::stop(wire_dtype), wire_dtype) .context("send binary stop")?; Ok(BinarySplitResult { token_id, predicted_token: reply.predicted, + second_predicted_token, + native_mtp, + native_mtp_verification_compute_us, activation_width, wire_dtype: args.activation_wire_dtype, boundary_producer_stage_index: boundary.desc.producer_stage_index, @@ -962,7 +1111,6 @@ fn run_binary_chain(args: BinaryChainConfig) -> Result { } let activation_width = activation_width(&boundary)?; - let direct_returns = CorrectnessDirectReturnServer::start("127.0.0.1:0")?; let run_id = generate_run_id(); let model_id = args.model_identity.model_id.clone(); let stage1_config_path = temp_config_path_for(&run_id, "stage-1"); @@ -1022,7 +1170,7 @@ fn run_binary_chain(args: BinaryChainConfig) -> Result { "upstream": { "stage_id": "stage-0", "stage_index": 0, - "endpoint": format!("tcp://{}", direct_returns.endpoint()) + "endpoint": "driver" }, "downstream": { "stage_id": "stage-2", @@ -1037,7 +1185,7 @@ fn run_binary_chain(args: BinaryChainConfig) -> Result { CorrectnessTopologyStage { stage_id: "stage-0", stage_index: 0, - endpoint: format!("tcp://{}", direct_returns.endpoint()), + endpoint: "driver".to_string(), layer_start: 0, layer_end: args.split_layer_1, load_mode: protocol_load_mode(args.stage_load_mode), @@ -1123,47 +1271,59 @@ fn run_binary_chain(args: BinaryChainConfig) -> Result { .context("stage 1 binary server did not become ready")?; let request_id = 2; let session_id = 2; - let direct_return = direct_returns.register(request_id, session_id)?; send_generation_config(&mut stream, wire_dtype, request_id, session_id, 1) .context("send binary chain generation config")?; - let mut state = StageStateHeader::new(WireMessageKind::DecodeEmbd, wire_dtype); - state.prompt_token_count = 0; - state.decode_step = 0; - state.current_token = token_id; - state.source_stage_index = 0; - state.flags |= activation_state_flags(&boundary); - let activation = skippy_protocol::binary::encode_f32_activation_payload_with_state_flags( + let message = binary_decode_message(BinaryDecodeMessageArgs { wire_dtype, - 1, + token_id, + decode_step: 0, + source_stage_index: 0, + boundary: &boundary, activation_width, - &boundary.payload, - activation_state_flags(&boundary), - ) - .context("failed to encode boundary activation for wire")?; - let message = StageWireMessage { - kind: WireMessageKind::DecodeEmbd, - pos_start: 0, - token_count: 1, - state, request_id, session_id, - sampling: None, - chat_sampling_metadata: None, - tokens: vec![token_id], - positions: vec![0], - activation, - raw_bytes: Vec::new(), - }; + })?; write_stage_message(&mut stream, &message, wire_dtype).context("send binary chain decode")?; - let reply = direct_return - .recv_expected(WireReplyKind::PredictedToken) - .context("receive direct binary chain reply")?; + let reply = recv_reply(&mut stream).context("receive binary chain prediction reply")?; + ensure_reply_kind(&reply, WireReplyKind::PredictedToken)?; + let native_mtp = native_mtp_sideband_report(&reply); + let (second_predicted_token, native_mtp_verification_compute_us) = + if args.native_mtp_verification { + let verification_timer = Instant::now(); + let (_boundary_prediction, second_boundary) = session0 + .decode_step_frame(reply.predicted, None, 0) + .context("stage 0 failed to produce second chain activation frame")?; + let second_message = binary_decode_message(BinaryDecodeMessageArgs { + wire_dtype, + token_id: reply.predicted, + decode_step: 1, + source_stage_index: 0, + boundary: &second_boundary, + activation_width, + request_id, + session_id, + })?; + write_stage_message(&mut stream, &second_message, wire_dtype) + .context("send second binary chain decode")?; + let second_reply = + recv_reply(&mut stream).context("receive second binary chain prediction reply")?; + ensure_reply_kind(&second_reply, WireReplyKind::PredictedToken)?; + ( + Some(second_reply.predicted), + Some(elapsed_us(verification_timer)), + ) + } else { + (None, None) + }; write_stage_message(&mut stream, &StageWireMessage::stop(wire_dtype), wire_dtype) .context("send binary chain stop")?; Ok(BinaryChainResult { token_id, predicted_token: reply.predicted, + second_predicted_token, + native_mtp, + native_mtp_verification_compute_us, activation_width, wire_dtype: args.activation_wire_dtype, stage0_wire_payload_bytes: message.activation.len(), @@ -1587,6 +1747,10 @@ fn elapsed_ms(started: Instant) -> f64 { started.elapsed().as_secs_f64() * 1000.0 } +fn elapsed_us(started: Instant) -> i64 { + (started.elapsed().as_secs_f64() * 1_000_000.0).round() as i64 +} + fn mean_pair_sum(left: &[f64], right: &[f64]) -> f64 { let count = left.len().min(right.len()); if count == 0 { @@ -2655,13 +2819,20 @@ fn baseline_report(result: FullModelResult) -> BaselineReport { BaselineReport { token_id: result.token_id, predicted_token: result.predicted_token, + second_predicted_token: result.second_predicted_token, } } -fn split_report(result: BinarySplitResult) -> SplitReport { +fn split_report( + result: BinarySplitResult, + native_mtp_n1: Option, +) -> SplitReport { SplitReport { token_id: result.token_id, predicted_token: result.predicted_token, + second_predicted_token: result.second_predicted_token, + native_mtp: result.native_mtp, + native_mtp_n1, activation_width: result.activation_width, wire_dtype: result.wire_dtype, boundary: BoundaryReport { @@ -2675,15 +2846,173 @@ fn split_report(result: BinarySplitResult) -> SplitReport { } } +fn native_mtp_n1_verification_report( + requested: bool, + first: &NativeMtpSidebandReport, + second_target_token: Option, + second_baseline_token: Option, + verification_compute_us: Option, +) -> Option { + if !requested && first.draft_token.is_none() { + return None; + } + + let drafted_tokens = u64::from(first.draft_token.is_some()); + let verification_count = + u64::from(first.draft_token.is_some() && second_target_token.is_some()); + let accepted_tokens = u64::from( + matches!((first.draft_token, second_target_token), (Some(draft), Some(target)) if draft == target), + ); + let rejected_tokens = verification_count.saturating_sub(accepted_tokens); + let pending_tokens = drafted_tokens.saturating_sub(verification_count); + let byte_identical = matches!((second_target_token, second_baseline_token), (Some(target), Some(baseline)) if target == baseline); + let accept_rate = if verification_count == 0 { + 0.0 + } else { + accepted_tokens as f64 / verification_count as f64 + }; + + Some(NativeMtpN1VerificationReport { + drafted_tokens, + accepted_tokens, + rejected_tokens, + pending_tokens, + verification_count, + accept_rate, + byte_identical, + draft_token: first.draft_token, + second_target_token, + second_baseline_token, + proposal_compute_us: first.proposal_compute_us, + verification_compute_us, + }) +} + +fn native_mtp_n1_satisfies_requirement( + report: &Option, + requirement: NativeMtpRequirement, +) -> bool { + if !requirement.require_draft { + return true; + } + report.as_ref().is_some_and(|report| { + report.drafted_tokens == 1 + && report.verification_count == 1 + && report.pending_tokens == 0 + && report.byte_identical + }) +} + +fn native_mtp_sideband_report(reply: &StageReply) -> NativeMtpSidebandReport { + let authoritative_token = reply.predicted_tokens.first().copied(); + let draft_token = reply.predicted_tokens.get(1).copied(); + let proposal_compute_us = reply + .predicted_tokens + .get(2) + .copied() + .map(|value| i64::from(value.max(0))); + NativeMtpSidebandReport { + sideband_present: draft_token.is_some(), + predicted_token_count: reply.predicted_tokens.len(), + authoritative_matches_reply: authoritative_token + .is_none_or(|token| token == reply.predicted), + authoritative_token, + draft_token, + proposal_compute_us, + } +} + +fn native_mtp_requirement(args: NativeMtpArgs) -> NativeMtpRequirement { + NativeMtpRequirement { + require_draft: args.require_native_mtp_draft, + } +} + +fn ensure_native_mtp_artifact_if_required( + runtime: &RuntimeArgs, + requirement: NativeMtpRequirement, +) -> Result<()> { + if !requirement.require_draft { + return Ok(()); + } + + let model_path = native_mtp_preflight_model_path(runtime); + if !model_path.is_file() { + return Ok(()); + } + + let summary = native_mtp_artifact_summary(model_path)?; + if summary.supports_native_mtp() { + return Ok(()); + } + + bail!( + "native MTP draft was required, but {} does not advertise a usable native MTP head: missing {}", + model_path.display(), + summary.missing_reasons().join(", ") + ); +} + +fn native_mtp_preflight_model_path(runtime: &RuntimeArgs) -> &Path { + runtime + .stage_model + .as_deref() + .filter(|path| path.is_file()) + .unwrap_or(runtime.model.as_path()) +} + +fn native_mtp_artifact_summary(model_path: &Path) -> Result { + let meta = scan_gguf_compact_meta(model_path) + .with_context(|| format!("inspect GGUF metadata for {}", model_path.display()))?; + let info = skippy_runtime::ModelInfo::open(model_path) + .with_context(|| format!("inspect GGUF tensors for {}", model_path.display()))?; + let tensors = info.tensors()?; + Ok(native_mtp_artifact_summary_from_names( + meta.nextn_predict_layers, + tensors.iter().map(|tensor| tensor.name.as_str()), + )) +} + +fn native_mtp_artifact_summary_from_names<'a>( + nextn_predict_layers: u32, + names: impl IntoIterator, +) -> NativeMtpArtifactSummary { + let mut summary = NativeMtpArtifactSummary { + nextn_predict_layers, + ..NativeMtpArtifactSummary::default() + }; + for name in names { + let name = name.to_ascii_lowercase(); + summary.has_eh_proj |= native_mtp_name_matches(&name, "eh_proj"); + summary.has_enorm |= native_mtp_name_matches(&name, "enorm"); + summary.has_hnorm |= native_mtp_name_matches(&name, "hnorm"); + } + summary +} + +fn native_mtp_name_matches(name: &str, suffix: &str) -> bool { + name.contains(&format!(".nextn.{suffix}")) + || name.contains(&format!(".{suffix}.")) + || name.ends_with(&format!(".{suffix}")) +} + +fn native_mtp_satisfies_requirement( + report: &NativeMtpSidebandReport, + requirement: NativeMtpRequirement, +) -> bool { + report.authoritative_matches_reply && (!requirement.require_draft || report.sideband_present) +} + fn emit_report(report: &T, report_out: Option<&Path>) -> Result<()> { let json = serde_json::to_string_pretty(report)?; println!("{json}"); if let Some(path) = report_out { - if let Some(parent) = path.parent() - && !parent.as_os_str().is_empty() - { - fs::create_dir_all(parent) - .with_context(|| format!("create report directory {}", parent.display()))?; + match path.parent() { + Some(parent) if !parent.as_os_str().is_empty() => { + fs::create_dir_all(parent) + .with_context(|| format!("create report directory {}", parent.display()))?; + } + _ => {} } fs::write(path, format!("{json}\n")) .with_context(|| format!("write correctness report {}", path.display()))?; @@ -2691,6 +3020,187 @@ fn emit_report(report: &T, report_out: Option<&Path>) -> Result<() Ok(()) } +#[cfg(test)] +mod tests { + use skippy_protocol::binary::{StageReplyStats, WireReplyKind}; + + use super::*; + + fn predicted_reply(predicted: i32, predicted_tokens: Vec) -> StageReply { + StageReply { + kind: WireReplyKind::PredictedToken, + predicted, + predicted_tokens, + stats: StageReplyStats::default(), + } + } + + #[test] + fn native_mtp_report_treats_plain_authoritative_token_as_no_draft() { + let report = native_mtp_sideband_report(&predicted_reply(11, vec![11])); + + assert!(!report.sideband_present); + assert_eq!(report.predicted_token_count, 1); + assert!(report.authoritative_matches_reply); + assert_eq!(report.authoritative_token, Some(11)); + assert_eq!(report.draft_token, None); + assert_eq!(report.proposal_compute_us, None); + } + + #[test] + fn native_mtp_report_extracts_draft_sideband() { + let report = native_mtp_sideband_report(&predicted_reply(11, vec![11, 12, 34])); + + assert!(report.sideband_present); + assert_eq!(report.predicted_token_count, 3); + assert!(report.authoritative_matches_reply); + assert_eq!(report.authoritative_token, Some(11)); + assert_eq!(report.draft_token, Some(12)); + assert_eq!(report.proposal_compute_us, Some(34)); + } + + #[test] + fn native_mtp_report_flags_authoritative_sideband_mismatch() { + let report = native_mtp_sideband_report(&predicted_reply(11, vec![10, 12, 34])); + + assert!(report.sideband_present); + assert!(!report.authoritative_matches_reply); + assert_eq!(report.authoritative_token, Some(10)); + assert_eq!(report.draft_token, Some(12)); + } + + #[test] + fn native_mtp_report_clamps_negative_proposal_time() { + let report = native_mtp_sideband_report(&predicted_reply(11, vec![11, 12, -34])); + + assert_eq!(report.proposal_compute_us, Some(0)); + } + + #[test] + fn native_mtp_requirement_can_require_draft_presence() { + let no_draft = native_mtp_sideband_report(&predicted_reply(11, vec![11])); + let draft = native_mtp_sideband_report(&predicted_reply(11, vec![11, 12, 34])); + let optional = NativeMtpRequirement { + require_draft: false, + }; + let required = NativeMtpRequirement { + require_draft: true, + }; + + assert!(native_mtp_satisfies_requirement(&no_draft, optional)); + assert!(!native_mtp_satisfies_requirement(&no_draft, required)); + assert!(native_mtp_satisfies_requirement(&draft, required)); + } + + #[test] + fn native_mtp_n1_report_accepts_matching_second_target() { + let first = native_mtp_sideband_report(&predicted_reply(11, vec![11, 12, 34])); + let report = native_mtp_n1_verification_report(true, &first, Some(12), Some(12), Some(9)) + .expect("verification report"); + + assert_eq!(report.drafted_tokens, 1); + assert_eq!(report.accepted_tokens, 1); + assert_eq!(report.rejected_tokens, 0); + assert_eq!(report.pending_tokens, 0); + assert_eq!(report.verification_count, 1); + assert_eq!(report.accept_rate, 1.0); + assert!(report.byte_identical); + assert_eq!(report.proposal_compute_us, Some(34)); + assert_eq!(report.verification_compute_us, Some(9)); + assert!(native_mtp_n1_satisfies_requirement( + &Some(report), + NativeMtpRequirement { + require_draft: true + } + )); + } + + #[test] + fn native_mtp_n1_report_rejects_mismatched_draft_without_failing_byte_identity() { + let first = native_mtp_sideband_report(&predicted_reply(11, vec![11, 12, 34])); + let report = native_mtp_n1_verification_report(true, &first, Some(13), Some(13), Some(9)) + .expect("verification report"); + + assert_eq!(report.drafted_tokens, 1); + assert_eq!(report.accepted_tokens, 0); + assert_eq!(report.rejected_tokens, 1); + assert_eq!(report.pending_tokens, 0); + assert_eq!(report.verification_count, 1); + assert_eq!(report.accept_rate, 0.0); + assert!(report.byte_identical); + assert!(native_mtp_n1_satisfies_requirement( + &Some(report), + NativeMtpRequirement { + require_draft: true + } + )); + } + + #[test] + fn native_mtp_n1_requirement_fails_when_required_draft_is_missing() { + let first = native_mtp_sideband_report(&predicted_reply(11, vec![11])); + let report = native_mtp_n1_verification_report(true, &first, Some(13), Some(13), Some(9)) + .expect("required verification report"); + + assert_eq!(report.drafted_tokens, 0); + assert_eq!(report.verification_count, 0); + assert!(report.byte_identical); + assert!(!native_mtp_n1_satisfies_requirement( + &Some(report), + NativeMtpRequirement { + require_draft: true + } + )); + } + + #[test] + fn native_mtp_artifact_summary_requires_metadata_and_tensors() { + let summary = native_mtp_artifact_summary_from_names( + 1, + [ + "blk.47.nextn.eh_proj", + "blk.47.nextn.enorm", + "blk.47.nextn.hnorm", + ], + ); + + assert!(summary.supports_native_mtp()); + assert!(summary.missing_reasons().is_empty()); + } + + #[test] + fn native_mtp_artifact_summary_accepts_source_style_tensor_names() { + let summary = native_mtp_artifact_summary_from_names( + 1, + [ + "model.layers.47.eh_proj.weight", + "model.layers.47.enorm.weight", + "model.layers.47.hnorm.weight", + ], + ); + + assert!(summary.supports_native_mtp()); + } + + #[test] + fn native_mtp_artifact_summary_rejects_missing_nextn_metadata() { + let summary = native_mtp_artifact_summary_from_names( + 0, + [ + "blk.47.nextn.eh_proj", + "blk.47.nextn.enorm", + "blk.47.nextn.hnorm", + ], + ); + + assert!(!summary.supports_native_mtp()); + assert_eq!( + summary.missing_reasons(), + vec!["*.nextn_predict_layers > 0"] + ); + } +} + #[derive(Clone, Copy)] struct PackageStageSpec { topology_id: &'static str, diff --git a/crates/skippy-ffi/README.md b/crates/skippy-ffi/README.md index fd6263ddcd..2a139d43e5 100644 --- a/crates/skippy-ffi/README.md +++ b/crates/skippy-ffi/README.md @@ -179,7 +179,6 @@ hook currently bound by this crate. | --- | --- | | `skippy_abi_version` | Returns the compiled stage ABI version. The C header exports it; this Rust crate mirrors the version through constants. | | `skippy_abi_features` | Returns the compiled feature bitmask. Rust binds this function in `skippy-ffi`; higher-level consumers can use it for feature probing or gating. | -| `skippy_status_string` | Converts a status enum to a static C string. | | `skippy_error_free` | Frees an allocated `skippy_error`. | ### Model and session lifecycle @@ -209,12 +208,9 @@ hook currently bound by this crate. | Function | Purpose | | --- | --- | | `skippy_prefill_chunk` | Prefills a token chunk using raw activation buffers for staged input/output. | -| `skippy_decode_step` | Decodes one token using raw activation buffers and optionally returns a predicted token. | | `skippy_verify_tokens` | Runs batched token verification and returns the model-selected tokens. | | `skippy_decode_step_sampled` | Decodes one token with `SamplingConfig`, including penalties and logit bias. | | `skippy_prefill_chunk_frame` | Prefills a token chunk using `ActivationDesc` plus payload buffers. | -| `skippy_decode_step_frame` | Decodes one token using activation-frame descriptors and payloads. | -| `skippy_verify_tokens_frame` | Runs batched verification with activation-frame descriptors and payloads. | | `skippy_decode_step_frame_sampled` | Decodes one token with activation-frame I/O and `SamplingConfig`. | ### Token and chat helpers diff --git a/crates/skippy-ffi/src/lib.rs b/crates/skippy-ffi/src/lib.rs index cb5bc3ed4c..9160dfc81a 100644 --- a/crates/skippy-ffi/src/lib.rs +++ b/crates/skippy-ffi/src/lib.rs @@ -1,8 +1,9 @@ pub const ABI_VERSION_MAJOR: u32 = 0; pub const ABI_VERSION_MINOR: u32 = 1; -pub const ABI_VERSION_PATCH: u32 = 26; +pub const ABI_VERSION_PATCH: u32 = 27; pub const FEATURE_BACKEND_DEVICES: u64 = 1 << 23; pub const FEATURE_RUNTIME_EVENTS: u64 = 1 << 24; +pub const FEATURE_NATIVE_MTP_N1: u64 = 1 << 25; use std::ffi::{c_char, c_int, c_void}; @@ -328,6 +329,146 @@ pub struct LogitBias { pub bias: f32, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(i32)] +pub enum LlamaFileType { + AllF32 = 0, + MostlyF16 = 1, + MostlyQ4_0 = 2, + MostlyQ4_1 = 3, + MostlyQ8_0 = 7, + MostlyQ5_0 = 8, + MostlyQ5_1 = 9, + MostlyQ2K = 10, + MostlyQ3KS = 11, + MostlyQ3KM = 12, + MostlyQ3KL = 13, + MostlyQ4KS = 14, + MostlyQ4KM = 15, + MostlyQ5KS = 16, + MostlyQ5KM = 17, + MostlyQ6K = 18, + MostlyIQ2XXS = 19, + MostlyIQ2XS = 20, + MostlyQ2KS = 21, + MostlyIQ3XS = 22, + MostlyIQ3XXS = 23, + MostlyIQ1S = 24, + MostlyIQ4NL = 25, + MostlyIQ3S = 26, + MostlyIQ3M = 27, + MostlyIQ2S = 28, + MostlyIQ2M = 29, + MostlyIQ4XS = 30, + MostlyIQ1M = 31, + MostlyBf16 = 32, + MostlyTQ1_0 = 36, + MostlyTQ2_0 = 37, + MostlyMxfp4Moe = 38, + MostlyNvfp4 = 39, + MostlyQ1_0 = 40, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(i32)] +pub enum GgmlType { + F32 = 0, + F16 = 1, + Q4_0 = 2, + Q4_1 = 3, + Q5_0 = 6, + Q5_1 = 7, + Q8_0 = 8, + Q8_1 = 9, + Q2K = 10, + Q3K = 11, + Q4K = 12, + Q5K = 13, + Q6K = 14, + Q8K = 15, + IQ2XXS = 16, + IQ2XS = 17, + IQ3XXS = 18, + IQ1S = 19, + IQ4NL = 20, + IQ3S = 21, + IQ2S = 22, + IQ4XS = 23, + I8 = 24, + I16 = 25, + I32 = 26, + I64 = 27, + F64 = 28, + IQ1M = 29, + Bf16 = 30, + TQ1_0 = 34, + TQ2_0 = 35, + Mxfp4 = 39, + Nvfp4 = 40, + Q1_0 = 41, + Count = 42, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(i32)] +pub enum LlamaModelKvOverrideType { + Int = 0, + Float = 1, + Bool = 2, + Str = 3, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub union LlamaModelKvOverrideValue { + pub val_i64: i64, + pub val_f64: f64, + pub val_bool: bool, + pub val_str: [c_char; 128], +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct LlamaModelKvOverride { + pub tag: LlamaModelKvOverrideType, + pub key: [c_char; 128], + pub value: LlamaModelKvOverrideValue, +} + +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct LlamaModelTensorOverride { + pub pattern: *const c_char, + pub tensor_type: GgmlType, +} + +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct LlamaModelImatrixData { + pub name: *const c_char, + pub data: *const f32, + pub size: usize, +} + +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct LlamaModelQuantizeParams { + pub nthread: i32, + pub ftype: LlamaFileType, + pub output_tensor_type: GgmlType, + pub token_embedding_type: GgmlType, + pub allow_requantize: bool, + pub quantize_output_tensor: bool, + pub only_copy: bool, + pub pure: bool, + pub keep_split: bool, + pub dry_run: bool, + pub imatrix: *const LlamaModelImatrixData, + pub kv_overrides: *const LlamaModelKvOverride, + pub tt_overrides: *const LlamaModelTensorOverride, + pub prune_layers: *const i32, +} + #[repr(C)] #[derive(Debug, Clone, Copy)] pub struct SamplingConfig { @@ -364,6 +505,15 @@ pub struct KvPageDesc { pub flags: u64, } +#[repr(C)] +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct NativeMtpDraft { + pub version: u32, + pub available: bool, + pub token_id: i32, + pub proposal_compute_us: i64, +} + #[repr(C)] #[derive(Debug, Clone, Copy, Default, PartialEq)] pub struct TokenSignal { @@ -565,7 +715,8 @@ mod dynamic { dynamic_symbols! { llama_log_set(log_callback: LlamaLogCallback, user_data: *mut c_void); ggml_log_set(log_callback: LlamaLogCallback, user_data: *mut c_void); - skippy_status_string(status: Status) -> *const c_char; + llama_model_quantize_default_params() -> LlamaModelQuantizeParams; + llama_model_quantize(fname_inp: *const c_char, fname_out: *const c_char, params: *const LlamaModelQuantizeParams) -> u32; skippy_error_free(error: *mut Error); skippy_backend_device_count(out_count: *mut usize, out_error: *mut *mut Error) -> Status; skippy_backend_device_at(index: usize, out_device: *mut BackendDevice, out_error: *mut *mut Error) -> Status; @@ -577,7 +728,6 @@ mod dynamic { skippy_session_create_from_resident_prefix(model: *mut Model, cache_seq_id: i32, token_ids: *const i32, token_count: usize, out_session: *mut *mut Session, out_error: *mut *mut Error) -> Status; skippy_session_llama_context(session: *mut Session) -> *mut Opaque; skippy_session_position(session: *const Session) -> i32; - skippy_session_native_seq_id(session: *const Session) -> i32; skippy_session_batch_size(session: *const Session) -> i32; skippy_session_begin_external_decode(session: *mut Session, out_error: *mut *mut Error) -> Status; skippy_session_end_external_decode(session: *mut Session, out_error: *mut *mut Error) -> Status; @@ -589,16 +739,17 @@ mod dynamic { skippy_restore_session_checkpoint(session: *mut Session, token_count: u64, out_error: *mut *mut Error) -> Status; skippy_session_free(session: *mut Session, out_error: *mut *mut Error) -> Status; skippy_prefill_chunk(session: *mut Session, token_ids: *const i32, token_count: usize, input_activations: *const c_void, input_activation_bytes: usize, output_activations: *mut c_void, output_activation_capacity: usize, out_output_activation_bytes: *mut usize, out_error: *mut *mut Error) -> Status; - skippy_decode_step(session: *mut Session, token_id: i32, input_activation: *const c_void, input_activation_bytes: usize, output_activation: *mut c_void, output_activation_capacity: usize, out_output_activation_bytes: *mut usize, out_predicted_token: *mut i32, out_error: *mut *mut Error) -> Status; skippy_verify_tokens(session: *mut Session, token_ids: *const i32, token_count: usize, output_tokens: *mut i32, output_token_capacity: usize, out_token_count: *mut usize, out_error: *mut *mut Error) -> Status; skippy_decode_step_sampled(session: *mut Session, token_id: i32, sampling: *const SamplingConfig, input_activation: *const c_void, input_activation_bytes: usize, output_activation: *mut c_void, output_activation_capacity: usize, out_output_activation_bytes: *mut usize, out_predicted_token: *mut i32, out_error: *mut *mut Error) -> Status; + skippy_decode_batch_sampled(sessions: *const *mut Session, token_ids: *const i32, sampling: *const *const SamplingConfig, request_count: usize, out_predicted_tokens: *mut i32, predicted_token_capacity: usize, out_error: *mut *mut Error) -> Status; skippy_prefill_chunk_frame(session: *mut Session, token_ids: *const i32, token_count: usize, input_desc: *const ActivationDesc, input_payload: *const c_void, output_desc: *mut ActivationDesc, output_payload: *mut c_void, output_payload_capacity: usize, out_output_payload_bytes: *mut usize, out_error: *mut *mut Error) -> Status; skippy_prefill_chunk_frame_sampled(session: *mut Session, token_ids: *const i32, token_count: usize, sampling: *const SamplingConfig, input_desc: *const ActivationDesc, input_payload: *const c_void, output_desc: *mut ActivationDesc, output_payload: *mut c_void, output_payload_capacity: usize, out_output_payload_bytes: *mut usize, out_predicted_token: *mut i32, out_error: *mut *mut Error) -> Status; skippy_prefill_chunk_frame_with_positions(session: *mut Session, token_ids: *const i32, token_count: usize, positions: *const i32, position_count: usize, input_desc: *const ActivationDesc, input_payload: *const c_void, output_desc: *mut ActivationDesc, output_payload: *mut c_void, output_payload_capacity: usize, out_output_payload_bytes: *mut usize, out_error: *mut *mut Error) -> Status; skippy_prefill_chunk_frame_sampled_with_positions(session: *mut Session, token_ids: *const i32, token_count: usize, positions: *const i32, position_count: usize, sampling: *const SamplingConfig, input_desc: *const ActivationDesc, input_payload: *const c_void, output_desc: *mut ActivationDesc, output_payload: *mut c_void, output_payload_capacity: usize, out_output_payload_bytes: *mut usize, out_predicted_token: *mut i32, out_error: *mut *mut Error) -> Status; - skippy_decode_step_frame(session: *mut Session, token_id: i32, input_desc: *const ActivationDesc, input_payload: *const c_void, output_desc: *mut ActivationDesc, output_payload: *mut c_void, output_payload_capacity: usize, out_output_payload_bytes: *mut usize, out_predicted_token: *mut i32, out_error: *mut *mut Error) -> Status; - skippy_verify_tokens_frame(session: *mut Session, token_ids: *const i32, token_count: usize, input_desc: *const ActivationDesc, input_payload: *const c_void, output_desc: *mut ActivationDesc, output_payload: *mut c_void, output_payload_capacity: usize, out_output_payload_bytes: *mut usize, output_tokens: *mut i32, output_token_capacity: usize, out_token_count: *mut usize, out_error: *mut *mut Error) -> Status; skippy_decode_step_frame_sampled(session: *mut Session, token_id: i32, sampling: *const SamplingConfig, input_desc: *const ActivationDesc, input_payload: *const c_void, output_desc: *mut ActivationDesc, output_payload: *mut c_void, output_payload_capacity: usize, out_output_payload_bytes: *mut usize, out_predicted_token: *mut i32, out_error: *mut *mut Error) -> Status; + skippy_decode_step_frame_sampled_mtp_n1(session: *mut Session, token_id: i32, sampling: *const SamplingConfig, input_desc: *const ActivationDesc, input_payload: *const c_void, output_desc: *mut ActivationDesc, output_payload: *mut c_void, output_payload_capacity: usize, out_output_payload_bytes: *mut usize, out_predicted_token: *mut i32, out_mtp_draft: *mut NativeMtpDraft, out_error: *mut *mut Error) -> Status; + skippy_decode_step_frame_batch_sampled(sessions: *const *mut Session, token_ids: *const i32, sampling: *const *const SamplingConfig, input_descs: *const *const ActivationDesc, input_payloads: *const *const c_void, output_descs: *mut ActivationDesc, output_payloads: *const *mut c_void, output_payload_capacities: *const usize, out_output_payload_bytes: *mut usize, out_predicted_tokens: *mut i32, predicted_token_capacity: usize, request_count: usize, out_error: *mut *mut Error) -> Status; + skippy_verify_tokens_frame_sampled(session: *mut Session, token_ids: *const i32, token_count: usize, sampling: *const SamplingConfig, input_desc: *const ActivationDesc, input_payload: *const c_void, output_desc: *mut ActivationDesc, output_payload: *mut c_void, output_payload_capacity: usize, out_output_payload_bytes: *mut usize, output_tokens: *mut i32, output_token_capacity: usize, out_token_count: *mut usize, out_error: *mut *mut Error) -> Status; skippy_session_copy_output_activation_frame(session: *mut Session, token_count: usize, output_desc: *mut ActivationDesc, output_payload: *mut c_void, output_payload_capacity: usize, out_output_payload_bytes: *mut usize, out_error: *mut *mut Error) -> Status; skippy_session_last_token_signal(session: *mut Session, out_signal: *mut TokenSignal, out_error: *mut *mut Error) -> Status; skippy_session_signal_window(session: *mut Session, window_tokens: u32, out_window: *mut GenerationSignalWindow, out_error: *mut *mut Error) -> Status; @@ -617,9 +768,6 @@ mod dynamic { skippy_tokenize(model: *mut Model, text: *const c_char, add_special: bool, output_tokens: *mut i32, output_token_capacity: usize, out_token_count: *mut usize, out_error: *mut *mut Error) -> Status; skippy_detokenize(model: *mut Model, tokens: *const i32, token_count: usize, output_text: *mut c_char, output_text_capacity: usize, out_text_bytes: *mut usize, out_error: *mut *mut Error) -> Status; skippy_token_is_eog(model: *mut Model, token_id: i32, out_is_eog: *mut bool, out_error: *mut *mut Error) -> Status; - skippy_apply_chat_template(model: *mut Model, messages: *const ChatMessage, message_count: usize, add_assistant: bool, override_enable_thinking: bool, enable_thinking: bool, output_text: *mut c_char, output_text_capacity: usize, out_text_bytes: *mut usize, out_error: *mut *mut Error) -> Status; - skippy_apply_chat_template_json(model: *mut Model, messages_json: *const c_char, tools_json: *const c_char, tool_choice_json: *const c_char, add_assistant: bool, override_enable_thinking: bool, enable_thinking: bool, parallel_tool_calls: bool, output_text: *mut c_char, output_text_capacity: usize, out_text_bytes: *mut usize, output_metadata_json: *mut c_char, output_metadata_json_capacity: usize, out_metadata_json_bytes: *mut usize, out_error: *mut *mut Error) -> Status; - skippy_parse_chat_response_json(generated_text: *const c_char, metadata_json: *const c_char, is_partial: bool, output_message_json: *mut c_char, output_message_json_capacity: usize, out_message_json_bytes: *mut usize, out_error: *mut *mut Error) -> Status; skippy_model_info_open(path: *const c_char, out_info: *mut *mut ModelInfo, out_error: *mut *mut Error) -> Status; skippy_model_info_free(info: *mut ModelInfo, out_error: *mut *mut Error) -> Status; skippy_model_info_tensor_count(info: *mut ModelInfo, out_count: *mut usize, out_error: *mut *mut Error) -> Status; @@ -673,6 +821,44 @@ mod dynamic { out_model: *mut *mut Model, out_error: *mut *mut Error, ) -> Status; + type SkippyApplyChatTemplateFn = unsafe extern "C" fn( + model: *mut Model, + messages: *const ChatMessage, + message_count: usize, + add_assistant: bool, + override_enable_thinking: bool, + enable_thinking: bool, + output_text: *mut c_char, + output_text_capacity: usize, + out_text_bytes: *mut usize, + out_error: *mut *mut Error, + ) -> Status; + type SkippyApplyChatTemplateJsonFn = unsafe extern "C" fn( + model: *mut Model, + messages_json: *const c_char, + tools_json: *const c_char, + tool_choice_json: *const c_char, + add_assistant: bool, + override_enable_thinking: bool, + enable_thinking: bool, + parallel_tool_calls: bool, + output_text: *mut c_char, + output_text_capacity: usize, + out_text_bytes: *mut usize, + output_metadata_json: *mut c_char, + output_metadata_json_capacity: usize, + out_metadata_json_bytes: *mut usize, + out_error: *mut *mut Error, + ) -> Status; + type SkippyParseChatResponseJsonFn = unsafe extern "C" fn( + generated_text: *const c_char, + metadata_json: *const c_char, + is_partial: bool, + output_message_json: *mut c_char, + output_message_json_capacity: usize, + out_message_json_bytes: *mut usize, + out_error: *mut *mut Error, + ) -> Status; impl Symbols { fn lookup_optional(&self, name: &[u8]) -> Option @@ -712,6 +898,131 @@ mod dynamic { ) }) } + + fn skippy_apply_chat_template_fn() -> Option { + static CACHE: OnceLock> = OnceLock::new(); + *CACHE.get_or_init(|| { + symbols().lookup_optional::(b"skippy_apply_chat_template\0") + }) + } + + fn skippy_apply_chat_template_json_fn() -> Option { + static CACHE: OnceLock> = OnceLock::new(); + *CACHE.get_or_init(|| { + symbols().lookup_optional::( + b"skippy_apply_chat_template_json\0", + ) + }) + } + + fn skippy_parse_chat_response_json_fn() -> Option { + static CACHE: OnceLock> = OnceLock::new(); + *CACHE.get_or_init(|| { + symbols().lookup_optional::( + b"skippy_parse_chat_response_json\0", + ) + }) + } + + #[allow(clippy::missing_safety_doc, clippy::too_many_arguments)] + pub unsafe fn skippy_apply_chat_template( + model: *mut Model, + messages: *const ChatMessage, + message_count: usize, + add_assistant: bool, + override_enable_thinking: bool, + enable_thinking: bool, + output_text: *mut c_char, + output_text_capacity: usize, + out_text_bytes: *mut usize, + out_error: *mut *mut Error, + ) -> Status { + let Some(function) = skippy_apply_chat_template_fn() else { + return Status::Unsupported; + }; + unsafe { + function( + model, + messages, + message_count, + add_assistant, + override_enable_thinking, + enable_thinking, + output_text, + output_text_capacity, + out_text_bytes, + out_error, + ) + } + } + + #[allow(clippy::missing_safety_doc, clippy::too_many_arguments)] + pub unsafe fn skippy_apply_chat_template_json( + model: *mut Model, + messages_json: *const c_char, + tools_json: *const c_char, + tool_choice_json: *const c_char, + add_assistant: bool, + override_enable_thinking: bool, + enable_thinking: bool, + parallel_tool_calls: bool, + output_text: *mut c_char, + output_text_capacity: usize, + out_text_bytes: *mut usize, + output_metadata_json: *mut c_char, + output_metadata_json_capacity: usize, + out_metadata_json_bytes: *mut usize, + out_error: *mut *mut Error, + ) -> Status { + let Some(function) = skippy_apply_chat_template_json_fn() else { + return Status::Unsupported; + }; + unsafe { + function( + model, + messages_json, + tools_json, + tool_choice_json, + add_assistant, + override_enable_thinking, + enable_thinking, + parallel_tool_calls, + output_text, + output_text_capacity, + out_text_bytes, + output_metadata_json, + output_metadata_json_capacity, + out_metadata_json_bytes, + out_error, + ) + } + } + + #[allow(clippy::missing_safety_doc, clippy::too_many_arguments)] + pub unsafe fn skippy_parse_chat_response_json( + generated_text: *const c_char, + metadata_json: *const c_char, + is_partial: bool, + output_message_json: *mut c_char, + output_message_json_capacity: usize, + out_message_json_bytes: *mut usize, + out_error: *mut *mut Error, + ) -> Status { + let Some(function) = skippy_parse_chat_response_json_fn() else { + return Status::Unsupported; + }; + unsafe { + function( + generated_text, + metadata_json, + is_partial, + output_message_json, + output_message_json_capacity, + out_message_json_bytes, + out_error, + ) + } + } } #[cfg(feature = "dynamic-runtime")] @@ -732,9 +1043,16 @@ unsafe extern "C" { pub fn ggml_log_set(log_callback: LlamaLogCallback, user_data: *mut c_void); + pub fn llama_model_quantize_default_params() -> LlamaModelQuantizeParams; + + pub fn llama_model_quantize( + fname_inp: *const c_char, + fname_out: *const c_char, + params: *const LlamaModelQuantizeParams, + ) -> u32; + pub fn skippy_abi_features() -> u64; - pub fn skippy_status_string(status: Status) -> *const c_char; pub fn skippy_error_free(error: *mut Error); pub fn skippy_backend_device_count(out_count: *mut usize, out_error: *mut *mut Error) @@ -784,8 +1102,6 @@ unsafe extern "C" { pub fn skippy_session_position(session: *const Session) -> i32; - pub fn skippy_session_native_seq_id(session: *const Session) -> i32; - pub fn skippy_session_batch_size(session: *const Session) -> i32; pub fn skippy_session_begin_external_decode( @@ -847,18 +1163,6 @@ unsafe extern "C" { out_error: *mut *mut Error, ) -> Status; - pub fn skippy_decode_step( - session: *mut Session, - token_id: i32, - input_activation: *const c_void, - input_activation_bytes: usize, - output_activation: *mut c_void, - output_activation_capacity: usize, - out_output_activation_bytes: *mut usize, - out_predicted_token: *mut i32, - out_error: *mut *mut Error, - ) -> Status; - pub fn skippy_verify_tokens( session: *mut Session, token_ids: *const i32, @@ -882,6 +1186,16 @@ unsafe extern "C" { out_error: *mut *mut Error, ) -> Status; + pub fn skippy_decode_batch_sampled( + sessions: *const *mut Session, + token_ids: *const i32, + sampling: *const *const SamplingConfig, + request_count: usize, + out_predicted_tokens: *mut i32, + predicted_token_capacity: usize, + out_error: *mut *mut Error, + ) -> Status; + pub fn skippy_prefill_chunk_frame( session: *mut Session, token_ids: *const i32, @@ -942,36 +1256,38 @@ unsafe extern "C" { out_error: *mut *mut Error, ) -> Status; - pub fn skippy_decode_step_frame( + pub fn skippy_verify_tokens_frame_sampled( session: *mut Session, - token_id: i32, + token_ids: *const i32, + token_count: usize, + sampling: *const SamplingConfig, input_desc: *const ActivationDesc, input_payload: *const c_void, output_desc: *mut ActivationDesc, output_payload: *mut c_void, output_payload_capacity: usize, out_output_payload_bytes: *mut usize, - out_predicted_token: *mut i32, + output_tokens: *mut i32, + output_token_capacity: usize, + out_token_count: *mut usize, out_error: *mut *mut Error, ) -> Status; - pub fn skippy_verify_tokens_frame( + pub fn skippy_decode_step_frame_sampled( session: *mut Session, - token_ids: *const i32, - token_count: usize, + token_id: i32, + sampling: *const SamplingConfig, input_desc: *const ActivationDesc, input_payload: *const c_void, output_desc: *mut ActivationDesc, output_payload: *mut c_void, output_payload_capacity: usize, out_output_payload_bytes: *mut usize, - output_tokens: *mut i32, - output_token_capacity: usize, - out_token_count: *mut usize, + out_predicted_token: *mut i32, out_error: *mut *mut Error, ) -> Status; - pub fn skippy_decode_step_frame_sampled( + pub fn skippy_decode_step_frame_sampled_mtp_n1( session: *mut Session, token_id: i32, sampling: *const SamplingConfig, @@ -982,6 +1298,23 @@ unsafe extern "C" { output_payload_capacity: usize, out_output_payload_bytes: *mut usize, out_predicted_token: *mut i32, + out_mtp_draft: *mut NativeMtpDraft, + out_error: *mut *mut Error, + ) -> Status; + + pub fn skippy_decode_step_frame_batch_sampled( + sessions: *const *mut Session, + token_ids: *const i32, + sampling: *const *const SamplingConfig, + input_descs: *const *const ActivationDesc, + input_payloads: *const *const c_void, + output_descs: *mut ActivationDesc, + output_payloads: *const *mut c_void, + output_payload_capacities: *const usize, + out_output_payload_bytes: *mut usize, + out_predicted_tokens: *mut i32, + predicted_token_capacity: usize, + request_count: usize, out_error: *mut *mut Error, ) -> Status; diff --git a/crates/skippy-model-package/src/main.rs b/crates/skippy-model-package/src/main.rs index c2270ddd4b..14d0a252aa 100644 --- a/crates/skippy-model-package/src/main.rs +++ b/crates/skippy-model-package/src/main.rs @@ -1154,16 +1154,22 @@ fn write_package_artifact( ); let path = out_dir.join(&spec.relative_path); write_stage_artifact(source, &stage, &path)?; + let relative_path = spec.relative_path.display().to_string(); + run_artifact_hook(artifact_hook, &path, &relative_path)?; + let artifact_info = ModelInfo::open(&path) + .with_context(|| format!("open package artifact {}", path.display()))?; + let artifact_tensors = artifact_info + .tensors() + .with_context(|| format!("read package artifact tensors {}", path.display()))?; let metadata = fs::metadata(&path) .with_context(|| format!("read artifact metadata {}", path.display()))?; let artifact = PackageArtifact { - path: spec.relative_path.display().to_string(), - tensor_count: stage.tensor_count, - tensor_bytes: stage.tensor_bytes, + path: relative_path, + tensor_count: artifact_tensors.len(), + tensor_bytes: artifact_tensors.iter().map(|tensor| tensor.byte_size).sum(), artifact_bytes: metadata.len(), sha256: file_sha256(&path)?, }; - run_artifact_hook(artifact_hook, &path, &artifact.path)?; Ok(artifact) } diff --git a/crates/skippy-prompt/src/prompt_cli/args.rs b/crates/skippy-prompt/src/prompt_cli/args.rs index a3b207e6f9..b878f01777 100644 --- a/crates/skippy-prompt/src/prompt_cli/args.rs +++ b/crates/skippy-prompt/src/prompt_cli/args.rs @@ -184,8 +184,6 @@ pub struct BinaryReplArgs { pub tokenizer_n_gpu_layers: i32, #[arg(long, default_value = "127.0.0.1:19031")] pub first_stage_addr: String, - #[arg(long, default_value = "127.0.0.1:19030")] - pub direct_return_bind_addr: SocketAddr, #[arg(long, default_value_t = 0)] pub tokenizer_layer_start: u32, #[arg(long, default_value_t = 10)] diff --git a/crates/skippy-prompt/src/prompt_cli/binary_repl.rs b/crates/skippy-prompt/src/prompt_cli/binary_repl.rs index 6f96a4266c..14fdf89287 100644 --- a/crates/skippy-prompt/src/prompt_cli/binary_repl.rs +++ b/crates/skippy-prompt/src/prompt_cli/binary_repl.rs @@ -108,11 +108,6 @@ pub fn binary_repl(args: BinaryReplArgs) -> Result<()> { format_prompt_max_new_tokens(args.max_new_tokens), args.prefill_chunk_size ); - let direct_returns = PromptDirectReturnServer::start(args.direct_return_bind_addr)?; - eprintln!( - "direct prediction return listener: {}", - direct_returns.endpoint() - ); if let Some(draft) = draft.as_ref() { eprintln!( "draft model enabled: {} speculative_window={}", @@ -211,7 +206,6 @@ pub fn binary_repl(args: BinaryReplArgs) -> Result<()> { prompt_index, prompt: &prompt, live_session: None, - direct_returns: &direct_returns, }) .or_else(|error| handle_prompt_error(error, &interrupt, prompt_index))?; prompt_index += 1; @@ -277,7 +271,6 @@ pub fn binary_repl(args: BinaryReplArgs) -> Result<()> { prompt_index, prompt: &prompt, live_session: None, - direct_returns: &direct_returns, }) .or_else(|error| handle_prompt_error(error, &interrupt, prompt_index))?; prompt_index += 1; @@ -298,7 +291,6 @@ pub fn binary_repl(args: BinaryReplArgs) -> Result<()> { prompt_index, prompt: input, live_session: append_transcript.then_some(&mut live_session), - direct_returns: &direct_returns, }); if prompt_result.is_err() { live_session.mark_dirty(); diff --git a/crates/skippy-prompt/src/prompt_cli/direct_return.rs b/crates/skippy-prompt/src/prompt_cli/direct_return.rs deleted file mode 100644 index 68e000cb16..0000000000 --- a/crates/skippy-prompt/src/prompt_cli/direct_return.rs +++ /dev/null @@ -1,151 +0,0 @@ -#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] -struct PromptDirectReturnKey { - request_id: u64, - session_id: u64, -} - -type PromptDirectReturnResult = Result; -type PromptDirectReturnSender = mpsc::Sender; -type PromptDirectReturnWaiters = - Arc>>; - -struct PromptDirectReturnServer { - local_addr: SocketAddr, - waiters: PromptDirectReturnWaiters, -} - -impl PromptDirectReturnServer { - fn start(bind_addr: SocketAddr) -> Result { - let listener = TcpListener::bind(bind_addr) - .with_context(|| format!("bind prompt direct-return listener {bind_addr}"))?; - let local_addr = listener - .local_addr() - .context("read prompt direct-return listener address")?; - let waiters = Arc::new(Mutex::new(HashMap::new())); - let thread_waiters = waiters.clone(); - thread::spawn(move || { - for stream in listener.incoming() { - match stream { - Ok(stream) => { - let waiters = thread_waiters.clone(); - thread::spawn(move || { - if let Err(error) = - handle_prompt_direct_return_connection(waiters, stream) - { - eprintln!("prompt direct-return connection failed: {error:#}"); - } - }); - } - Err(error) if error.kind() == io::ErrorKind::Interrupted => continue, - Err(error) => { - eprintln!("prompt direct-return listener failed: {error}"); - break; - } - } - } - }); - Ok(Self { - local_addr, - waiters, - }) - } - - fn endpoint(&self) -> SocketAddr { - self.local_addr - } - - fn register( - &self, - request_id: u64, - session_id: u64, - timeout: Duration, - ) -> Result { - let key = PromptDirectReturnKey { - request_id, - session_id, - }; - let (sender, receiver) = mpsc::channel(); - self.waiters - .lock() - .map_err(|_| anyhow!("prompt direct-return hub lock poisoned"))? - .insert(key, sender); - Ok(PromptDirectReturnReceiver { - key, - waiters: self.waiters.clone(), - receiver, - timeout, - }) - } -} - -struct PromptDirectReturnReceiver { - key: PromptDirectReturnKey, - waiters: PromptDirectReturnWaiters, - receiver: mpsc::Receiver, - timeout: Duration, -} - -impl PromptDirectReturnReceiver { - fn recv_expected(&self, expected: WireReplyKind) -> Result { - let reply = self - .receiver - .recv_timeout(self.timeout) - .context("timed out waiting for prompt direct prediction return")? - .map_err(|error| anyhow!(error))?; - if reply.kind != expected { - bail!( - "expected {expected:?} direct prediction return, got {:?}", - reply.kind - ); - } - Ok(reply) - } -} - -impl Drop for PromptDirectReturnReceiver { - fn drop(&mut self) { - if let Ok(mut waiters) = self.waiters.lock() { - waiters.remove(&self.key); - } - } -} - -fn handle_prompt_direct_return_connection( - waiters: PromptDirectReturnWaiters, - mut stream: TcpStream, -) -> Result<()> { - send_ready(&mut stream).context("send prompt direct-return ready")?; - let open = read_stage_message(&mut stream, 0).context("read prompt direct-return open")?; - if open.kind != WireMessageKind::PredictionReturnOpen { - bail!("expected prediction-return-open message"); - } - let key = PromptDirectReturnKey { - request_id: open.request_id, - session_id: open.session_id, - }; - let sender = waiters - .lock() - .map_err(|_| anyhow!("prompt direct-return hub lock poisoned"))? - .get(&key) - .cloned() - .ok_or_else(|| { - anyhow!( - "no prompt direct-return waiter for request {}", - key.request_id - ) - })?; - loop { - match recv_reply(&mut stream) { - Ok(reply) => { - if sender.send(Ok(reply)).is_err() { - return Ok(()); - } - } - Err(error) if error.kind() == io::ErrorKind::UnexpectedEof => return Ok(()), - Err(error) => { - let _ = sender.send(Err(error.to_string())); - return Err(error).context("read prompt direct prediction return"); - } - } - } -} diff --git a/crates/skippy-prompt/src/prompt_cli/generation.rs b/crates/skippy-prompt/src/prompt_cli/generation.rs index 5183c419c3..eaf01404fb 100644 --- a/crates/skippy-prompt/src/prompt_cli/generation.rs +++ b/crates/skippy-prompt/src/prompt_cli/generation.rs @@ -11,7 +11,6 @@ struct PromptRun<'a> { prompt_index: usize, prompt: &'a str, live_session: Option<&'a mut PromptLiveSession>, - direct_returns: &'a PromptDirectReturnServer, } fn run_prompt(run: PromptRun<'_>) -> Result<()> { @@ -28,7 +27,6 @@ fn run_prompt(run: PromptRun<'_>) -> Result<()> { prompt_index, prompt, mut live_session, - direct_returns, } = run; if args.prefill_chunk_size == 0 { @@ -69,9 +67,6 @@ fn run_prompt(run: PromptRun<'_>) -> Result<()> { ); let prompt_index_bytes = prompt_index.to_le_bytes(); let request_id = stable_wire_id(&[session_id.as_bytes(), &prompt_index_bytes]); - let direct_return_timeout = Duration::from_secs(args.decode_timeout_secs.max(1)); - let direct_return = - direct_returns.register(request_id, wire_session_id, direct_return_timeout)?; let mut session_reuse = PromptSessionReuseStats::default(); let mut one_shot_stream = None; @@ -337,7 +332,6 @@ fn run_prompt(run: PromptRun<'_>) -> Result<()> { prefill_token_count, decode_index, current, - &direct_return, ) .with_context(|| stage_chain_error_context(args))?; decode_ms += reply.elapsed_ms; @@ -385,7 +379,6 @@ fn run_prompt(run: PromptRun<'_>) -> Result<()> { decode_index, &verify_inputs, true, - &direct_return, ) .with_context(|| stage_chain_error_context(args))?; decode_ms += reply.elapsed_ms; @@ -439,7 +432,6 @@ fn run_prompt(run: PromptRun<'_>) -> Result<()> { prefill_token_count, decode_index, current, - &direct_return, ) .with_context(|| stage_chain_error_context(args))?; commit_tokens = vec![repair.predicted]; @@ -461,7 +453,6 @@ fn run_prompt(run: PromptRun<'_>) -> Result<()> { decode_index, repair_inputs, false, - &direct_return, ) .with_context(|| stage_chain_error_context(args))?; commit_tokens = repaired_commit_tokens( diff --git a/crates/skippy-prompt/src/prompt_cli/launch.rs b/crates/skippy-prompt/src/prompt_cli/launch.rs index 9e9c5478d6..1fb762a279 100644 --- a/crates/skippy-prompt/src/prompt_cli/launch.rs +++ b/crates/skippy-prompt/src/prompt_cli/launch.rs @@ -345,10 +345,6 @@ fn prompt_repl_launch(args: PromptArgs) -> Result<()> { tokenizer_load_mode, tokenizer_n_gpu_layers: 0, first_stage_addr: stages[0].endpoint_addr.clone(), - direct_return_bind_addr: SocketAddr::from(( - [127, 0, 0, 1], - args.first_stage_port.saturating_sub(1), - )), tokenizer_layer_start: stages[0].layer_start as u32, tokenizer_layer_end: stages[0].layer_end, ctx_size: args.ctx_size, diff --git a/crates/skippy-prompt/src/prompt_cli/mod.rs b/crates/skippy-prompt/src/prompt_cli/mod.rs index 83a5c5964f..8fee132d62 100644 --- a/crates/skippy-prompt/src/prompt_cli/mod.rs +++ b/crates/skippy-prompt/src/prompt_cli/mod.rs @@ -1,9 +1,9 @@ use std::{ - collections::{BTreeMap, BTreeSet, HashMap, VecDeque, hash_map::DefaultHasher}, + collections::{BTreeMap, BTreeSet, VecDeque, hash_map::DefaultHasher}, fs, hash::{Hash, Hasher}, io::{self, BufRead, BufReader, IsTerminal, Read, Write}, - net::{Shutdown, SocketAddr, TcpListener, TcpStream}, + net::{Shutdown, SocketAddr, TcpStream}, path::{Component, Path, PathBuf}, process::{Child, Command, Stdio}, sync::{ @@ -23,8 +23,8 @@ use rustyline::{DefaultEditor, error::ReadlineError}; use serde_json::Value; use skippy_protocol::binary::{ LLAMA_TOKEN_NULL, READY_MAGIC, StageReply, StageReplyStats, StageStateHeader, StageWireMessage, - WireActivationDType, WireMessageKind, WireReplyKind, read_stage_message, recv_reply, - send_ready, state_flags, write_stage_message, + WireActivationDType, WireMessageKind, WireReplyKind, recv_reply, state_flags, + write_stage_message, }; use skippy_protocol::{ FlashAttentionType as StageFlashAttentionType, LoadMode, PeerConfig, StageConfig, @@ -57,7 +57,6 @@ include!("args.rs"); include!("command.rs"); include!("launch.rs"); include!("interrupt.rs"); -include!("direct_return.rs"); include!("binary_repl.rs"); include!("logs.rs"); include!("prompt_format.rs"); diff --git a/crates/skippy-prompt/src/prompt_cli/wire_messages.rs b/crates/skippy-prompt/src/prompt_cli/wire_messages.rs index aaf4471011..3d7098ba09 100644 --- a/crates/skippy-prompt/src/prompt_cli/wire_messages.rs +++ b/crates/skippy-prompt/src/prompt_cli/wire_messages.rs @@ -28,7 +28,6 @@ fn send_decode_step( prefill_token_count: usize, decode_index: usize, current: i32, - direct_return: &PromptDirectReturnReceiver, ) -> Result { let decode_started = Instant::now(); let mut state = StageStateHeader::new(WireMessageKind::DecodeEmbd, wire_dtype); @@ -55,9 +54,9 @@ fn send_decode_step( }; write_stage_message(&mut *stream, &message, wire_dtype) .with_context(|| format!("send decode step {decode_index}"))?; - let reply = direct_return - .recv_expected(WireReplyKind::PredictedToken) + let reply = recv_reply(&mut *stream) .with_context(|| format!("receive decode step {decode_index} reply"))?; + ensure_reply_kind(&reply, WireReplyKind::PredictedToken)?; Ok(DecodeStepReply { predicted: reply.predicted, stats: reply.stats, @@ -77,7 +76,6 @@ fn send_verify_span( decode_index: usize, tokens: &[i32], checkpoint: bool, - direct_return: &PromptDirectReturnReceiver, ) -> Result { if tokens.is_empty() { bail!("verify span requires at least one token"); @@ -112,9 +110,9 @@ fn send_verify_span( .with_context(|| format!("send verify span at decode step {decode_index}"))?; let write_ms = elapsed_ms(write_started); let wait_started = Instant::now(); - let reply = direct_return - .recv_expected(WireReplyKind::PredictedTokens) + let reply = recv_reply(&mut *stream) .with_context(|| format!("receive verify span {decode_index} reply"))?; + ensure_reply_kind(&reply, WireReplyKind::PredictedTokens)?; let wait_ms = elapsed_ms(wait_started); Ok(VerifySpanReply { predicted_tokens: reply.predicted_tokens, @@ -125,6 +123,13 @@ fn send_verify_span( }) } +fn ensure_reply_kind(reply: &StageReply, expected: WireReplyKind) -> Result<()> { + if reply.kind != expected { + bail!("expected {expected:?} reply, got {:?}", reply.kind); + } + Ok(()) +} + fn send_generation_config( stream: &mut TcpStream, wire_dtype: WireActivationDType, diff --git a/crates/skippy-protocol/src/binary/codec.rs b/crates/skippy-protocol/src/binary/codec.rs index 328cde78fc..59739a9d6f 100644 --- a/crates/skippy-protocol/src/binary/codec.rs +++ b/crates/skippy-protocol/src/binary/codec.rs @@ -44,10 +44,28 @@ pub fn send_reply_predicted_with_stats( predicted: i32, stats: StageReplyStats, ) -> io::Result<()> { + send_reply_predicted_with_tokens_and_stats(&mut writer, predicted, &[predicted], stats) +} + +pub fn send_reply_predicted_with_tokens_and_stats( + mut writer: impl Write, + predicted: i32, + predicted_tokens: &[i32], + stats: StageReplyStats, +) -> io::Result<()> { + if predicted_tokens.len() > MAX_STAGE_PREDICTED_TOKENS { + return Err(invalid_input("too many predicted tokens")); + } write_i32(&mut writer, WireReplyKind::PredictedToken as i32)?; write_i32(&mut writer, predicted)?; - write_i32(&mut writer, 1)?; - write_i32(&mut writer, predicted)?; + write_i32( + &mut writer, + i32::try_from(predicted_tokens.len()) + .map_err(|_| invalid_input("too many predicted tokens"))?, + )?; + for token in predicted_tokens { + write_i32(&mut writer, *token)?; + } write_reply_stats(&mut writer, stats) } @@ -417,90 +435,120 @@ fn read_sampling_config(mut reader: impl Read) -> io::Result(); + fn write_reply_stats(mut writer: impl Write, stats: StageReplyStats) -> io::Result<()> { - write_i64(&mut writer, stats.kv_lookup_hits)?; - write_i64(&mut writer, stats.kv_lookup_misses)?; - write_i64(&mut writer, stats.kv_lookup_errors)?; - write_i64(&mut writer, stats.kv_imported_pages)?; - write_i64(&mut writer, stats.kv_imported_tokens)?; - write_i64(&mut writer, stats.kv_recorded_pages)?; - write_i64(&mut writer, stats.kv_recorded_bytes)?; - write_i64(&mut writer, stats.kv_hit_stage_mask)?; - write_i64(&mut writer, stats.kv_record_stage_mask)?; - write_i64(&mut writer, stats.checkpoint_flush_us)?; - write_i64(&mut writer, stats.checkpoint_prefill_drain_us)?; - write_i64(&mut writer, stats.checkpoint_local_us)?; - write_i64(&mut writer, stats.checkpoint_downstream_write_us)?; - write_i64(&mut writer, stats.checkpoint_downstream_wait_us)?; - write_i64(&mut writer, stats.checkpoint_total_us)?; - write_i64(&mut writer, stats.checkpoint_prefill_drained_replies)?; - write_i64(&mut writer, stats.restore_flush_us)?; - write_i64(&mut writer, stats.restore_prefill_drain_us)?; - write_i64(&mut writer, stats.restore_local_us)?; - write_i64(&mut writer, stats.restore_downstream_write_us)?; - write_i64(&mut writer, stats.restore_downstream_wait_us)?; - write_i64(&mut writer, stats.restore_total_us)?; - write_i64(&mut writer, stats.restore_prefill_drained_replies)?; - write_i64(&mut writer, stats.verify_span_compute_us)?; - write_i64(&mut writer, stats.verify_span_forward_write_us)?; - write_i64(&mut writer, stats.verify_span_downstream_wait_us)?; - write_i64(&mut writer, stats.verify_span_total_us)?; - write_i64(&mut writer, stats.verify_span_stage_count)?; - write_i64(&mut writer, stats.verify_span_request_count)?; - write_i64(&mut writer, stats.verify_span_token_count)?; - write_i64(&mut writer, stats.verify_span_max_tokens)?; - write_i64(&mut writer, stats.verify_span_checkpointed_requests)?; - write_i64(&mut writer, stats.verify_span_skip_checkpoint_requests)?; - write_i64(&mut writer, stats.prefill_edge_write_us_max)?; - write_i64(&mut writer, stats.prefill_edge_wait_us_max)?; - write_i64(&mut writer, stats.prefill_edge_total_us_max)?; - write_i64(&mut writer, stats.prefill_edge_stage_index)?; - write_i64(&mut writer, stats.prefill_edge_activation_bytes_max)?; - write_i64(&mut writer, stats.prefill_edge_observation_count) + let fields = reply_stats_fields(stats); + let mut bytes = [0_u8; REPLY_STATS_WIRE_BYTES]; + for (chunk, value) in bytes + .chunks_exact_mut(std::mem::size_of::()) + .zip(fields) + { + chunk.copy_from_slice(&value.to_le_bytes()); + } + writer.write_all(&bytes) } fn read_reply_stats(mut reader: impl Read) -> io::Result { - Ok(StageReplyStats { - kv_lookup_hits: read_i64(&mut reader)?, - kv_lookup_misses: read_i64(&mut reader)?, - kv_lookup_errors: read_i64(&mut reader)?, - kv_imported_pages: read_i64(&mut reader)?, - kv_imported_tokens: read_i64(&mut reader)?, - kv_recorded_pages: read_i64(&mut reader)?, - kv_recorded_bytes: read_i64(&mut reader)?, - kv_hit_stage_mask: read_i64(&mut reader)?, - kv_record_stage_mask: read_i64(&mut reader)?, - checkpoint_flush_us: read_i64(&mut reader)?, - checkpoint_prefill_drain_us: read_i64(&mut reader)?, - checkpoint_local_us: read_i64(&mut reader)?, - checkpoint_downstream_write_us: read_i64(&mut reader)?, - checkpoint_downstream_wait_us: read_i64(&mut reader)?, - checkpoint_total_us: read_i64(&mut reader)?, - checkpoint_prefill_drained_replies: read_i64(&mut reader)?, - restore_flush_us: read_i64(&mut reader)?, - restore_prefill_drain_us: read_i64(&mut reader)?, - restore_local_us: read_i64(&mut reader)?, - restore_downstream_write_us: read_i64(&mut reader)?, - restore_downstream_wait_us: read_i64(&mut reader)?, - restore_total_us: read_i64(&mut reader)?, - restore_prefill_drained_replies: read_i64(&mut reader)?, - verify_span_compute_us: read_i64(&mut reader)?, - verify_span_forward_write_us: read_i64(&mut reader)?, - verify_span_downstream_wait_us: read_i64(&mut reader)?, - verify_span_total_us: read_i64(&mut reader)?, - verify_span_stage_count: read_i64(&mut reader)?, - verify_span_request_count: read_i64(&mut reader)?, - verify_span_token_count: read_i64(&mut reader)?, - verify_span_max_tokens: read_i64(&mut reader)?, - verify_span_checkpointed_requests: read_i64(&mut reader)?, - verify_span_skip_checkpoint_requests: read_i64(&mut reader)?, - prefill_edge_write_us_max: read_i64(&mut reader)?, - prefill_edge_wait_us_max: read_i64(&mut reader)?, - prefill_edge_total_us_max: read_i64(&mut reader)?, - prefill_edge_stage_index: read_i64(&mut reader)?, - prefill_edge_activation_bytes_max: read_i64(&mut reader)?, - prefill_edge_observation_count: read_i64(&mut reader)?, - }) + let mut bytes = [0_u8; REPLY_STATS_WIRE_BYTES]; + reader.read_exact(&mut bytes)?; + let mut fields = [0_i64; REPLY_STATS_FIELD_COUNT]; + for (field, chunk) in fields + .iter_mut() + .zip(bytes.chunks_exact(std::mem::size_of::())) + { + *field = i64::from_le_bytes(chunk.try_into().expect("i64 chunk size")); + } + Ok(reply_stats_from_fields(fields)) +} + +fn reply_stats_fields(stats: StageReplyStats) -> [i64; REPLY_STATS_FIELD_COUNT] { + [ + stats.kv_lookup_hits, + stats.kv_lookup_misses, + stats.kv_lookup_errors, + stats.kv_imported_pages, + stats.kv_imported_tokens, + stats.kv_recorded_pages, + stats.kv_recorded_bytes, + stats.kv_hit_stage_mask, + stats.kv_record_stage_mask, + stats.checkpoint_flush_us, + stats.checkpoint_prefill_drain_us, + stats.checkpoint_local_us, + stats.checkpoint_downstream_write_us, + stats.checkpoint_downstream_wait_us, + stats.checkpoint_total_us, + stats.checkpoint_prefill_drained_replies, + stats.restore_flush_us, + stats.restore_prefill_drain_us, + stats.restore_local_us, + stats.restore_downstream_write_us, + stats.restore_downstream_wait_us, + stats.restore_total_us, + stats.restore_prefill_drained_replies, + stats.verify_span_compute_us, + stats.verify_span_forward_write_us, + stats.verify_span_downstream_wait_us, + stats.verify_span_total_us, + stats.verify_span_stage_count, + stats.verify_span_request_count, + stats.verify_span_token_count, + stats.verify_span_max_tokens, + stats.verify_span_checkpointed_requests, + stats.verify_span_skip_checkpoint_requests, + stats.prefill_edge_write_us_max, + stats.prefill_edge_wait_us_max, + stats.prefill_edge_total_us_max, + stats.prefill_edge_stage_index, + stats.prefill_edge_activation_bytes_max, + stats.prefill_edge_observation_count, + ] +} + +fn reply_stats_from_fields(fields: [i64; REPLY_STATS_FIELD_COUNT]) -> StageReplyStats { + StageReplyStats { + kv_lookup_hits: fields[0], + kv_lookup_misses: fields[1], + kv_lookup_errors: fields[2], + kv_imported_pages: fields[3], + kv_imported_tokens: fields[4], + kv_recorded_pages: fields[5], + kv_recorded_bytes: fields[6], + kv_hit_stage_mask: fields[7], + kv_record_stage_mask: fields[8], + checkpoint_flush_us: fields[9], + checkpoint_prefill_drain_us: fields[10], + checkpoint_local_us: fields[11], + checkpoint_downstream_write_us: fields[12], + checkpoint_downstream_wait_us: fields[13], + checkpoint_total_us: fields[14], + checkpoint_prefill_drained_replies: fields[15], + restore_flush_us: fields[16], + restore_prefill_drain_us: fields[17], + restore_local_us: fields[18], + restore_downstream_write_us: fields[19], + restore_downstream_wait_us: fields[20], + restore_total_us: fields[21], + restore_prefill_drained_replies: fields[22], + verify_span_compute_us: fields[23], + verify_span_forward_write_us: fields[24], + verify_span_downstream_wait_us: fields[25], + verify_span_total_us: fields[26], + verify_span_stage_count: fields[27], + verify_span_request_count: fields[28], + verify_span_token_count: fields[29], + verify_span_max_tokens: fields[30], + verify_span_checkpointed_requests: fields[31], + verify_span_skip_checkpoint_requests: fields[32], + prefill_edge_write_us_max: fields[33], + prefill_edge_wait_us_max: fields[34], + prefill_edge_total_us_max: fields[35], + prefill_edge_stage_index: fields[36], + prefill_edge_activation_bytes_max: fields[37], + prefill_edge_observation_count: fields[38], + } } fn read_i32(mut reader: impl Read) -> io::Result { @@ -533,16 +581,6 @@ fn write_f32(mut writer: impl Write, value: f32) -> io::Result<()> { writer.write_all(&value.to_le_bytes()) } -fn read_i64(mut reader: impl Read) -> io::Result { - let mut bytes = [0_u8; 8]; - reader.read_exact(&mut bytes)?; - Ok(i64::from_le_bytes(bytes)) -} - -fn write_i64(mut writer: impl Write, value: i64) -> io::Result<()> { - writer.write_all(&value.to_le_bytes()) -} - fn read_u64(mut reader: impl Read) -> io::Result { let mut bytes = [0_u8; 8]; reader.read_exact(&mut bytes)?; diff --git a/crates/skippy-protocol/src/binary/mod.rs b/crates/skippy-protocol/src/binary/mod.rs index 501bb3924e..0eaac71cf8 100644 --- a/crates/skippy-protocol/src/binary/mod.rs +++ b/crates/skippy-protocol/src/binary/mod.rs @@ -10,7 +10,8 @@ pub use activation::{ pub use codec::{ read_stage_message, recv_ready, recv_reply, send_ready, send_reply_ack, send_reply_ack_with_stats, send_reply_predicted, send_reply_predicted_tokens_with_stats, - send_reply_predicted_with_stats, write_stage_message, + send_reply_predicted_with_stats, send_reply_predicted_with_tokens_and_stats, + write_stage_message, }; pub use types::{ ACTIVATION_FLAG_GEMMA3N_ALTUP, ACTIVATION_FLAG_RWKV7_V_FIRST, LLAMA_TOKEN_NULL, @@ -104,6 +105,22 @@ mod tests { assert_eq!(reply.predicted_tokens, vec![42]); } + #[test] + fn predicted_token_reply_preserves_sideband_tokens() { + let mut bytes = Vec::new(); + send_reply_predicted_with_tokens_and_stats( + &mut bytes, + 42, + &[42, 43, 123], + StageReplyStats::default(), + ) + .unwrap(); + let reply = recv_reply(Cursor::new(bytes)).unwrap(); + assert_eq!(reply.kind, WireReplyKind::PredictedToken); + assert_eq!(reply.predicted, 42); + assert_eq!(reply.predicted_tokens, vec![42, 43, 123]); + } + #[test] fn reply_stats_preserve_prefill_edge_transport() { let mut stats = StageReplyStats::default(); diff --git a/crates/skippy-quantize/Cargo.toml b/crates/skippy-quantize/Cargo.toml new file mode 100644 index 0000000000..9a6d2f1b23 --- /dev/null +++ b/crates/skippy-quantize/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "skippy-quantize" +edition.workspace = true +license.workspace = true +version.workspace = true +description = "Resumable GGUF conversion and quantization CLI for Skippy workflows" +repository = "https://github.com/Mesh-LLM/mesh-llm" +homepage = "https://github.com/Mesh-LLM/mesh-llm" +readme = "README.md" +publish = false + +[features] +default = [] +dynamic-llama-quant = ["llama-quant-ffi/dynamic-runtime"] + +[dependencies] +anyhow.workspace = true +clap.workspace = true +libc = "0.2" +llama-quant-ffi = { path = "../llama-quant-ffi" } +serde.workspace = true +serde_json.workspace = true +skippy-ffi = { path = "../skippy-ffi" } diff --git a/crates/skippy-quantize/README.md b/crates/skippy-quantize/README.md new file mode 100644 index 0000000000..4451c293a3 --- /dev/null +++ b/crates/skippy-quantize/README.md @@ -0,0 +1,539 @@ +# skippy-quantize + +`skippy-quantize` is the native Rust control plane for resumable GGUF +conversion and quantization jobs used by Skippy workflows. It replaces the old +Python converter and external `llama-quantize` process orchestration for this +pipeline; it does not install compatibility shims or shell out to those tools. + +The crate owns: + +- durable conversion and quantization manifests; +- split-GGUF progress detection and next-window planning; +- native SafeTensors-to-GGUF conversion for supported checkpoint families; +- in-process GGUF quantization through the linked llama quantization ABI; +- bounded source staging for quantization windows; +- optional output spooling with per-window publish and cleanup; +- successful-window records and JSON status/preflight output; +- tensor-type recipes for custom quant profiles such as `UD-Q3_K_S`; +- exact split artifact validation and optional llama load verification. + +Build through the repo recipes: + +```bash +just skippy-quantize-build +just skippy-quantize-release-build +just skippy-quantize-standalone-release-build +``` + +## Output + +Human-readable output is the default. It uses compact emoji-labelled status +lines and progress bars for shard progress. Commands that expose `--json` +emit structured JSON instead, including `status`, `next-window`, validation, +preflight, and the resumable `run-convert[-window]` / `run-quant[-window]` +commands. + +### CLI output examples + +Backend inspection defaults to human-readable capability summaries: + +```bash +skippy-quantize backends +``` + +```text +✅ native-rust conversion: available +ℹ️ native-rust: resumable_windows=true low_residency_streaming=true +✅ llama-api quantization: available +ℹ️ llama-api runtime: available +⚠️ skippy-abi runtime not loaded +ℹ️ skippy-abi: model_introspection=false gguf_slice_write=false feature_mask=unknown +``` + +Preflight shows the job shape, backend readiness, and source/target shard +progress: + +```bash +skippy-quantize quantize \ + --preflight-only \ + --backend llama-api \ + --tensor-type-file /mnt/recipe/glm-5.2-ud-q3-k-s.txt \ + /mnt/bf16/BF16/GLM-5.2-BF16-00001-of-00306.gguf \ + /mnt/quant/UD-Q3_K_S/GLM-5.2-UD-Q3_K_S.gguf \ + UD-Q3_K_S +``` + +```text +ℹ️ Preflight QuantizeGguf with backend llama-api +📊 target: [██░░░░░░░░░░░░░░░░░░░░░░] 28/306 shards (9.15%) +✅ Manifest is compatible +✅ Backend is ready +📊 source: [████████████████████████] 306/306 shards (100.00%) +✅ Source artifact is complete +ℹ️ Target missing ranges: 29..306 +``` + +`status` and `next-window` are concise for operators: + +```bash +skippy-quantize status --manifest /tmp/skippy-quantize.json +skippy-quantize next-window --manifest /tmp/skippy-quantize.json +``` + +```text +📊 job status: [██████░░░░░░░░░░░░░░░░░░] 77/306 shards (25.16%) +⚠️ Missing shards: 229 +ℹ️ Missing ranges: 78..306 +ℹ️ Next window: 78 + +ℹ️ Next window: 78 +``` + +Conversion and quantization window runners print the selected window and the +effective command in human mode. Add `--dry-run` to direct, job, `run-*`, or +`run-*-window` commands to plan the next missing window without writing output +artifacts, creating spool/output directories, staging source shards, recording +window records, publishing shards, or running completion verification. + +For long jobs, add `--json-event-file PATH` to write a compact periodic JSON +snapshot for agents to poll. The file is overwritten in place and keeps only a +bounded recent-event window, so agents do not need to ingest every log line: + +```bash +skippy-quantize run-quant \ + --manifest /tmp/skippy-quantize.json \ + --backend llama-api \ + --max-memory 32G \ + --json-event-file /tmp/skippy-quantize-status.json \ + --json-event-interval-seconds 120 \ + --json-event-window 8 +``` + +The snapshot has event type `skippy_quantize_periodic_status`, current phase, +current split window, timestamps, and the last N high-level events. + +```bash +skippy-quantize run-convert-window \ + --manifest /tmp/skippy-convert.json \ + --max-memory 32G \ + --spool-dir /tmp/skippy-convert-output +``` + +```text +🔒 Manifest lock acquired: /tmp/skippy-convert.json.lock +🪟 convert window: 42 +ℹ️ Output prefix: /tmp/skippy-convert-output/BF16/GLM-5.2-BF16.gguf +ℹ️ Command: skippy-quantize run-convert-window --backend native-rust --source /mnt/checkpoint --outfile /tmp/skippy-convert-output/BF16/GLM-5.2-BF16.gguf --first-split 42 --last-split 42 --expected-splits 306 +⚠️ convert memory budget: hard cap 32.00 GiB +ℹ️ Writing native convert shard 42/306 -> /tmp/skippy-convert-output/BF16/GLM-5.2-BF16-00042-of-00306.gguf (buffer 8.00 MiB, estimated working set 16.00 MiB) +✅ Published /mnt/target/BF16/GLM-5.2-BF16-00042-of-00306.gguf (49.87 GiB) +🔓 Manifest lock released: /tmp/skippy-convert.json.lock +``` + +Dry-run mode stops after the same plan and memory-budget output: + +```bash +skippy-quantize run-convert-window \ + --manifest /tmp/skippy-convert.json \ + --max-memory 32G \ + --spool-dir /tmp/skippy-convert-output \ + --dry-run +``` + +```text +🪟 convert window: 42 +ℹ️ Output prefix: /tmp/skippy-convert-output/BF16/GLM-5.2-BF16.gguf +ℹ️ Command: skippy-quantize run-convert-window --backend native-rust --source /mnt/checkpoint --outfile /tmp/skippy-convert-output/BF16/GLM-5.2-BF16.gguf --first-split 42 --last-split 42 --expected-splits 306 +⚠️ convert memory budget: hard cap 32.00 GiB +⚠️ convert dry run: no files were written, cleaned, recorded, or published +``` + +```bash +skippy-quantize run-quant-window \ + --manifest /tmp/skippy-quantize.json \ + --backend llama-api \ + --work-dir /tmp/skippy-quantize-work \ + --spool-dir /tmp/skippy-quantize-output +``` + +```text +🔒 Manifest lock acquired: /tmp/skippy-quantize.json.lock +📤 Copying /mnt/bf16/BF16/GLM-5.2-BF16-00001-of-00306.gguf -> /tmp/skippy-quantize-work/source-window/BF16/GLM-5.2-BF16-00001-of-00306.gguf (49.87 GiB) +ℹ️ Staged source window 1..306 at /tmp/skippy-quantize-work/source-window +🪟 quant window: 1..306 +ℹ️ Staged first shard: /tmp/skippy-quantize-work/source-window/BF16/GLM-5.2-BF16-00001-of-00306.gguf +ℹ️ Output prefix: /tmp/skippy-quantize-output/UD-Q3_K_S/GLM-5.2-UD-Q3_K_S.gguf +ℹ️ Command: llama-api-quantize --tensor-type-file /mnt/recipe/glm-5.2-ud-q3-k-s.txt --keep-split /tmp/skippy-quantize-work/source-window/BF16/GLM-5.2-BF16-00001-of-00306.gguf /tmp/skippy-quantize-output/UD-Q3_K_S/GLM-5.2-UD-Q3_K_S Q3_K_S +✅ Published /mnt/quant/UD-Q3_K_S/GLM-5.2-UD-Q3_K_S-00001-of-00306.gguf (13.42 GiB) +🧹 Cleaned staged source: /tmp/skippy-quantize-work/source-window +🔓 Manifest lock released: /tmp/skippy-quantize.json.lock +``` + +The llama API quantization backend now uses the unpatched llama.cpp quantizer. +It can preserve split output with `--keep-split`, but it cannot process only a +partial split window. Use a quant manifest window covering all expected splits +when selecting `--backend llama-api` or `--backend skippy-abi`. + +Validation commands also use the same progress-bar formatter: + +```bash +skippy-quantize validate-splits \ + --root /mnt/quant \ + --prefix UD-Q3_K_S \ + --basename GLM-5.2-UD-Q3_K_S \ + --expected-splits 306 +``` + +```text +📊 split artifact: [████████████████████████] 306/306 shards (100.00%) +✅ Split artifact is complete +``` + +Every workflow above can emit JSON for job automation: + +```bash +skippy-quantize run-quant-window \ + --manifest /tmp/skippy-quantize.json \ + --backend llama-api \ + --json +``` + +```json +{ + "event": "quant_window", + "plan": { + "first_split": 1, + "last_split": 306, + "staged_first_shard": "/tmp/skippy-quantize-work/source-window/BF16/GLM-5.2-BF16-00001-of-00306.gguf", + "output_prefix": "/mnt/quant/UD-Q3_K_S/GLM-5.2-UD-Q3_K_S.gguf", + "command": [ + "llama-api-quantize", + "--keep-split", + "/tmp/skippy-quantize-work/source-window/BF16/GLM-5.2-BF16-00001-of-00306.gguf", + "/mnt/quant/UD-Q3_K_S/GLM-5.2-UD-Q3_K_S", + "Q3_K_S" + ] + } +} +``` + +## Backends + +Inspect backend capabilities: + +```bash +skippy-quantize backends --json +``` + +`native-rust` is the HF checkpoint conversion backend. It streams tensor +payloads from SafeTensors into GGUF shards without materializing the whole model +or an output shard in memory. It currently requires tokenizer metadata from +`tokenizer.json`; checkpoints that only provide SentencePiece `tokenizer.model` +are rejected with a clear error until native SentencePiece support lands. + +`llama-api` and `skippy-abi` are quantization backends. The normal +`skippy-quantize` build links the pinned llama.cpp quantization ABI into the +binary, so `llama-api` can call `llama_model_quantize` in-process without a +separate `llama-quantize` executable or a dynamic library flag: + +```bash +skippy-quantize quantize \ + --backend llama-api \ + /mnt/source/BF16/model-00001-of-00002.gguf \ + /mnt/target/Q2_K/model-q2.gguf \ + Q2_K +``` + +`--native-runtime-library PATH` remains available for development builds that +intentionally load a dynamic llama.cpp runtime instead of using the linked ABI; +build that path with `--features dynamic-llama-quant`. Use +`--backend skippy-abi` when probing or loading the Skippy-patched runtime used +by mesh-llm. + +## Convert + +Create a conversion manifest: + +```bash +skippy-quantize init-convert \ + --source /mnt/checkpoint \ + --target /mnt/target \ + --target-prefix BF16 \ + --output-basename GLM-5.2-BF16 \ + --output-type bf16 \ + --expected-splits 306 \ + --window-size 1 \ + --manifest /tmp/skippy-convert.json +``` + +Run the next missing conversion window: + +```bash +skippy-quantize run-convert-window \ + --manifest /tmp/skippy-convert.json \ + --split-max-size 50G \ + --stream-buffer-bytes 8388608 \ + --spool-dir /tmp/skippy-convert-output \ + --record-dir /tmp/skippy-convert-records +``` + +Run conversion windows until complete: + +```bash +skippy-quantize run-convert \ + --manifest /tmp/skippy-convert.json \ + --max-memory 32G \ + --stream-buffer-bytes 8388608 \ + --spool-dir /tmp/skippy-convert-output \ + --record-dir /tmp/skippy-convert-records +``` + +For a direct native conversion command, pass the checkpoint and desired GGUF +output path. The command derives the target prefix, output basename, manifest +path, and then runs the same resumable loop: + +```bash +skippy-quantize convert \ + --output-type bf16 \ + --expected-splits 306 \ + --window-size 1 \ + --spool-dir /tmp/skippy-convert-output \ + /mnt/checkpoint \ + /mnt/target/BF16/GLM-5.2-BF16.gguf +``` + +Important conversion flags: + +- `--output-type {auto,bf16,f16,f32}` controls the emitted GGUF tensor type. +- `--expected-splits N` declares how many output shards the job should produce. +- `--window-size N` controls how many output shards each resumable run may + materialize. +- `--split-max-size SIZE` mirrors the intended split size in the native writer. +- `--stream-buffer-bytes BYTES` controls tensor streaming chunk size. +- `--max-memory SIZE` reduces native stream buffers and records the budget in + job logs. +- `--mtp` writes only appended MTP draft layers where supported. +- `--no-mtp` writes the trunk and drops appended MTP draft layers. +- `--spool-dir DIR` writes window outputs to a local spool before publishing. +- `--keep-spool` keeps the spooled window after publishing. +- `--record-dir DIR` writes per-window run records. +- `--print-only` prints the planned command/report for one window without + executing. +- `--dry-run` plans the next window without creating manifests, output/spool + directories, records, or artifacts. Loop commands plan only the next missing + window because no shard is written to advance progress. +- `--json-event-file PATH` writes a compact periodically refreshed status + snapshot for agent polling. +- `--json-event-interval-seconds N` controls the refresh period, default `120`. +- `--json-event-window N` controls how many recent high-level events are kept, + default `8`. + +## Quantize + +Create a quantization manifest from an existing split BF16/FP16 GGUF artifact: + +```bash +skippy-quantize init-quant \ + --source /mnt/bf16 \ + --source-prefix BF16 \ + --target /mnt/quant \ + --target-prefix UD-Q3_K_S \ + --output-basename GLM-5.2-UD-Q3_K_S \ + --quant UD-Q3_K_S \ + --tensor-type-file /mnt/recipe/tensor-types.txt \ + --window-size 1 \ + --manifest /tmp/skippy-quantize.json +``` + +Run one quantization window: + +```bash +skippy-quantize run-quant-window \ + --manifest /tmp/skippy-quantize.json \ + --backend llama-api \ + --work-dir /tmp/skippy-quantize-work \ + --spool-dir /tmp/skippy-quantize-output \ + --record-dir /tmp/skippy-quantize-records +``` + +Run until complete: + +```bash +skippy-quantize run-quant \ + --manifest /tmp/skippy-quantize.json \ + --backend llama-api \ + --max-memory 32G \ + --work-dir /tmp/skippy-quantize-work \ + --spool-dir /tmp/skippy-quantize-output +``` + +Important quantization flags: + +- `--backend {llama-api,skippy-abi}` selects the in-process quant backend. +- `--native-runtime-library PATH` optionally loads a dynamic native runtime + exposing `llama_model_quantize`; normal standalone builds do not need it. +- `--max-memory SIZE` applies to native Rust conversion memory planning. The + unpatched llama API quantization backend rejects it because llama.cpp does not + expose a quantization memory-budget knob. +- `--tensor-type-file PATH` applies per-tensor recipe overrides. +- `--tensor-type NAME=TYPE` adds an inline per-tensor override. +- `--imatrix PATH` loads legacy `.dat` or GGUF imatrix data. +- `--include-weights PATTERN` and `--exclude-weights PATTERN` filter imatrix + weights. +- `--output-tensor-type TYPE` and `--token-embedding-type TYPE` override key + tensor types. +- `--prune-layers SPEC` forwards layer-pruning metadata to the native quant + API. +- `--override-kv KEY=TYPE:VALUE` adds GGUF metadata overrides. +- `--allow-requantize`, `--pure`, and `--leave-output-tensor` mirror native + quantization parameters. +- `--dry-run` plans the next quant window without creating manifests, output or + spool directories, staging source shards, records, or artifacts. Loop commands + plan only the next missing window because no shard is written to advance + progress. +- `--json-event-file PATH` writes a compact periodically refreshed status + snapshot for agent polling. +- `--json-event-interval-seconds N` controls the refresh period, default `120`. +- `--json-event-window N` controls how many recent high-level events are kept, + default `8`. +- `--keep-split`, `--first-split`, and `--last-split` can request a manual + split window for direct `quantize`; llama API quantization accepts only the + full split range after the mesh-llm llama-quantize split-window patches were + removed. +- `--no-stage-source` skips local source-window staging. +- `--keep-staged-source` keeps the staged source window after success. + +## Recipes + +Top-level quantization modes intentionally mirror the pinned llama.cpp quant +table. Custom profile names such as `Q2_K-MTP-Q8`, `UD-Q3_K_S`, or `Q4_K_XL` +belong in artifact names such as `--target-prefix` and `--output-basename`, not +in `--quant`. Pass the base llama quant with `--quant` and express any +per-tensor policy with `--tensor-type-file` or repeated `--tensor-type`. + +The tensor recipe format is one override per line: + +```text +blk.*.ffn_gate_exps.weight=Q2_K +blk.*.ffn_down_exps.weight=Q3_K +mtp.*=Q8_0 +``` + +Inspect supported modes and raw tensor override types: + +```bash +skippy-quantize list-quants --json +skippy-quantize list-tensor-types --json +``` + +## BF16 to layer package + +For lab workflows, keep one reusable split BF16 GGUF artifact as the durable +source of truth, then build quantized layer packages from it. The quantized +GGUF shards can be treated as disposable staging once the package preflight +passes: + +```bash +skippy-quantize quantize-layer-package \ + --source /Users/lab/glm52-work/bf16-gguf \ + --source-prefix BF16 \ + --target /Users/lab/glm52-work/quantized \ + --target-prefix Q2_K-MTP-Q8 \ + --manifest /Users/lab/glm52-work/work/q2-k-mtp-q8-package/quant-manifest.json \ + --package-dir /Users/lab/glm52-work/packages/GLM-5.2-Q2_K-MTP-Q8-layers \ + --package-model-id meshllm/GLM-5.2-Q2_K-MTP-Q8-GGUF:Q2_K-MTP-Q8 \ + --package-source-repo meshllm/GLM-5.2-Q2_K-MTP-Q8-GGUF \ + --package-source-revision local \ + --work-dir /Users/lab/glm52-work/work/q2-k-mtp-q8-package/native-work \ + --spool-dir /Users/lab/glm52-work/work/q2-k-mtp-q8-package/spool \ + --record-dir /Users/lab/glm52-work/work/q2-k-mtp-q8-package/records \ + --json-event-file /Users/lab/glm52-work/work/q2-k-mtp-q8-package/status.json \ + --quant Q2_K \ + --tensor-type-file /Users/lab/glm52-work/recipes/glm-5.2-q2-k-mtp-q8.tensor-types.txt \ + --output-basename GLM-5.2-Q2_K-MTP-Q8 \ + --stages 2 \ + --replace-package \ + --watchdog-seconds 120 +``` + +Build prerequisites: + +```bash +just skippy-quantize-standalone-release-build +cargo build --release --locked -p skippy-model-package +``` + +The command validates the source split, writes the package artifacts from the +BF16 GGUF source, quantizes each artifact in place, then runs package preflight. +It does not materialize a complete quantized GGUF repo first. By default, the +temporary quant scratch directory is deleted after package preflight passes; +pass `--keep-quant` to retain it. It does not pass `--max-memory` to +quantization because the unpatched llama API quant backend does not expose a +memory-budget knob. + +## Validation + +Useful checks: + +```bash +skippy-quantize status --manifest /tmp/skippy-quantize.json --json +skippy-quantize next-window --manifest /tmp/skippy-quantize.json --json +skippy-quantize verify-job --manifest /tmp/skippy-quantize.json --llama-load +skippy-quantize validate-tensor-types /mnt/recipe/tensor-types.txt +skippy-quantize validate-splits --root /mnt/target --prefix UD-Q3_K_S --json +``` + +### Reference parity smoke + +Use `scripts/compare-reference-quantization.py` when changing native +conversion or quantization behavior. It compares the native Rust path against +the pinned llama.cpp reference tools: + +- SafeTensors conversion: upstream `convert_hf_to_gguf.py` and + `skippy-quantize convert` must emit the same tensor name set, shapes, types, + and tensor payload bytes. Whole-file GGUF byte equality is not required here + because the two writers may emit metadata and tensors in different order. +- Quantization: standalone `llama-quantize --keep-split` and + `skippy-quantize quantize --backend llama-api` must emit byte-identical split + GGUF outputs for every mode reported by `skippy-quantize list-quants --json`. + +Conversion-only smoke: + +```bash +uv run --python 3.12 \ + --with torch \ + --with transformers \ + --with numpy \ + --with sentencepiece \ + --with protobuf \ + --with gguf \ + --no-project \ + crates/skippy-quantize/scripts/compare-reference-quantization.py \ + --work-dir /tmp/skippy-quantize-conversion-parity \ + --clean \ + --skippy-quantize ./target/debug/skippy-quantize \ + --llama-quantize ./.deps/llama.cpp/build-cli/bin/llama-quantize \ + --python-converter ./.deps/llama.cpp/convert_hf_to_gguf.py \ + --checkpoint /tmp/qwen2-safetensors-fixture \ + --skip-quantization +``` + +All advertised quant modes: + +```bash +uv run --python 3.12 \ + --with gguf \ + --with numpy \ + --no-project \ + crates/skippy-quantize/scripts/compare-reference-quantization.py \ + --work-dir /tmp/skippy-quantize-allmodes \ + --clean \ + --skippy-quantize ./target/debug/skippy-quantize \ + --llama-quantize ./.deps/llama.cpp/build-cli/bin/llama-quantize \ + --quant-input /tmp/qwen2-bf16-fixture.gguf \ + --generate-imatrix +``` + +`--generate-imatrix` creates a deterministic all-ones legacy imatrix from the +GGUF tensor metadata so very low-bit and IQ modes are tested instead of being +accepted as matching failures. diff --git a/crates/skippy-quantize/scripts/compare-reference-quantization.py b/crates/skippy-quantize/scripts/compare-reference-quantization.py new file mode 100755 index 0000000000..897b880ba8 --- /dev/null +++ b/crates/skippy-quantize/scripts/compare-reference-quantization.py @@ -0,0 +1,465 @@ +#!/usr/bin/env python3 +"""Compare native skippy-quantize output with llama.cpp reference tools. + +The script has two independent checks: + +* conversion: convert a SafeTensors checkpoint with upstream + convert_hf_to_gguf.py and with `skippy-quantize convert`, then compare GGUF + tensor names, shapes, types, and payload bytes. +* quantization: quantize a BF16/FP16 GGUF with standalone `llama-quantize` and + with `skippy-quantize quantize --backend llama-api` for every requested quant + mode, then require byte-identical split outputs. + +It intentionally records failures instead of stopping at the first unsupported +mode so a full quant catalog run produces actionable evidence. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import shutil +import struct +import subprocess +import sys +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + + +@dataclass +class CommandResult: + argv: list[str] + returncode: int + log: str + + +@dataclass +class ConversionResult: + status: str + python_gguf: str | None + skippy_gguf: str | None + tensor_order_equal: bool | None + tensor_name_set_equal: bool | None + tensor_mismatch_count: int | None + byte_equal: bool | None + error: str | None + + +@dataclass +class QuantResult: + quant: str + status: str + reference_output: str | None + skippy_output: str | None + reference_sha256: str | None + skippy_sha256: str | None + byte_equal: bool | None + reference_returncode: int | None + skippy_returncode: int | None + error: str | None + + +def main() -> int: + args = parse_args() + work_dir = args.work_dir.resolve() + if args.clean and work_dir.exists(): + shutil.rmtree(work_dir) + work_dir.mkdir(parents=True, exist_ok=True) + + report: dict[str, Any] = { + "conversion": None, + "quantization": [], + } + + if args.checkpoint is not None: + report["conversion"] = asdict(run_conversion(args, work_dir)) + + quant_input = args.quant_input + if quant_input is None and report["conversion"] is not None: + conversion = report["conversion"] + if conversion["status"] == "ok": + quant_input = Path(conversion["skippy_gguf"]) + + if quant_input is not None and not args.skip_quantization: + for quant in selected_quants(args): + result = run_quant(args, work_dir, quant, quant_input.resolve()) + report["quantization"].append(asdict(result)) + + report_path = work_dir / "comparison-report.json" + report_path.write_text(json.dumps(report, indent=2), encoding="utf-8") + print(json.dumps(report, indent=2)) + print(f"comparison_report={report_path}") + + return 0 if report_passed(report, args.allow_matching_failures) else 1 + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", type=Path) + parser.add_argument("--quant-input", type=Path) + parser.add_argument("--work-dir", type=Path, required=True) + parser.add_argument("--skippy-quantize", type=Path, required=True) + parser.add_argument("--llama-quantize", type=Path, required=True) + parser.add_argument("--python-converter", type=Path) + parser.add_argument("--python", default=sys.executable) + parser.add_argument("--output-type", default="bf16") + parser.add_argument("--nthreads", default="8") + parser.add_argument("--quant", action="append") + parser.add_argument("--skip-quant", action="append", default=[]) + parser.add_argument("--skip-quantization", action="store_true") + parser.add_argument("--imatrix", type=Path) + parser.add_argument( + "--generate-imatrix", + action="store_true", + help="Generate a deterministic all-ones legacy imatrix from --quant-input.", + ) + parser.add_argument("--native-runtime-library", action="append", type=Path, default=[]) + parser.add_argument("--clean", action="store_true") + parser.add_argument( + "--allow-matching-failures", + action="store_true", + help="Treat modes where both reference and native fail as non-fatal.", + ) + return parser.parse_args() + + +def run_conversion(args: argparse.Namespace, work_dir: Path) -> ConversionResult: + if args.python_converter is None: + return ConversionResult( + status="skipped", + python_gguf=None, + skippy_gguf=None, + tensor_order_equal=None, + tensor_name_set_equal=None, + tensor_mismatch_count=None, + byte_equal=None, + error="--python-converter not provided", + ) + + conversion_dir = work_dir / "conversion" + python_dir = conversion_dir / "python" + skippy_dir = conversion_dir / "skippy" + python_dir.mkdir(parents=True, exist_ok=True) + skippy_dir.mkdir(parents=True, exist_ok=True) + + basename = args.checkpoint.name + python_gguf = python_dir / f"{basename}-{args.output_type}.gguf" + skippy_gguf = skippy_dir / f"{basename}-{args.output_type}.gguf" + + python_cmd = [ + args.python, + str(args.python_converter), + "--outtype", + args.output_type, + "--outfile", + str(python_gguf), + str(args.checkpoint), + ] + result = run_logged(python_cmd, conversion_dir / "python-convert.log") + if result.returncode != 0: + return conversion_error("python_conversion_failed", python_gguf, skippy_gguf, result) + + skippy_cmd = [ + str(args.skippy_quantize), + "convert", + "--output-type", + args.output_type, + "--expected-splits", + "1", + "--no-verify-on-complete", + str(args.checkpoint), + str(skippy_gguf), + ] + result = run_logged(skippy_cmd, conversion_dir / "skippy-convert.log") + if result.returncode != 0: + return conversion_error("skippy_conversion_failed", python_gguf, skippy_gguf, result) + + comparison = compare_gguf_tensors(python_gguf, skippy_gguf) + return ConversionResult( + status="ok" if comparison["tensor_mismatch_count"] == 0 else "tensor_mismatch", + python_gguf=str(python_gguf), + skippy_gguf=str(skippy_gguf), + tensor_order_equal=comparison["tensor_order_equal"], + tensor_name_set_equal=comparison["tensor_name_set_equal"], + tensor_mismatch_count=comparison["tensor_mismatch_count"], + byte_equal=same_file_bytes(python_gguf, skippy_gguf), + error=None, + ) + + +def conversion_error( + status: str, python_gguf: Path, skippy_gguf: Path, result: CommandResult +) -> ConversionResult: + return ConversionResult( + status=status, + python_gguf=str(python_gguf), + skippy_gguf=str(skippy_gguf), + tensor_order_equal=None, + tensor_name_set_equal=None, + tensor_mismatch_count=None, + byte_equal=None, + error=f"returncode={result.returncode} log={result.log}", + ) + + +def selected_quants(args: argparse.Namespace) -> list[str]: + if args.quant: + names = args.quant + else: + proc = subprocess.run( + [str(args.skippy_quantize), "list-quants", "--json"], + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=True, + ) + names = json.loads(proc.stdout)["whole_model_quant_modes"] + skipped = set(args.skip_quant) + return [name for name in names if name not in skipped] + + +def run_quant( + args: argparse.Namespace, work_dir: Path, quant: str, quant_input: Path +) -> QuantResult: + quant_dir = work_dir / "quant" / safe_name(quant) + ref_dir = quant_dir / "reference" + skippy_dir = quant_dir / "skippy" + ref_dir.mkdir(parents=True, exist_ok=True) + skippy_dir.mkdir(parents=True, exist_ok=True) + + reference_prefix = ref_dir / f"model-{safe_name(quant)}.gguf" + skippy_prefix = skippy_dir / f"model-{safe_name(quant)}.gguf" + imatrix = imatrix_path(args, work_dir, quant_input) + reference_cmd = [ + str(args.llama_quantize), + "--keep-split", + "--first-split", + "1", + "--last-split", + "1", + ] + if imatrix is not None: + reference_cmd.extend(["--imatrix", str(imatrix)]) + reference_cmd.extend([ + str(quant_input), + str(reference_prefix), + quant, + str(args.nthreads), + ]) + reference = run_logged(reference_cmd, quant_dir / "reference.log") + + skippy_cmd = [ + str(args.skippy_quantize), + "quantize", + "--backend", + "llama-api", + "--no-stage-source", + "--no-verify-on-complete", + "--work-dir", + str(quant_dir / "work"), + ] + for library in args.native_runtime_library: + skippy_cmd.extend(["--native-runtime-library", str(library)]) + if imatrix is not None: + skippy_cmd.extend(["--imatrix", str(imatrix)]) + skippy_cmd.extend([str(quant_input), str(skippy_prefix), quant, str(args.nthreads)]) + skippy = run_logged(skippy_cmd, quant_dir / "skippy.log") + + ref_output = split_output_path(reference_prefix) + skippy_output = split_output_path(strip_gguf_suffix(skippy_prefix)) + + if reference.returncode != 0 or skippy.returncode != 0: + status = "matching_failure" if reference.returncode == skippy.returncode else "failure" + return QuantResult( + quant=quant, + status=status, + reference_output=str(ref_output) if ref_output.exists() else None, + skippy_output=str(skippy_output) if skippy_output.exists() else None, + reference_sha256=None, + skippy_sha256=None, + byte_equal=None, + reference_returncode=reference.returncode, + skippy_returncode=skippy.returncode, + error=f"reference_log={reference.log} skippy_log={skippy.log}", + ) + + if not ref_output.exists() or not skippy_output.exists(): + return QuantResult( + quant=quant, + status="missing_output", + reference_output=str(ref_output), + skippy_output=str(skippy_output), + reference_sha256=None, + skippy_sha256=None, + byte_equal=None, + reference_returncode=reference.returncode, + skippy_returncode=skippy.returncode, + error="expected split output missing", + ) + + ref_sha = sha256(ref_output) + skippy_sha = sha256(skippy_output) + byte_equal = ref_sha == skippy_sha and same_file_bytes(ref_output, skippy_output) + return QuantResult( + quant=quant, + status="ok" if byte_equal else "byte_mismatch", + reference_output=str(ref_output), + skippy_output=str(skippy_output), + reference_sha256=ref_sha, + skippy_sha256=skippy_sha, + byte_equal=byte_equal, + reference_returncode=reference.returncode, + skippy_returncode=skippy.returncode, + error=None, + ) + + +def run_logged(argv: list[str], log_path: Path) -> CommandResult: + log_path.parent.mkdir(parents=True, exist_ok=True) + with log_path.open("w", encoding="utf-8") as log: + proc = subprocess.run( + argv, + text=True, + stdout=log, + stderr=subprocess.STDOUT, + env=os.environ.copy(), + ) + return CommandResult(argv=argv, returncode=proc.returncode, log=str(log_path)) + + +def compare_gguf_tensors(left: Path, right: Path) -> dict[str, Any]: + try: + import numpy as np + from gguf import GGUFReader + except Exception as exc: # pragma: no cover - dependency error path + raise SystemExit(f"install gguf and numpy for tensor comparison: {exc}") from exc + + left_reader = GGUFReader(left) + right_reader = GGUFReader(right) + left_tensors = {tensor.name: tensor for tensor in left_reader.tensors} + right_tensors = {tensor.name: tensor for tensor in right_reader.tensors} + tensor_order_equal = list(left_tensors) == list(right_tensors) + tensor_name_set_equal = set(left_tensors) == set(right_tensors) + mismatches = [] + for name, left_tensor in left_tensors.items(): + right_tensor = right_tensors.get(name) + if right_tensor is None: + mismatches.append(name) + continue + same_shape = left_tensor.shape.tolist() == right_tensor.shape.tolist() + same_type = left_tensor.tensor_type == right_tensor.tensor_type + same_data = np.array_equal(left_tensor.data, right_tensor.data) + if not (same_shape and same_type and same_data): + mismatches.append(name) + for name in right_tensors: + if name not in left_tensors: + mismatches.append(name) + return { + "tensor_order_equal": tensor_order_equal, + "tensor_name_set_equal": tensor_name_set_equal, + "tensor_mismatch_count": len(mismatches), + "first_mismatches": mismatches[:20], + } + + +def report_passed(report: dict[str, Any], allow_matching_failures: bool) -> bool: + conversion = report.get("conversion") + if conversion is not None and conversion["status"] not in ("ok", "skipped"): + return False + for result in report["quantization"]: + if result["status"] == "ok": + continue + if allow_matching_failures and result["status"] == "matching_failure": + continue + return False + return True + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def same_file_bytes(left: Path, right: Path) -> bool: + if left.stat().st_size != right.stat().st_size: + return False + with left.open("rb") as left_handle, right.open("rb") as right_handle: + while True: + left_chunk = left_handle.read(1024 * 1024) + right_chunk = right_handle.read(1024 * 1024) + if left_chunk != right_chunk: + return False + if not left_chunk: + return True + + +def imatrix_path(args: argparse.Namespace, work_dir: Path, quant_input: Path) -> Path | None: + if args.imatrix is not None: + return args.imatrix.resolve() + if not args.generate_imatrix: + return None + path = work_dir / "generated-imatrix.dat" + if not path.exists(): + write_legacy_all_ones_imatrix(quant_input, path) + return path + + +def write_legacy_all_ones_imatrix(quant_input: Path, output: Path) -> None: + try: + from gguf import GGUFReader + except Exception as exc: # pragma: no cover - dependency error path + raise SystemExit(f"install gguf to generate imatrix fixture: {exc}") from exc + + reader = GGUFReader(quant_input) + entries: list[tuple[str, int]] = [] + for tensor in reader.tensors: + shape = tensor.shape.tolist() + if len(shape) < 2: + continue + imatrix_width = int(shape[0]) + if len(shape) >= 3: + imatrix_width *= int(shape[2]) + entries.append((tensor.name, imatrix_width)) + + if not entries: + raise SystemExit(f"no rank >= 2 tensors found for imatrix in {quant_input}") + + output.parent.mkdir(parents=True, exist_ok=True) + with output.open("wb") as handle: + handle.write(struct.pack(" Path: + prefix = strip_gguf_suffix(prefix_or_path) + return prefix.with_name(f"{prefix.name}-00001-of-00001.gguf") + + +def strip_gguf_suffix(path: Path) -> Path: + if path.suffix == ".gguf": + return path.with_suffix("") + return path + + +def safe_name(name: str) -> str: + return name.replace("/", "_").replace(" ", "_").replace("-", "_") + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/crates/skippy-quantize/src/artifacts.rs b/crates/skippy-quantize/src/artifacts.rs new file mode 100644 index 0000000000..213ed69991 --- /dev/null +++ b/crates/skippy-quantize/src/artifacts.rs @@ -0,0 +1,277 @@ +use std::fs; +use std::path::{Path, PathBuf}; + +use anyhow::{Context, Result, ensure}; + +use crate::output::{format_bytes, print_copy, print_info, print_path_event, print_success}; +use crate::residency::copy_file_bounded_with_label; +use crate::splits::SplitWindow; + +pub fn execution_root(target: &Path, target_prefix: &str, spool_dir: Option<&Path>) -> PathBuf { + spool_dir.unwrap_or(target).join(target_prefix) +} + +pub fn publish_spooled_window( + spool_dir: Option<&Path>, + target: &Path, + target_prefix: &str, + output_basename: &str, + expected_splits: u32, + window: SplitWindow, + keep_spool: bool, +) -> Result<()> { + let Some(spool_dir) = spool_dir else { + return Ok(()); + }; + + let source_root = spool_dir.join(target_prefix); + let target_root = target.join(target_prefix); + fs::create_dir_all(&target_root) + .with_context(|| format!("create {}", target_root.display()))?; + + for index in window.first_split..=window.last_split { + publish_one_shard( + &source_root, + &target_root, + output_basename, + index, + expected_splits, + keep_spool, + )?; + } + Ok(()) +} + +pub fn clean_spooled_window( + spool_dir: Option<&Path>, + target_prefix: &str, + output_basename: &str, + expected_splits: u32, + window: SplitWindow, +) -> Result<()> { + let Some(spool_dir) = spool_dir else { + return Ok(()); + }; + let source_root = spool_dir.join(target_prefix); + for index in window.first_split..=window.last_split { + for path in stale_spool_candidates(&source_root, output_basename, index, expected_splits) { + if path.exists() { + fs::remove_file(&path).with_context(|| format!("remove {}", path.display()))?; + print_path_event("🧹", "Removed stale spool shard", &path); + } + } + } + Ok(()) +} + +fn publish_one_shard( + source_root: &Path, + target_root: &Path, + output_basename: &str, + index: u32, + expected_splits: u32, + keep_spool: bool, +) -> Result<()> { + let source = source_candidate(source_root, output_basename, index, expected_splits); + let target = target_root.join( + source + .file_name() + .context("spooled source shard has invalid file name")?, + ); + ensure!( + source.is_file(), + "spooled output shard does not exist: {}", + source.display() + ); + publish_file(&source, &target)?; + if !keep_spool { + fs::remove_file(&source).with_context(|| format!("remove {}", source.display()))?; + print_path_event("🧹", "Removed spooled shard", &source); + } + Ok(()) +} + +fn publish_file(source: &Path, target: &Path) -> Result<()> { + let source_len = source + .metadata() + .with_context(|| format!("stat {}", source.display()))? + .len(); + if target.exists() { + let target_len = target + .metadata() + .with_context(|| format!("stat {}", target.display()))? + .len(); + ensure!( + source_len == target_len, + "target shard already exists with different size: {} source_bytes={} target_bytes={}", + target.display(), + source_len, + target_len + ); + print_info(format!( + "Publish target already exists: {} ({})", + target.display(), + format_bytes(source_len) + )); + return Ok(()); + } + + let temp_name = format!( + "{}.part", + target + .file_name() + .and_then(|name| name.to_str()) + .context("target shard has invalid file name")? + ); + let temp = target.with_file_name(temp_name); + if temp.exists() { + fs::remove_file(&temp).with_context(|| format!("remove {}", temp.display()))?; + } + print_copy(source, target, Some(source_len)); + copy_file_bounded_with_label("publish_copy", source, &temp)?; + fs::rename(&temp, target) + .with_context(|| format!("rename {} -> {}", temp.display(), target.display()))?; + print_success(format!( + "Published {} ({})", + target.display(), + format_bytes(source_len) + )); + Ok(()) +} + +fn split_shard_name(output_basename: &str, index: u32, expected_splits: u32) -> String { + format!("{output_basename}-{index:05}-of-{expected_splits:05}.gguf") +} + +fn source_candidate( + source_root: &Path, + output_basename: &str, + index: u32, + expected_splits: u32, +) -> PathBuf { + let split = source_root.join(split_shard_name(output_basename, index, expected_splits)); + if split.is_file() || expected_splits != 1 || index != 1 { + return split; + } + let unsplit = source_root.join(format!("{output_basename}.gguf")); + if unsplit.is_file() { + return unsplit; + } + split +} + +fn stale_spool_candidates( + source_root: &Path, + output_basename: &str, + index: u32, + expected_splits: u32, +) -> Vec { + let mut candidates = + vec![source_root.join(split_shard_name(output_basename, index, expected_splits))]; + if expected_splits == 1 && index == 1 { + candidates.push(source_root.join(format!("{output_basename}.gguf"))); + } + candidates +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::records::unix_timestamp_ms; + + #[test] + fn publishes_spooled_window_and_cleans_source() { + let root = std::env::temp_dir().join(format!( + "skippy-quantize-publish-test-{}", + unix_timestamp_ms() + )); + let spool = root.join("spool"); + let target = root.join("target"); + let spool_root = spool.join("Q2_K"); + fs::create_dir_all(&spool_root).unwrap(); + fs::write(spool_root.join("model-00002-of-00003.gguf"), b"two").unwrap(); + + publish_spooled_window( + Some(&spool), + &target, + "Q2_K", + "model", + 3, + SplitWindow { + first_split: 2, + last_split: 2, + }, + false, + ) + .unwrap(); + + assert_eq!( + fs::read(target.join("Q2_K/model-00002-of-00003.gguf")).unwrap(), + b"two" + ); + assert!(!spool_root.join("model-00002-of-00003.gguf").exists()); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn cleans_only_current_spooled_window() { + let root = std::env::temp_dir().join(format!( + "skippy-quantize-spool-clean-test-{}", + unix_timestamp_ms() + )); + let spool = root.join("spool"); + let spool_root = spool.join("Q2_K"); + fs::create_dir_all(&spool_root).unwrap(); + fs::write(spool_root.join("model-00001-of-00003.gguf"), b"one").unwrap(); + fs::write(spool_root.join("model-00002-of-00003.gguf"), b"two").unwrap(); + fs::write(spool_root.join("model-00003-of-00003.gguf"), b"three").unwrap(); + + clean_spooled_window( + Some(&spool), + "Q2_K", + "model", + 3, + SplitWindow { + first_split: 2, + last_split: 2, + }, + ) + .unwrap(); + + assert!(spool_root.join("model-00001-of-00003.gguf").exists()); + assert!(!spool_root.join("model-00002-of-00003.gguf").exists()); + assert!(spool_root.join("model-00003-of-00003.gguf").exists()); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn publishes_unsplit_single_file_output() { + let root = std::env::temp_dir().join(format!( + "skippy-quantize-unsplit-publish-test-{}", + unix_timestamp_ms() + )); + let spool = root.join("spool"); + let target = root.join("target"); + let spool_root = spool.join("Q2_K"); + fs::create_dir_all(&spool_root).unwrap(); + fs::write(spool_root.join("model.gguf"), b"one").unwrap(); + + publish_spooled_window( + Some(&spool), + &target, + "Q2_K", + "model", + 1, + SplitWindow { + first_split: 1, + last_split: 1, + }, + false, + ) + .unwrap(); + + assert_eq!(fs::read(target.join("Q2_K/model.gguf")).unwrap(), b"one"); + assert!(!spool_root.join("model.gguf").exists()); + fs::remove_dir_all(root).unwrap(); + } +} diff --git a/crates/skippy-quantize/src/backend.rs b/crates/skippy-quantize/src/backend.rs new file mode 100644 index 0000000000..6907a8eb6a --- /dev/null +++ b/crates/skippy-quantize/src/backend.rs @@ -0,0 +1,292 @@ +use std::path::PathBuf; + +use anyhow::{Result, ensure}; +use clap::{Parser, ValueEnum}; +use serde::{Deserialize, Serialize}; + +use crate::output::{print_info, print_json_pretty, print_success, print_warn}; + +const SKIPPY_FEATURE_MODEL_INTROSPECTION: u64 = 1 << 3; +const SKIPPY_FEATURE_GGUF_SLICE_WRITE: u64 = 1 << 4; + +#[derive(Debug, Parser)] +pub struct BackendArgs { + #[arg(long = "skippy-runtime-library", value_name = "PATH")] + skippy_runtime_libraries: Vec, + #[arg(long)] + json: bool, +} + +#[derive(Debug, Serialize)] +pub struct BackendCapabilities { + pub native_rust: NativeRustCapabilities, + pub llama_api: LlamaApiCapabilities, + pub skippy_abi: SkippyAbiCapabilities, +} + +#[derive(Debug, Serialize)] +pub struct NativeRustCapabilities { + pub convert_hf_to_gguf: bool, + pub llama_quantize: bool, + pub resumable_windows: bool, + pub low_residency_streaming: bool, + pub reason: String, +} + +#[derive(Debug, Serialize)] +pub struct LlamaApiCapabilities { + pub convert_hf_to_gguf: bool, + pub llama_quantize: bool, + pub runtime_loaded: bool, + pub reason: String, +} + +#[derive(Debug, Serialize)] +pub struct SkippyAbiCapabilities { + pub convert_hf_to_gguf: bool, + pub llama_quantize: bool, + pub runtime_loaded: bool, + pub feature_mask: Option, + pub model_introspection: bool, + pub gguf_slice_write: bool, + pub load_error: Option, + pub reason: String, +} + +pub fn capabilities(skippy_runtime_libraries: &[PathBuf]) -> BackendCapabilities { + let skippy_abi = skippy_abi_capabilities(skippy_runtime_libraries); + let llama_api = LlamaApiCapabilities { + convert_hf_to_gguf: false, + llama_quantize: llama_quant_ffi::native_runtime_loaded(), + runtime_loaded: llama_quant_ffi::native_runtime_loaded(), + reason: if llama_quant_ffi::native_runtime_loaded() { + "linked or loaded llama quant runtime exposes llama_model_quantize".to_string() + } else { + "no linked llama quant runtime and no native runtime library was loaded for llama API probing".to_string() + }, + }; + BackendCapabilities { + native_rust: NativeRustCapabilities { + convert_hf_to_gguf: true, + llama_quantize: false, + resumable_windows: true, + low_residency_streaming: true, + reason: "Rust SafeTensors-to-GGUF writer streams tensor payloads and materializes one split window per run".to_string(), + }, + llama_api, + skippy_abi, + } +} + +pub fn run_backends(args: BackendArgs) -> Result<()> { + let capabilities = capabilities(&args.skippy_runtime_libraries); + if args.json { + print_json_pretty(&capabilities)?; + } else { + print_success(format!( + "native-rust conversion: {}", + bool_word(capabilities.native_rust.convert_hf_to_gguf) + )); + print_info(format!( + "native-rust: resumable_windows={} low_residency_streaming={}", + capabilities.native_rust.resumable_windows, + capabilities.native_rust.low_residency_streaming + )); + print_success(format!( + "llama-api quantization: {}", + bool_word(capabilities.llama_api.llama_quantize) + )); + print_info(format!( + "llama-api runtime: {}", + bool_word(capabilities.llama_api.runtime_loaded) + )); + if capabilities.skippy_abi.runtime_loaded { + print_success("skippy-abi runtime loaded"); + } else { + print_warn("skippy-abi runtime not loaded"); + } + print_info(format!( + "skippy-abi: model_introspection={} gguf_slice_write={} feature_mask={}", + capabilities.skippy_abi.model_introspection, + capabilities.skippy_abi.gguf_slice_write, + capabilities + .skippy_abi + .feature_mask + .map_or_else(|| "unknown".to_string(), |value| format!("{value:#x}")) + )); + if let Some(load_error) = capabilities.skippy_abi.load_error.as_deref() { + print_warn(format!("skippy-abi load error: {load_error}")); + } + } + Ok(()) +} + +fn bool_word(value: bool) -> &'static str { + if value { "available" } else { "unavailable" } +} + +fn skippy_abi_capabilities(skippy_runtime_libraries: &[PathBuf]) -> SkippyAbiCapabilities { + let load_error = load_skippy_runtime_for_probe(skippy_runtime_libraries); + let runtime_loaded = skippy_ffi::native_runtime_loaded(); + let feature_mask = if runtime_loaded { + std::panic::catch_unwind(skippy_ffi::skippy_abi_features).ok() + } else { + None + }; + let model_introspection = feature_mask.is_some_and(|mask| { + mask & SKIPPY_FEATURE_MODEL_INTROSPECTION == SKIPPY_FEATURE_MODEL_INTROSPECTION + }); + let gguf_slice_write = feature_mask.is_some_and(|mask| { + mask & SKIPPY_FEATURE_GGUF_SLICE_WRITE == SKIPPY_FEATURE_GGUF_SLICE_WRITE + }); + SkippyAbiCapabilities { + convert_hf_to_gguf: false, + llama_quantize: runtime_loaded, + runtime_loaded, + feature_mask, + model_introspection, + gguf_slice_write, + load_error, + reason: skippy_abi_reason(runtime_loaded, feature_mask, gguf_slice_write), + } +} + +fn load_skippy_runtime_for_probe(skippy_runtime_libraries: &[PathBuf]) -> Option { + if skippy_runtime_libraries.is_empty() || skippy_ffi::native_runtime_loaded() { + return None; + } + // The caller explicitly supplied these native runtime libraries for probing. + // Loading arbitrary libraries would be unsafe, so the command never guesses. + let result = unsafe { skippy_ffi::load_native_runtime_libraries(skippy_runtime_libraries) }; + result.err().map(|err| err.to_string()) +} + +fn skippy_abi_reason( + runtime_loaded: bool, + feature_mask: Option, + gguf_slice_write: bool, +) -> String { + if !runtime_loaded { + return "no Skippy native runtime library was loaded for ABI probing".to_string(); + } + if feature_mask.is_none() { + return "loaded Skippy runtime does not expose skippy_abi_features".to_string(); + } + if gguf_slice_write { + return "loaded Skippy ABI exposes GGUF slice writing and the linked llama symbols can be used for GGUF quantization, but not HF checkpoint conversion".to_string(); + } + "loaded Skippy ABI exposes staged inference/runtime entry points and linked llama symbols can be used for GGUF quantization, but not HF checkpoint conversion".to_string() +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] +pub enum BackendKind { + NativeRust, + LlamaApi, + SkippyAbi, +} + +impl BackendKind { + pub fn as_str(self) -> &'static str { + match self { + Self::NativeRust => "native-rust", + Self::LlamaApi => "llama-api", + Self::SkippyAbi => "skippy-abi", + } + } +} + +pub fn ensure_convert_backend(kind: BackendKind) -> Result<()> { + ensure!( + matches!(kind, BackendKind::NativeRust), + "backend {} cannot convert HF checkpoints yet: {}", + kind.as_str(), + capabilities(&[]).skippy_abi.reason + ); + Ok(()) +} + +pub fn ensure_quant_backend(kind: BackendKind) -> Result<()> { + ensure!( + matches!(kind, BackendKind::LlamaApi | BackendKind::SkippyAbi), + "backend {} cannot quantize GGUFs yet: {}", + kind.as_str(), + capabilities(&[]).skippy_abi.reason + ); + Ok(()) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct BackendRunStatus { + pub status_code: Option, + pub success: bool, +} + +impl BackendRunStatus { + pub fn from_code(status_code: i32) -> Self { + Self { + status_code: Some(status_code), + success: status_code == 0, + } + } +} + +pub fn ensure_success(status: BackendRunStatus, command: &[String]) -> Result<()> { + ensure!( + status.success, + "command failed with status_code {:?}: {}", + status.status_code, + shell_words(command) + ); + Ok(()) +} + +fn shell_words(command: &[String]) -> String { + command.join(" ") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn reports_current_backend_capabilities() { + let capabilities = capabilities(&[]); + assert!(capabilities.native_rust.convert_hf_to_gguf); + assert!(!capabilities.native_rust.llama_quantize); + assert!(!capabilities.llama_api.convert_hf_to_gguf); + assert!(capabilities.llama_api.llama_quantize); + assert!(capabilities.llama_api.runtime_loaded); + assert!(!capabilities.skippy_abi.convert_hf_to_gguf); + assert!(!capabilities.skippy_abi.llama_quantize); + assert!(!capabilities.skippy_abi.runtime_loaded); + assert_eq!(capabilities.skippy_abi.feature_mask, None); + assert!(!capabilities.skippy_abi.model_introspection); + assert!(!capabilities.skippy_abi.gguf_slice_write); + assert_eq!(capabilities.skippy_abi.load_error, None); + assert!(capabilities.skippy_abi.reason.contains("Skippy")); + } + + #[test] + fn rejects_skippy_abi_conversion_backend_until_supported() { + assert!(ensure_convert_backend(BackendKind::SkippyAbi).is_err()); + } + + #[test] + fn accepts_skippy_abi_quant_backend() { + assert!(ensure_quant_backend(BackendKind::SkippyAbi).is_ok()); + } + + #[test] + fn validates_backend_run_status() { + let ok = BackendRunStatus { + status_code: Some(0), + success: true, + }; + let failed = BackendRunStatus { + status_code: Some(2), + success: false, + }; + assert!(ensure_success(ok, &["tool".to_string()]).is_ok()); + assert!(ensure_success(failed, &["tool".to_string()]).is_err()); + } +} diff --git a/crates/skippy-quantize/src/command_reports.rs b/crates/skippy-quantize/src/command_reports.rs new file mode 100644 index 0000000000..445d325389 --- /dev/null +++ b/crates/skippy-quantize/src/command_reports.rs @@ -0,0 +1,37 @@ +use std::path::PathBuf; + +use serde::Serialize; + +#[derive(Debug, Serialize)] +pub(crate) struct TensorTypeValidation { + pub(crate) valid: bool, + pub(crate) entry_count: usize, +} + +#[derive(Debug, Serialize)] +pub(crate) struct QuantWindowPlan { + pub(crate) first_split: u32, + pub(crate) last_split: u32, + pub(crate) staged_first_shard: PathBuf, + pub(crate) output_prefix: PathBuf, + pub(crate) command: Vec, +} + +#[derive(Debug, Serialize)] +pub(crate) struct ConvertWindowPlan { + pub(crate) first_split: u32, + pub(crate) last_split: u32, + pub(crate) output_prefix: PathBuf, + pub(crate) command: Vec, +} + +#[derive(Debug, Serialize)] +pub(crate) struct SplitValidation { + pub(crate) root: PathBuf, + pub(crate) prefix: String, + pub(crate) expected_splits: u32, + pub(crate) completed_count: usize, + pub(crate) first_missing: Option, + pub(crate) last_present: Option, + pub(crate) complete: bool, +} diff --git a/crates/skippy-quantize/src/direct_convert.rs b/crates/skippy-quantize/src/direct_convert.rs new file mode 100644 index 0000000000..a07d8619cb --- /dev/null +++ b/crates/skippy-quantize/src/direct_convert.rs @@ -0,0 +1,463 @@ +use std::path::{Path, PathBuf}; + +use anyhow::{Context, Result, ensure}; +use clap::Parser; + +use crate::hf_checkpoint::resolve_auto_output_type; +use crate::locking::with_manifest_lock; +use crate::manifest::ensure_manifest; +use crate::preflight::run_job_preflight; +use crate::splits::parse_split_file_name; +use crate::types::ConvertOutputType; +use crate::verify::print_verify_on_complete; +use crate::{ + ConvertRunnerArgs, InitConvertArgs, RunConvertArgs, RunConvertWindowArgs, VerifyLoadArgs, + convert_manifest_from_args, prepare_convert_runner, run_convert_unlocked, + run_convert_window_once_with_manifest, +}; + +#[derive(Debug, Parser)] +pub(crate) struct DirectConvertArgs { + #[command(flatten)] + runner: ConvertRunnerArgs, + #[arg(long)] + target_prefix: Option, + #[arg(long)] + output_basename: Option, + #[arg(long, alias = "outtype", value_enum, default_value_t = ConvertOutputType::Auto)] + output_type: ConvertOutputType, + #[arg(short = 'o', long)] + outfile: Option, + #[arg(long, default_value_t = 1)] + expected_splits: u32, + #[arg(long, default_value_t = 1)] + window_size: u32, + #[arg(long)] + max_windows: Option, + #[arg(long)] + manifest: Option, + #[arg(long = "no-verify-on-complete", action = clap::ArgAction::SetFalse, default_value_t = true)] + verify_on_complete: bool, + #[command(flatten)] + verify_load: VerifyLoadArgs, + #[arg(long)] + preflight_only: bool, + #[arg(long)] + json: bool, + source: Option, + output: Option, +} + +pub(crate) fn run_direct_convert(args: DirectConvertArgs) -> Result<()> { + let runner = prepare_convert_runner(args.runner.clone())?; + let source = args + .source + .clone() + .context("missing source path: provide MODEL")?; + let output_type = direct_output_type(&runner, &args, &source)?; + let output = + if let Some(output) = resolved_output(args.output.as_deref(), args.outfile.as_deref())? { + output.to_path_buf() + } else { + default_output_path(&source, output_type)? + }; + ensure!( + !runner.has_upstream_shard_controls(), + "--skip-output-shards-before/--stop-output-shards-after are not accepted by direct native conversion; use convert-job/run-convert-window windowing" + ); + ensure!( + !is_templated_output_path(&output), + "templated output paths are not supported by the native converter" + ); + let target = derive_output( + &output, + args.target_prefix.as_deref(), + args.output_basename.as_deref(), + args.expected_splits, + )?; + let manifest_path = args + .manifest + .clone() + .unwrap_or_else(|| default_manifest_path(&target, output_type)); + let manifest_args = InitConvertArgs { + source, + target: target.root, + target_prefix: target.prefix, + output_basename: target.output_basename, + output_type, + expected_splits: args.expected_splits, + window_size: args.window_size, + manifest: manifest_path.clone(), + }; + let manifest = convert_manifest_from_args(&manifest_args)?; + if args.preflight_only { + return run_job_preflight( + &manifest_path, + &manifest, + None, + None, + runner.backend, + None, + args.json, + ); + } + if runner.dry_run { + return run_convert_window_once_with_manifest( + &RunConvertWindowArgs { + manifest: manifest_path, + runner, + json: args.json, + }, + &manifest, + ) + .map(|_| ()); + } + with_manifest_lock(&manifest_path, || { + ensure_manifest(&manifest_path, &manifest)?; + run_convert_unlocked(RunConvertArgs { + window: RunConvertWindowArgs { + manifest: manifest_path.clone(), + runner, + json: args.json, + }, + max_windows: args.max_windows, + })?; + print_verify_on_complete( + &manifest_path, + args.verify_load.options(args.verify_on_complete), + ) + }) +} + +fn direct_output_type( + _runner: &ConvertRunnerArgs, + args: &DirectConvertArgs, + source: &Path, +) -> Result { + resolve_auto_output_type(source, args.output_type) +} + +fn resolved_output<'a>( + positional: Option<&'a Path>, + outfile: Option<&'a Path>, +) -> Result> { + match (positional, outfile) { + (Some(_), Some(_)) => { + anyhow::bail!("provide either positional OUTPUT or --outfile, not both") + } + (Some(output), None) | (None, Some(output)) => Ok(Some(output)), + (None, None) => Ok(None), + } +} + +fn default_output_path(source: &Path, output_type: ConvertOutputType) -> Result { + let model_name = source + .file_name() + .and_then(|value| value.to_str()) + .with_context(|| { + format!( + "cannot derive default output name from {}", + source.display() + ) + })?; + let parent = source.parent().unwrap_or_else(|| Path::new(".")); + Ok(parent + .join(model_name) + .join(format!("{model_name}-{}.gguf", output_type.as_arg()))) +} + +fn is_templated_output_path(path: &Path) -> bool { + let Some(name) = path.file_name().and_then(|value| value.to_str()) else { + return false; + }; + ["{}", "{ftype}", "{outtype}", "{FTYPE}", "{OUTTYPE}"] + .iter() + .any(|marker| name.contains(marker)) +} + +#[derive(Debug)] +struct OutputLocation { + root: PathBuf, + prefix: String, + output_basename: String, +} + +fn derive_output( + path: &Path, + prefix_override: Option<&str>, + basename_override: Option<&str>, + expected_splits: u32, +) -> Result { + let (root, prefix) = derive_root_and_prefix(path, prefix_override)?; + let output_basename = match basename_override { + Some(value) => value.to_string(), + None => output_basename(path, expected_splits)?, + }; + Ok(OutputLocation { + root, + prefix, + output_basename, + }) +} + +fn derive_root_and_prefix(path: &Path, prefix_override: Option<&str>) -> Result<(PathBuf, String)> { + let parent = path + .parent() + .with_context(|| format!("path has no parent directory: {}", path.display()))?; + if parent.as_os_str().is_empty() || parent == Path::new(".") { + return Ok(( + PathBuf::from("."), + prefix_override.unwrap_or("").to_string(), + )); + } + let prefix = match prefix_override { + Some(value) => value.to_string(), + None => parent + .file_name() + .and_then(|value| value.to_str()) + .with_context(|| format!("cannot derive prefix from {}", path.display()))? + .to_string(), + }; + let root = if prefix.is_empty() { + parent.to_path_buf() + } else { + parent + .parent() + .with_context(|| format!("path has no root above prefix: {}", path.display()))? + .to_path_buf() + }; + Ok((root, prefix)) +} + +fn output_basename(path: &Path, expected_splits: u32) -> Result { + let file_name = path + .file_name() + .and_then(|value| value.to_str()) + .with_context(|| format!("invalid output file name: {}", path.display()))?; + let stem = file_name + .strip_suffix(".gguf") + .with_context(|| format!("output must be a GGUF path: {}", path.display()))?; + if let Some((_, total)) = parse_split_file_name(file_name) { + ensure!( + total == expected_splits, + "output split total {total} does not match --expected-splits {expected_splits}" + ); + let (before_total, _) = stem.rsplit_once("-of-").with_context(|| { + format!( + "invalid split output file name after parse: {}", + path.display() + ) + })?; + let (base, _) = before_total.rsplit_once('-').with_context(|| { + format!( + "invalid split output file name after parse: {}", + path.display() + ) + })?; + return Ok(base.to_string()); + } + Ok(stem.to_string()) +} + +fn default_manifest_path(target: &OutputLocation, output_type: ConvertOutputType) -> PathBuf { + target.root.join(&target.prefix).join(format!( + ".{}.{}.skippy-convert.json", + target.output_basename, + output_type.as_arg() + )) +} + +#[cfg(test)] +mod tests { + use std::fs; + + use super::*; + + #[test] + fn parses_short_outfile_and_upstream_auto_default() { + let args = DirectConvertArgs::try_parse_from([ + "skippy-quantize convert", + "-o", + "/repo/auto/model.gguf", + "/models/source", + ]) + .unwrap(); + + assert_eq!(args.outfile, Some(PathBuf::from("/repo/auto/model.gguf"))); + assert_eq!(args.source, Some(PathBuf::from("/models/source"))); + assert_eq!(args.output_type, ConvertOutputType::Auto); + } + + #[test] + fn parses_source_without_output_for_default_filename_shape() { + let args = DirectConvertArgs::try_parse_from(["skippy-quantize convert", "/models/source"]) + .unwrap(); + + assert_eq!(args.source, Some(PathBuf::from("/models/source"))); + assert!(args.output.is_none()); + assert!(args.outfile.is_none()); + assert_eq!(args.runner.split_max_size, "0"); + } + + #[test] + fn native_convert_rejects_unsupported_python_converter_flags() { + let args = DirectConvertArgs::try_parse_from([ + "skippy-quantize convert", + "--vocab-only", + "/models/source", + "/repo/BF16/model.gguf", + ]) + .unwrap(); + let error = prepare_convert_runner(args.runner.clone()).unwrap_err(); + + assert!(error.to_string().contains("--vocab-only")); + } + + #[test] + fn native_convert_accepts_supported_runner_flags() { + let args = DirectConvertArgs::try_parse_from([ + "skippy-quantize convert", + "--mtp", + "--max-memory", + "32G", + "--stream-buffer-bytes", + "1024", + "/models/source", + "/repo/BF16/model.gguf", + ]) + .unwrap(); + + assert!(prepare_convert_runner(args.runner.clone()).is_ok()); + } + + #[test] + fn native_convert_rejects_conflicting_mtp_flags() { + let args = DirectConvertArgs::try_parse_from([ + "skippy-quantize convert", + "--mtp", + "--no-mtp", + "/models/source", + "/repo/BF16/model.gguf", + ]) + .unwrap(); + let error = prepare_convert_runner(args.runner.clone()).unwrap_err(); + + assert!(error.to_string().contains("mutually exclusive")); + } + + #[test] + fn detects_templated_output_paths() { + assert!(is_templated_output_path(Path::new( + "/repo/model-{ftype}.gguf" + ))); + assert!(is_templated_output_path(Path::new( + "/repo/model-{OUTTYPE}.gguf" + ))); + assert!(is_templated_output_path(Path::new("/repo/model-{}.gguf"))); + assert!(!is_templated_output_path(Path::new( + "/repo/model-bf16.gguf" + ))); + } + + #[test] + fn derives_output_basename_from_unsplit_path() { + let output = Path::new("/repo/BF16/model-bf16.gguf"); + let location = derive_output(output, None, None, 3).unwrap(); + + assert_eq!(location.root, PathBuf::from("/repo")); + assert_eq!(location.prefix, "BF16"); + assert_eq!(location.output_basename, "model-bf16"); + } + + #[test] + fn derives_current_directory_output_location() { + let location = derive_output(Path::new("model-bf16.gguf"), None, None, 1).unwrap(); + + assert_eq!(location.root, PathBuf::from(".")); + assert_eq!(location.prefix, ""); + assert_eq!(location.output_basename, "model-bf16"); + } + + #[test] + fn derives_output_basename_from_split_path() { + let output = Path::new("/repo/BF16/model-bf16-00001-of-00003.gguf"); + let location = derive_output(output, None, None, 3).unwrap(); + + assert_eq!(location.output_basename, "model-bf16"); + } + + #[test] + fn rejects_output_with_wrong_split_total() { + let output = Path::new("/repo/BF16/model-bf16-00001-of-00002.gguf"); + assert!(derive_output(output, None, None, 3).is_err()); + } + + #[test] + fn resolves_outfile_without_positional_output() { + let outfile = Path::new("/repo/BF16/model.gguf"); + assert_eq!(resolved_output(None, Some(outfile)).unwrap(), Some(outfile)); + } + + #[test] + fn resolves_missing_output_as_passthrough_default() { + assert_eq!(resolved_output(None, None).unwrap(), None); + } + + #[test] + fn derives_default_output_path_for_source_only_resumable_convert() { + assert_eq!( + default_output_path(Path::new("/models/source"), ConvertOutputType::Bf16).unwrap(), + PathBuf::from("/models/source/source-bf16.gguf") + ); + assert_eq!( + default_output_path(Path::new("source"), ConvertOutputType::Auto).unwrap(), + PathBuf::from("source/source-auto.gguf") + ); + } + + #[test] + fn rejects_conflicting_output_forms() { + assert!( + resolved_output( + Some(Path::new("/repo/BF16/a.gguf")), + Some(Path::new("/repo/BF16/b.gguf")), + ) + .is_err() + ); + } + + #[test] + fn direct_convert_dry_run_does_not_write_manifest_or_output() { + let root = unique_temp_dir("direct-convert-dry-run"); + let source = root.join("checkpoint"); + let output = root.join("BF16").join("model-bf16.gguf"); + let manifest = root.join("manifest.json"); + let args = DirectConvertArgs::try_parse_from([ + "skippy-quantize convert", + "--dry-run", + "--output-type", + "bf16", + "--manifest", + manifest.to_str().unwrap(), + source.to_str().unwrap(), + output.to_str().unwrap(), + ]) + .unwrap(); + + run_direct_convert(args).unwrap(); + + assert!(!manifest.exists()); + assert!(!root.join("BF16").exists()); + fs::remove_dir_all(root).ok(); + } + + fn unique_temp_dir(name: &str) -> PathBuf { + static NEXT_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let id = NEXT_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + std::env::temp_dir().join(format!("skippy-quantize-{name}-{nanos}-{id}")) + } +} diff --git a/crates/skippy-quantize/src/direct_quantize.rs b/crates/skippy-quantize/src/direct_quantize.rs new file mode 100644 index 0000000000..20be8c4e26 --- /dev/null +++ b/crates/skippy-quantize/src/direct_quantize.rs @@ -0,0 +1,644 @@ +use std::path::{Path, PathBuf}; + +use anyhow::{Context, Result, ensure}; +use clap::Parser; + +use crate::locking::with_manifest_lock; +use crate::manifest::ensure_manifest; +use crate::preflight::run_job_preflight; +use crate::splits::{SplitWindow, parse_split_file_name, validate_split_window}; +use crate::types::QuantSpec; +use crate::verify::print_verify_on_complete; +use crate::{ + InitQuantArgs, QuantRunnerArgs, RunQuantArgs, RunQuantWindowArgs, VerifyLoadArgs, + prepare_quant_runner, quant_backend_path, quant_manifest_from_args, run_quant_unlocked, + run_quant_window_once_with_manifest, +}; + +#[derive(Debug, Parser)] +pub(crate) struct DirectQuantizeArgs { + #[command(flatten)] + runner: QuantRunnerArgs, + #[arg(long)] + source_prefix: Option, + #[arg(long)] + target_prefix: Option, + #[arg(long)] + output_basename: Option, + #[arg(long)] + tensor_type_file: Option, + #[arg(long, default_value_t = 1)] + window_size: u32, + #[arg(long)] + max_windows: Option, + #[arg(long)] + manifest: Option, + #[arg(long = "no-verify-on-complete", action = clap::ArgAction::SetFalse, default_value_t = true)] + verify_on_complete: bool, + #[command(flatten)] + verify_load: VerifyLoadArgs, + #[arg(long)] + preflight_only: bool, + #[arg(long)] + json: bool, + #[arg(long)] + keep_split: bool, + #[arg(long)] + first_split: Option, + #[arg(long)] + last_split: Option, + input: PathBuf, + #[arg(value_name = "OUTPUT_OR_QUANT", num_args = 1..=3)] + positional: Vec, +} + +pub(crate) fn run_direct_quantize(args: DirectQuantizeArgs) -> Result<()> { + let mut runner = prepare_quant_runner(args.runner.clone())?; + let positional = parse_direct_quantize_positionals(&args.positional)?; + positional + .quant + .validate_recipe_requirements(args.tensor_type_file.is_some()) + .map_err(anyhow::Error::msg)?; + apply_positional_nthreads(&mut runner, positional.nthreads)?; + let source = derive_input(&args.input, args.source_prefix.as_deref())?; + let window_override = args.manual_split_window(source.expected_splits)?; + let output = if let Some(output) = positional.output.as_deref() { + output.to_path_buf() + } else { + default_output_path(&args.input, &positional.quant)? + }; + let target = derive_output( + &output, + args.target_prefix.as_deref(), + args.output_basename.as_deref(), + source.expected_splits, + )?; + let manifest_path = args + .manifest + .clone() + .unwrap_or_else(|| default_manifest_path(&target, &positional.quant)); + let manifest_args = InitQuantArgs { + source: source.root, + source_prefix: source.prefix, + target: target.root, + target_prefix: target.prefix, + output_basename: target.output_basename, + quant: positional.quant.clone(), + tensor_type_file: args.tensor_type_file, + window_size: args.window_size, + manifest: manifest_path.clone(), + }; + let manifest = quant_manifest_from_args(&manifest_args)?; + if args.preflight_only { + return run_job_preflight( + &manifest_path, + &manifest, + Some((&manifest_args.source, &manifest_args.source_prefix)), + window_override, + runner.backend, + quant_backend_path(&runner), + args.json, + ); + } + if runner.dry_run { + return run_quant_window_once_with_manifest( + &RunQuantWindowArgs { + manifest: manifest_path, + runner, + json: args.json, + }, + &manifest, + window_override, + ) + .map(|_| ()); + } + with_manifest_lock(&manifest_path, || { + ensure_manifest(&manifest_path, &manifest)?; + run_quant_unlocked(RunQuantArgs { + window: RunQuantWindowArgs { + manifest: manifest_path.clone(), + runner, + json: args.json, + }, + window_override, + max_windows: args.max_windows, + })?; + print_verify_on_complete( + &manifest_path, + args.verify_load.options(args.verify_on_complete), + ) + }) +} + +impl DirectQuantizeArgs { + fn has_upstream_split_controls(&self) -> bool { + self.keep_split || self.first_split.is_some() || self.last_split.is_some() + } + + fn manual_split_window(&self, expected_splits: u32) -> Result> { + if !self.has_upstream_split_controls() { + return Ok(None); + } + ensure!( + self.keep_split, + "--first-split and --last-split require --keep-split" + ); + let window = SplitWindow { + first_split: self.first_split.unwrap_or(1), + last_split: self.last_split.unwrap_or(expected_splits), + }; + validate_split_window(window, expected_splits)?; + Ok(Some(window)) + } +} + +#[derive(Debug, PartialEq, Eq)] +struct DirectQuantizePositionals { + output: Option, + quant: QuantSpec, + nthreads: Option, +} + +fn parse_direct_quantize_positionals(tokens: &[String]) -> Result { + match tokens { + [quant] => Ok(DirectQuantizePositionals { + output: None, + quant: parse_quant_type(quant)?, + nthreads: None, + }), + [first, second] => { + if let Ok(quant) = first.parse() { + return Ok(DirectQuantizePositionals { + output: None, + quant, + nthreads: Some(parse_nthreads(second)?), + }); + } + Ok(DirectQuantizePositionals { + output: Some(PathBuf::from(first)), + quant: parse_quant_type(second)?, + nthreads: None, + }) + } + [output, quant, nthreads] => Ok(DirectQuantizePositionals { + output: Some(PathBuf::from(output)), + quant: parse_quant_type(quant)?, + nthreads: Some(parse_nthreads(nthreads)?), + }), + _ => { + anyhow::bail!("expected QUANT, QUANT NTHREADS, OUTPUT QUANT, or OUTPUT QUANT NTHREADS") + } + } +} + +fn parse_quant_type(raw: &str) -> Result { + raw.parse::() + .map_err(|error| anyhow::anyhow!(error)) +} + +fn parse_nthreads(raw: &str) -> Result { + raw.parse::() + .with_context(|| format!("invalid nthreads {raw:?}")) +} + +#[derive(Debug)] +struct InputLocation { + root: PathBuf, + prefix: String, + expected_splits: u32, +} + +#[derive(Debug)] +struct OutputLocation { + root: PathBuf, + prefix: String, + output_basename: String, +} + +fn derive_input(path: &Path, prefix_override: Option<&str>) -> Result { + let file_name = file_name(path)?; + let expected_splits = if let Some((index, expected_splits)) = parse_split_file_name(file_name) { + ensure!( + index == 1, + "direct quantize requires the first split shard, got shard {index}: {}", + path.display() + ); + expected_splits + } else { + ensure!( + file_name.ends_with(".gguf"), + "input must be a GGUF file: {}", + path.display() + ); + 1 + }; + let (root, prefix) = derive_root_and_prefix(path, prefix_override)?; + Ok(InputLocation { + root, + prefix, + expected_splits, + }) +} + +fn derive_output( + path: &Path, + prefix_override: Option<&str>, + basename_override: Option<&str>, + expected_splits: u32, +) -> Result { + let (root, prefix) = derive_root_and_prefix(path, prefix_override)?; + let output_basename = match basename_override { + Some(value) => value.to_string(), + None => output_basename(path, expected_splits)?, + }; + Ok(OutputLocation { + root, + prefix, + output_basename, + }) +} + +fn derive_root_and_prefix(path: &Path, prefix_override: Option<&str>) -> Result<(PathBuf, String)> { + let parent = path + .parent() + .with_context(|| format!("path has no parent directory: {}", path.display()))?; + if parent.as_os_str().is_empty() || parent == Path::new(".") { + return Ok(( + PathBuf::from("."), + prefix_override.unwrap_or("").to_string(), + )); + } + let prefix = match prefix_override { + Some(value) => value.to_string(), + None => parent + .file_name() + .and_then(|value| value.to_str()) + .with_context(|| format!("cannot derive prefix from {}", path.display()))? + .to_string(), + }; + let root = if prefix.is_empty() { + parent.to_path_buf() + } else { + parent + .parent() + .with_context(|| format!("path has no root above prefix: {}", path.display()))? + .to_path_buf() + }; + Ok((root, prefix)) +} + +fn output_basename(path: &Path, expected_splits: u32) -> Result { + let file_name = file_name(path)?; + let stem = file_name + .strip_suffix(".gguf") + .with_context(|| format!("output must be a GGUF path: {}", path.display()))?; + if let Some((_, total)) = parse_split_file_name(file_name) { + ensure!( + total == expected_splits, + "output split total {total} does not match input split total {expected_splits}" + ); + let (before_total, _) = stem.rsplit_once("-of-").with_context(|| { + format!( + "invalid split output file name after parse: {}", + path.display() + ) + })?; + let (base, _) = before_total.rsplit_once('-').with_context(|| { + format!( + "invalid split output file name after parse: {}", + path.display() + ) + })?; + return Ok(base.to_string()); + } + Ok(stem.to_string()) +} + +fn default_manifest_path(target: &OutputLocation, quant: &QuantSpec) -> PathBuf { + target.root.join(&target.prefix).join(format!( + ".{}.{}.skippy-quantize.json", + target.output_basename, + quant.output_name() + )) +} + +fn default_output_path(input: &Path, quant: &QuantSpec) -> Result { + let parent = input.parent().unwrap_or_else(|| Path::new(".")); + Ok(parent.join(format!("ggml-model-{}.gguf", quant.output_name()))) +} + +fn file_name(path: &Path) -> Result<&str> { + path.file_name() + .and_then(|value| value.to_str()) + .with_context(|| format!("invalid path file name: {}", path.display())) +} + +fn apply_positional_nthreads( + runner: &mut QuantRunnerArgs, + positional_nthreads: Option, +) -> Result<()> { + if let Some(positional_nthreads) = positional_nthreads { + if let Some(flag_nthreads) = runner.nthreads { + ensure!( + flag_nthreads == positional_nthreads, + "positional nthreads {positional_nthreads} conflicts with --nthreads {flag_nthreads}" + ); + } + runner.nthreads = Some(positional_nthreads); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::fs; + + use crate::types::QuantType; + + use super::*; + + #[test] + fn parses_quant_without_output_for_upstream_dry_run_shape() { + let parsed = parse_direct_quantize_positionals(&["Q4_K".to_string()]).unwrap(); + + assert_eq!( + parsed, + DirectQuantizePositionals { + output: None, + quant: QuantType::Q4K.into(), + nthreads: None, + } + ); + } + + #[test] + fn parses_quant_without_output_for_upstream_default_output_shape() { + let args = DirectQuantizeArgs::try_parse_from([ + "skippy-quantize quantize", + "/repo/model.gguf", + "Q4_K", + ]) + .unwrap(); + + assert_eq!(args.input, PathBuf::from("/repo/model.gguf")); + assert_eq!( + parse_direct_quantize_positionals(&args.positional).unwrap(), + DirectQuantizePositionals { + output: None, + quant: QuantType::Q4K.into(), + nthreads: None, + } + ); + assert!(!args.runner.dry_run); + assert!(!args.runner.leave_output_tensor); + } + + #[test] + fn parses_numeric_ftype_like_upstream_llama_quantize() { + let args = DirectQuantizeArgs::try_parse_from([ + "skippy-quantize quantize", + "/repo/model.gguf", + "15", + ]) + .unwrap(); + + assert_eq!( + parse_direct_quantize_positionals(&args.positional).unwrap(), + DirectQuantizePositionals { + output: None, + quant: QuantType::Q4K.into(), + nthreads: None, + } + ); + } + + #[test] + fn rejects_profile_quant_label_for_direct_quantize() { + let error = parse_direct_quantize_positionals(&["UD-Q3_K_S".to_string()]).unwrap_err(); + + assert!( + error.to_string().contains("custom tensor-type recipes"), + "profile labels should not be accepted as quant modes: {error}" + ); + } + + #[test] + fn manual_split_window_matches_upstream_keep_split_defaults() { + let args = DirectQuantizeArgs::try_parse_from([ + "skippy-quantize quantize", + "--backend", + "llama-api", + "--keep-split", + "--first-split", + "3", + "/repo/model-00001-of-00005.gguf", + "/repo/q4/model-q4.gguf", + "Q4_K", + ]) + .unwrap(); + + let window = args.manual_split_window(5).unwrap().unwrap(); + + assert_eq!(window.first_split, 3); + assert_eq!(window.last_split, 5); + } + + #[test] + fn manual_split_window_rejects_first_or_last_without_keep_split() { + let args = DirectQuantizeArgs::try_parse_from([ + "skippy-quantize quantize", + "--backend", + "llama-api", + "--last-split", + "3", + "/repo/model-00001-of-00005.gguf", + "/repo/q4/model-q4.gguf", + "Q4_K", + ]) + .unwrap(); + + let error = args.manual_split_window(5).unwrap_err(); + + assert!( + error + .to_string() + .contains("--first-split and --last-split require --keep-split") + ); + } + + #[test] + fn manual_split_window_rejects_out_of_range_bounds() { + let args = DirectQuantizeArgs::try_parse_from([ + "skippy-quantize quantize", + "--backend", + "llama-api", + "--keep-split", + "--first-split", + "4", + "--last-split", + "6", + "/repo/model-00001-of-00005.gguf", + "/repo/q4/model-q4.gguf", + "Q4_K", + ]) + .unwrap(); + + assert!(args.manual_split_window(5).is_err()); + } + + #[test] + fn parses_quant_and_threads_without_output() { + let parsed = + parse_direct_quantize_positionals(&["Q4_K".to_string(), "8".to_string()]).unwrap(); + + assert_eq!( + parsed, + DirectQuantizePositionals { + output: None, + quant: QuantType::Q4K.into(), + nthreads: Some(8), + } + ); + } + + #[test] + fn parses_output_quant_and_threads() { + let parsed = parse_direct_quantize_positionals(&[ + "/repo/Q4/model.gguf".to_string(), + "Q4_K".to_string(), + "8".to_string(), + ]) + .unwrap(); + + assert_eq!( + parsed, + DirectQuantizePositionals { + output: Some(PathBuf::from("/repo/Q4/model.gguf")), + quant: QuantType::Q4K.into(), + nthreads: Some(8), + } + ); + } + + #[test] + fn direct_quantize_dry_run_does_not_write_manifest_stage_or_output() { + let root = unique_temp_dir("direct-quant-dry-run"); + let source_dir = root.join("source").join("BF16"); + fs::create_dir_all(&source_dir).unwrap(); + let first = source_dir.join("model-bf16-00001-of-00002.gguf"); + fs::write(&first, b"not-a-real-gguf").unwrap(); + fs::write(source_dir.join("model-bf16-00002-of-00002.gguf"), b"").unwrap(); + let output = root.join("target").join("Q4_K").join("model-q4.gguf"); + let manifest = root.join("manifest.json"); + let work_dir = root.join("work"); + let spool_dir = root.join("spool"); + let args = DirectQuantizeArgs::try_parse_from([ + "skippy-quantize quantize", + "--dry-run", + "--manifest", + manifest.to_str().unwrap(), + "--work-dir", + work_dir.to_str().unwrap(), + "--spool-dir", + spool_dir.to_str().unwrap(), + "--keep-split", + first.to_str().unwrap(), + output.to_str().unwrap(), + "Q4_K", + ]) + .unwrap(); + + run_direct_quantize(args).unwrap(); + + assert!(!manifest.exists()); + assert!(!work_dir.exists()); + assert!(!spool_dir.exists()); + assert!(!root.join("target").exists()); + fs::remove_dir_all(root).ok(); + } + + fn unique_temp_dir(name: &str) -> PathBuf { + static NEXT_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let id = NEXT_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + std::env::temp_dir().join(format!("skippy-quantize-{name}-{nanos}-{id}")) + } + + #[test] + fn derives_split_input_location() { + let input = Path::new("/repo/BF16/model-00001-of-00003.gguf"); + let location = derive_input(input, None).unwrap(); + + assert_eq!(location.root, PathBuf::from("/repo")); + assert_eq!(location.prefix, "BF16"); + assert_eq!(location.expected_splits, 3); + } + + #[test] + fn derives_unsplit_input_location_as_one_shard() { + let input = Path::new("/repo/BF16/model.gguf"); + let location = derive_input(input, None).unwrap(); + + assert_eq!(location.root, PathBuf::from("/repo")); + assert_eq!(location.prefix, "BF16"); + assert_eq!(location.expected_splits, 1); + } + + #[test] + fn derives_current_directory_input_and_output_locations() { + let input = derive_input(Path::new("model.gguf"), None).unwrap(); + let output = derive_output(Path::new("ggml-model-Q4_K.gguf"), None, None, 1).unwrap(); + + assert_eq!(input.root, PathBuf::from(".")); + assert_eq!(input.prefix, ""); + assert_eq!(input.expected_splits, 1); + assert_eq!(output.root, PathBuf::from(".")); + assert_eq!(output.prefix, ""); + assert_eq!(output.output_basename, "ggml-model-Q4_K"); + } + + #[test] + fn derives_default_output_path_for_no_output_quantize_shape() { + assert_eq!( + default_output_path(Path::new("/repo/BF16/model.gguf"), &QuantType::Q4K.into()) + .unwrap(), + PathBuf::from("/repo/BF16/ggml-model-Q4_K.gguf") + ); + assert_eq!( + default_output_path(Path::new("model.gguf"), &QuantType::Q2KS.into()).unwrap(), + PathBuf::from("ggml-model-Q2_K_S.gguf") + ); + assert_eq!( + default_output_path(Path::new("/repo/BF16/model.gguf"), &QuantType::Q3KS.into()) + .unwrap(), + PathBuf::from("/repo/BF16/ggml-model-Q3_K_S.gguf") + ); + } + + #[test] + fn derives_output_basename_from_unsplit_output_path() { + let output = Path::new("/repo/Q2_K/model-q2.gguf"); + let location = derive_output(output, None, None, 3).unwrap(); + + assert_eq!(location.root, PathBuf::from("/repo")); + assert_eq!(location.prefix, "Q2_K"); + assert_eq!(location.output_basename, "model-q2"); + } + + #[test] + fn derives_output_basename_from_split_output_path() { + let output = Path::new("/repo/Q2_K/model-q2-00001-of-00003.gguf"); + let location = derive_output(output, None, None, 3).unwrap(); + + assert_eq!(location.output_basename, "model-q2"); + } + + #[test] + fn rejects_non_first_input_shard() { + let input = Path::new("/repo/BF16/model-00002-of-00003.gguf"); + assert!(derive_input(input, None).is_err()); + } +} diff --git a/crates/skippy-quantize/src/float_convert.rs b/crates/skippy-quantize/src/float_convert.rs new file mode 100644 index 0000000000..9044b88989 --- /dev/null +++ b/crates/skippy-quantize/src/float_convert.rs @@ -0,0 +1,217 @@ +use std::io::Write; + +use anyhow::{Context, Result}; + +use crate::types::ConvertOutputType; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum FloatDType { + F32, + F16, + Bf16, +} + +impl FloatDType { + pub(crate) fn from_safetensor(dtype: &str) -> Option { + match dtype { + "F32" => Some(Self::F32), + "F16" => Some(Self::F16), + "BF16" => Some(Self::Bf16), + _ => None, + } + } + + pub(crate) fn byte_size(self) -> u64 { + match self { + Self::F32 => 4, + Self::F16 | Self::Bf16 => 2, + } + } +} + +pub(crate) fn target_dtype_for( + source_dtype: FloatDType, + output_type: Option, +) -> Result { + match output_type { + None => Ok(source_dtype), + Some(ConvertOutputType::F32) => Ok(FloatDType::F32), + Some(ConvertOutputType::F16) => Ok(FloatDType::F16), + Some(ConvertOutputType::Bf16) => Ok(FloatDType::Bf16), + Some(other) => { + anyhow::bail!( + "native conversion does not support output type {}", + other.as_arg() + ) + } + } +} + +pub(crate) fn target_dtype_for_tensor( + source_dtype: FloatDType, + output_type: Option, + shape: &[u64], +) -> Result { + if shape.len() <= 1 + && matches!( + output_type, + Some(ConvertOutputType::F16 | ConvertOutputType::Bf16) + ) + { + return Ok(FloatDType::F32); + } + target_dtype_for(source_dtype, output_type) +} + +pub(crate) fn convert_float_chunk( + input: &[u8], + source_dtype: FloatDType, + target_dtype: FloatDType, + writer: &mut W, +) -> Result { + let element_count = input.len() / source_dtype.byte_size() as usize; + let output_len = element_count + .checked_mul(target_dtype.byte_size() as usize) + .context("converted chunk byte length overflow")?; + let mut output = Vec::with_capacity(output_len); + for index in 0..element_count { + let value = read_float_element(input, source_dtype, index); + write_float_element(&mut output, target_dtype, value); + } + writer.write_all(&output)?; + Ok(output.len() as u64) +} + +fn read_float_element(input: &[u8], dtype: FloatDType, index: usize) -> f32 { + match dtype { + FloatDType::F32 => { + let start = index * 4; + f32::from_le_bytes(input[start..start + 4].try_into().expect("slice length")) + } + FloatDType::F16 => { + let start = index * 2; + f16_bits_to_f32(u16::from_le_bytes( + input[start..start + 2].try_into().expect("slice length"), + )) + } + FloatDType::Bf16 => { + let start = index * 2; + f32::from_bits( + u32::from(u16::from_le_bytes( + input[start..start + 2].try_into().expect("slice length"), + )) << 16, + ) + } + } +} + +fn write_float_element(output: &mut Vec, dtype: FloatDType, value: f32) { + match dtype { + FloatDType::F32 => output.extend_from_slice(&value.to_le_bytes()), + FloatDType::F16 => output.extend_from_slice(&f32_to_f16_bits(value).to_le_bytes()), + FloatDType::Bf16 => output.extend_from_slice(&f32_to_bf16_bits(value).to_le_bytes()), + } +} + +fn f32_to_bf16_bits(value: f32) -> u16 { + let bits = value.to_bits(); + let rounding_bias = ((bits >> 16) & 1) + 0x7fff; + ((bits.wrapping_add(rounding_bias)) >> 16) as u16 +} + +fn f16_bits_to_f32(bits: u16) -> f32 { + let sign = (u32::from(bits & 0x8000)) << 16; + let exp = (bits >> 10) & 0x1f; + let frac = u32::from(bits & 0x03ff); + let value = if exp == 0 { + if frac == 0 { + sign + } else { + let mut frac = frac; + let mut exp_shift = -14_i32; + while frac & 0x0400 == 0 { + frac <<= 1; + exp_shift -= 1; + } + frac &= 0x03ff; + sign | (u32::try_from(exp_shift + 127).unwrap() << 23) | (frac << 13) + } + } else if exp == 0x1f { + sign | 0x7f80_0000 | (frac << 13) + } else { + sign | (u32::from(exp + 112) << 23) | (frac << 13) + }; + f32::from_bits(value) +} + +fn f32_to_f16_bits(value: f32) -> u16 { + let bits = value.to_bits(); + let sign = ((bits >> 16) & 0x8000) as u16; + let exp = ((bits >> 23) & 0xff) as i32; + let mant = bits & 0x007f_ffff; + + if exp == 0xff { + if mant == 0 { + return sign | 0x7c00; + } + return sign | 0x7e00; + } + + let half_exp = exp - 127 + 15; + if half_exp >= 0x1f { + return sign | 0x7c00; + } + if half_exp <= 0 { + if half_exp < -10 { + return sign; + } + let mantissa = mant | 0x0080_0000; + let shift = u32::try_from(14 - half_exp).unwrap(); + let mut half_mant = (mantissa >> shift) as u16; + if (mantissa >> (shift - 1)) & 1 != 0 { + half_mant = half_mant.saturating_add(1); + } + return sign | half_mant; + } + + let mut half = sign | ((half_exp as u16) << 10) | ((mant >> 13) as u16); + if (mant & 0x0000_1000) != 0 { + half = half.saturating_add(1); + } + half +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn target_dtype_rejects_unresolved_auto() { + assert_eq!( + target_dtype_for(FloatDType::F32, None).unwrap(), + FloatDType::F32 + ); + assert_eq!( + target_dtype_for(FloatDType::F32, Some(ConvertOutputType::Bf16)).unwrap(), + FloatDType::Bf16 + ); + assert!(target_dtype_for(FloatDType::F32, Some(ConvertOutputType::Auto)).is_err()); + } + + #[test] + fn target_dtype_keeps_rank_one_tensors_f32_for_float16_outputs() { + assert_eq!( + target_dtype_for_tensor(FloatDType::Bf16, Some(ConvertOutputType::Bf16), &[8]).unwrap(), + FloatDType::F32 + ); + assert_eq!( + target_dtype_for_tensor(FloatDType::Bf16, Some(ConvertOutputType::F16), &[8]).unwrap(), + FloatDType::F32 + ); + assert_eq!( + target_dtype_for_tensor(FloatDType::Bf16, Some(ConvertOutputType::Bf16), &[8, 8]) + .unwrap(), + FloatDType::Bf16 + ); + } +} diff --git a/crates/skippy-quantize/src/gguf_template.rs b/crates/skippy-quantize/src/gguf_template.rs new file mode 100644 index 0000000000..982347af02 --- /dev/null +++ b/crates/skippy-quantize/src/gguf_template.rs @@ -0,0 +1,709 @@ +use std::fs; +use std::path::Path; + +use anyhow::{Context, Result, ensure}; +use serde_json::Value; + +use crate::gguf_writer::GgufKv; +use crate::tokenizer_metadata::push_tokenizer_metadata; + +#[derive(Debug, Clone, Copy)] +pub(crate) struct MetadataOptions { + pub(crate) include_mtp: bool, +} + +impl Default for MetadataOptions { + fn default() -> Self { + Self { include_mtp: true } + } +} + +pub(crate) fn metadata_from_hf_config(source: &Path, tensor_count: usize) -> Result> { + metadata_from_hf_config_with_options(source, tensor_count, MetadataOptions::default()) +} + +pub(crate) fn mtp_layer_start_from_hf_config(source: &Path) -> Result> { + let config = read_hf_config(source)?; + let Some(nextn_layers) = optional_u32(&config, "num_nextn_predict_layers") + .or_else(|| optional_u32(&config, "mtp_num_hidden_layers")) + else { + return Ok(None); + }; + if nextn_layers == 0 { + return Ok(None); + } + required_u32(&config, "num_hidden_layers").map(Some) +} + +pub(crate) fn metadata_from_hf_config_with_options( + source: &Path, + tensor_count: usize, + options: MetadataOptions, +) -> Result> { + let config = read_hf_config(source)?; + let arch = architecture_name(&config)?; + let mut metadata = vec![ + GgufKv::string("general.architecture", arch), + GgufKv::string("general.name", model_name(source)), + GgufKv::bool("skippy.convert.raw_safetensors", false), + GgufKv::u64("skippy.convert.tensor_count", tensor_count as u64), + ]; + push_common_llm_metadata(&mut metadata, arch, &config, options)?; + push_attention_metadata(&mut metadata, arch, &config)?; + push_glm_dsa_indexer_metadata(&mut metadata, arch, &config)?; + push_moe_metadata(&mut metadata, arch, &config); + push_tokenizer_metadata(&mut metadata, source, &config)?; + if options.include_mtp { + push_if_u32( + &mut metadata, + arch, + "nextn_predict_layers", + &config, + "num_nextn_predict_layers", + ); + } + Ok(metadata) +} + +fn read_hf_config(source: &Path) -> Result { + let config_path = source.join("config.json"); + serde_json::from_slice( + &fs::read(&config_path).with_context(|| format!("read {}", config_path.display()))?, + ) + .with_context(|| format!("parse {}", config_path.display())) +} + +fn architecture_name(config: &Value) -> Result<&'static str> { + let model_type = config + .get("model_type") + .and_then(Value::as_str) + .unwrap_or_default(); + if matches!(model_type, "glm4_moe_lite" | "deepseek_v2") { + return Ok("deepseek2"); + } + if matches!(model_type, "glm4_moe" | "glm4v_moe") { + return Ok("glm4moe"); + } + if matches!(model_type, "glm_moe_dsa" | "glm-dsa") { + return Ok("glm-dsa"); + } + if is_unsupported_qwen3_variant(model_type) { + anyhow::bail!( + "native GGUF metadata for model_type={model_type:?} requires \ + Qwen3.5/Qwen3Next/Qwen3VL-specific metadata and tensor support; \ + use the external convert_hf_to_gguf.py backend" + ); + } + if model_type.starts_with("qwen3_moe") { + return Ok("qwen3moe"); + } + if model_type.starts_with("qwen2_moe") { + return Ok("qwen2moe"); + } + if model_type.starts_with("qwen3") { + return Ok("qwen3"); + } + if model_type.starts_with("qwen2") { + return Ok("qwen2"); + } + if matches!(model_type, "llama" | "mistral") { + return Ok("llama"); + } + anyhow::bail!("unsupported native GGUF metadata template for model_type={model_type:?}") +} + +fn is_unsupported_qwen3_variant(model_type: &str) -> bool { + let normalized = model_type + .chars() + .filter(|ch| ch.is_ascii_alphanumeric()) + .flat_map(char::to_lowercase) + .collect::(); + normalized.starts_with("qwen3next") + || normalized.starts_with("qwen3vl") + || normalized.starts_with("qwen35") +} + +fn push_common_llm_metadata( + metadata: &mut Vec, + arch: &str, + config: &Value, + options: MetadataOptions, +) -> Result<()> { + push_required_u32(metadata, arch, "vocab_size", config, "vocab_size")?; + push_required_u32( + metadata, + arch, + "context_length", + config, + "max_position_embeddings", + )?; + push_required_u32(metadata, arch, "embedding_length", config, "hidden_size")?; + let block_count = required_u32(config, "num_hidden_layers")? + + if options.include_mtp { + optional_u32(config, "num_nextn_predict_layers").unwrap_or(0) + } else { + 0 + }; + metadata.push(GgufKv::u32(&format!("{arch}.block_count"), block_count)); + push_if_u32( + metadata, + arch, + "feed_forward_length", + config, + "intermediate_size", + ); + if let Some(eps) = optional_f32(config, "rms_norm_eps") { + metadata.push(GgufKv::f32( + &format!("{arch}.attention.layer_norm_rms_epsilon"), + eps, + )); + } + Ok(()) +} + +fn push_attention_metadata(metadata: &mut Vec, arch: &str, config: &Value) -> Result<()> { + let head_count = required_u32(config, "num_attention_heads")?; + metadata.push(GgufKv::u32( + &format!("{arch}.attention.head_count"), + head_count, + )); + metadata.push(GgufKv::u32( + &format!("{arch}.attention.head_count_kv"), + optional_u32(config, "num_key_value_heads").unwrap_or(head_count), + )); + if let Some((nope, rope)) = + optional_u32(config, "qk_nope_head_dim").zip(optional_u32(config, "qk_rope_head_dim")) + { + metadata.push(GgufKv::u32( + &format!("{arch}.attention.key_length"), + nope + rope, + )); + metadata.push(GgufKv::u32(&format!("{arch}.rope.dimension_count"), rope)); + } else { + let head_dim = optional_u32(config, "head_dim").unwrap_or( + required_u32(config, "hidden_size")? / required_u32(config, "num_attention_heads")?, + ); + metadata.push(GgufKv::u32( + &format!("{arch}.attention.key_length"), + head_dim, + )); + metadata.push(GgufKv::u32( + &format!("{arch}.rope.dimension_count"), + head_dim, + )); + } + let value_len = optional_u32(config, "v_head_dim") + .or_else(|| optional_u32(config, "head_dim")) + .unwrap_or( + required_u32(config, "hidden_size")? / required_u32(config, "num_attention_heads")?, + ); + metadata.push(GgufKv::u32( + &format!("{arch}.attention.value_length"), + value_len, + )); + if let Some(theta) = optional_f32(config, "rope_theta") { + metadata.push(GgufKv::f32(&format!("{arch}.rope.freq_base"), theta)); + } + push_if_u32( + metadata, + arch, + "attention.q_lora_rank", + config, + "q_lora_rank", + ); + push_if_u32( + metadata, + arch, + "attention.kv_lora_rank", + config, + "kv_lora_rank", + ); + Ok(()) +} + +fn push_glm_dsa_indexer_metadata( + metadata: &mut Vec, + arch: &str, + config: &Value, +) -> Result<()> { + if arch != "glm-dsa" { + return Ok(()); + } + push_required_first_u32( + metadata, + arch, + "attention.indexer.head_count", + config, + &["index_n_heads", "indexer_n_head"], + )?; + push_required_first_u32( + metadata, + arch, + "attention.indexer.key_length", + config, + &["index_head_dim", "indexer_head_size"], + )?; + push_required_first_u32( + metadata, + arch, + "attention.indexer.top_k", + config, + &["index_topk", "indexer_top_k"], + )?; + Ok(()) +} + +fn push_moe_metadata(metadata: &mut Vec, arch: &str, config: &Value) { + push_first_u32( + metadata, + arch, + "expert_count", + config, + &["n_routed_experts", "num_experts"], + ); + push_first_u32( + metadata, + arch, + "expert_used_count", + config, + &["num_experts_per_tok"], + ); + push_first_u32( + metadata, + arch, + "expert_shared_count", + config, + &["n_shared_experts"], + ); + push_first_u32( + metadata, + arch, + "expert_feed_forward_length", + config, + &["moe_intermediate_size"], + ); + push_first_u32( + metadata, + arch, + "expert_shared_feed_forward_length", + config, + &["shared_expert_intermediate_size"], + ); + push_first_u32( + metadata, + arch, + "leading_dense_block_count", + config, + &["first_k_dense_replace"], + ); + if let Some(scale) = optional_f32(config, "routed_scaling_factor") { + metadata.push(GgufKv::f32(&format!("{arch}.expert_weights_scale"), scale)); + } + if let Some(norm) = optional_bool(config, "norm_topk_prob") { + metadata.push(GgufKv::bool(&format!("{arch}.expert_weights_norm"), norm)); + } +} + +fn push_required_first_u32( + metadata: &mut Vec, + arch: &str, + gguf_suffix: &str, + config: &Value, + config_keys: &[&str], +) -> Result<()> { + for config_key in config_keys { + if let Some(value) = optional_u32(config, config_key) { + ensure!( + value > 0, + "config value {config_key:?} must be greater than zero" + ); + metadata.push(GgufKv::u32(&format!("{arch}.{gguf_suffix}"), value)); + return Ok(()); + } + } + anyhow::bail!("config missing one of {config_keys:?}") +} + +fn push_first_u32( + metadata: &mut Vec, + arch: &str, + gguf_suffix: &str, + config: &Value, + config_keys: &[&str], +) { + for config_key in config_keys { + if let Some(value) = optional_u32(config, config_key) { + metadata.push(GgufKv::u32(&format!("{arch}.{gguf_suffix}"), value)); + return; + } + } +} + +fn push_required_u32( + metadata: &mut Vec, + arch: &str, + gguf_suffix: &str, + config: &Value, + config_key: &str, +) -> Result<()> { + metadata.push(GgufKv::u32( + &format!("{arch}.{gguf_suffix}"), + required_u32(config, config_key)?, + )); + Ok(()) +} + +fn push_if_u32( + metadata: &mut Vec, + arch: &str, + gguf_suffix: &str, + config: &Value, + config_key: &str, +) { + if let Some(value) = optional_u32(config, config_key) { + metadata.push(GgufKv::u32(&format!("{arch}.{gguf_suffix}"), value)); + } +} + +fn required_u32(config: &Value, key: &str) -> Result { + let value = optional_u32(config, key).with_context(|| format!("config missing {key:?}"))?; + ensure!(value > 0, "config value {key:?} must be greater than zero"); + Ok(value) +} + +fn optional_u32(config: &Value, key: &str) -> Option { + config + .get(key) + .and_then(Value::as_u64) + .and_then(|value| u32::try_from(value).ok()) +} + +fn optional_f32(config: &Value, key: &str) -> Option { + config + .get(key) + .and_then(Value::as_f64) + .map(|value| value as f32) +} + +fn optional_bool(config: &Value, key: &str) -> Option { + config.get(key).and_then(Value::as_bool) +} + +fn model_name(source: &Path) -> &str { + source + .file_name() + .and_then(|value| value.to_str()) + .unwrap_or("checkpoint") +} + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + + use super::*; + + #[test] + fn builds_glm_moe_lite_metadata_from_config() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + fs::write( + root.join("config.json"), + r#"{ + "model_type": "glm4_moe_lite", + "vocab_size": 154880, + "max_position_embeddings": 202752, + "hidden_size": 2048, + "intermediate_size": 10240, + "num_hidden_layers": 47, + "num_nextn_predict_layers": 1, + "num_attention_heads": 20, + "num_key_value_heads": 20, + "qk_nope_head_dim": 192, + "qk_rope_head_dim": 64, + "v_head_dim": 256, + "rope_theta": 1000000, + "q_lora_rank": 768, + "kv_lora_rank": 512, + "n_routed_experts": 64, + "num_experts_per_tok": 4, + "n_shared_experts": 1, + "moe_intermediate_size": 1536, + "first_k_dense_replace": 1, + "routed_scaling_factor": 1.8, + "norm_topk_prob": true, + "rms_norm_eps": 1e-5 + }"#, + ) + .unwrap(); + + let metadata = metadata_from_hf_config(&root, 3).unwrap(); + let text = format!("{metadata:?}"); + + assert!(text.contains("general.architecture")); + assert!(text.contains("deepseek2")); + assert!(text.contains("deepseek2.block_count")); + assert!(text.contains("48")); + assert!(text.contains("deepseek2.attention.key_length")); + assert!(text.contains("256")); + assert!(text.contains("deepseek2.nextn_predict_layers")); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn omits_mtp_metadata_when_requested() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + fs::write( + root.join("config.json"), + r#"{ + "model_type": "glm4_moe_lite", + "vocab_size": 154880, + "max_position_embeddings": 202752, + "hidden_size": 2048, + "intermediate_size": 10240, + "num_hidden_layers": 47, + "num_nextn_predict_layers": 1, + "num_attention_heads": 20, + "num_key_value_heads": 20, + "qk_nope_head_dim": 192, + "qk_rope_head_dim": 64, + "v_head_dim": 256 + }"#, + ) + .unwrap(); + + let metadata = + metadata_from_hf_config_with_options(&root, 3, MetadataOptions { include_mtp: false }) + .unwrap(); + + assert!(metadata.iter().any(|kv| { + matches!( + kv, + GgufKv::U32 { key, value } + if key == "deepseek2.block_count" && *value == 47 + ) + })); + assert!(!metadata.iter().any(|kv| { + matches!( + kv, + GgufKv::U32 { key, .. } if key == "deepseek2.nextn_predict_layers" + ) + })); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn builds_glm_dsa_indexer_metadata_from_config() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + fs::write( + root.join("config.json"), + r#"{ + "model_type": "glm_moe_dsa", + "vocab_size": 154880, + "max_position_embeddings": 1048576, + "hidden_size": 6144, + "intermediate_size": 12288, + "num_hidden_layers": 78, + "num_nextn_predict_layers": 1, + "num_attention_heads": 64, + "num_key_value_heads": 64, + "qk_nope_head_dim": 192, + "qk_rope_head_dim": 64, + "v_head_dim": 256, + "q_lora_rank": 2048, + "kv_lora_rank": 512, + "index_n_heads": 32, + "index_head_dim": 128, + "index_topk": 2048, + "n_routed_experts": 256, + "num_experts_per_tok": 8, + "n_shared_experts": 1, + "moe_intermediate_size": 2048, + "first_k_dense_replace": 3, + "routed_scaling_factor": 2.5, + "norm_topk_prob": true, + "rms_norm_eps": 1e-5 + }"#, + ) + .unwrap(); + + let metadata = metadata_from_hf_config(&root, 3).unwrap(); + + assert!(metadata.iter().any(|kv| { + matches!( + kv, + GgufKv::U32 { key, value } + if key == "glm-dsa.attention.indexer.head_count" && *value == 32 + ) + })); + assert!(metadata.iter().any(|kv| { + matches!( + kv, + GgufKv::U32 { key, value } + if key == "glm-dsa.attention.indexer.key_length" && *value == 128 + ) + })); + assert!(metadata.iter().any(|kv| { + matches!( + kv, + GgufKv::U32 { key, value } + if key == "glm-dsa.attention.indexer.top_k" && *value == 2048 + ) + })); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn builds_llama_metadata_from_config() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + fs::write( + root.join("config.json"), + r#"{ + "model_type": "llama", + "vocab_size": 128256, + "max_position_embeddings": 131072, + "hidden_size": 4096, + "intermediate_size": 14336, + "num_hidden_layers": 32, + "num_attention_heads": 32, + "num_key_value_heads": 8, + "head_dim": 128, + "rope_theta": 500000, + "rms_norm_eps": 1e-5 + }"#, + ) + .unwrap(); + + let metadata = metadata_from_hf_config(&root, 3).unwrap(); + let text = format!("{metadata:?}"); + + assert!(text.contains("general.architecture")); + assert!(text.contains("llama")); + assert!(text.contains("llama.block_count")); + assert!(text.contains("32")); + assert!(text.contains("llama.attention.head_count_kv")); + assert!(text.contains("8")); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn defaults_kv_heads_to_attention_heads() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + fs::write( + root.join("config.json"), + r#"{ + "model_type": "llama", + "vocab_size": 32000, + "max_position_embeddings": 4096, + "hidden_size": 4096, + "intermediate_size": 11008, + "num_hidden_layers": 32, + "num_attention_heads": 32 + }"#, + ) + .unwrap(); + + let metadata = metadata_from_hf_config(&root, 3).unwrap(); + let text = format!("{metadata:?}"); + + assert!(text.contains("llama.attention.head_count_kv")); + assert!(text.contains("32")); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn builds_qwen2_moe_metadata_from_qwen_config_keys() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + fs::write( + root.join("config.json"), + r#"{ + "model_type": "qwen2_moe", + "vocab_size": 151936, + "max_position_embeddings": 32768, + "hidden_size": 2048, + "intermediate_size": 5632, + "num_hidden_layers": 24, + "num_attention_heads": 16, + "num_key_value_heads": 16, + "head_dim": 128, + "num_experts": 60, + "num_experts_per_tok": 4, + "moe_intermediate_size": 1408, + "shared_expert_intermediate_size": 5632, + "rope_theta": 1000000, + "rms_norm_eps": 1e-6 + }"#, + ) + .unwrap(); + + let metadata = metadata_from_hf_config(&root, 3).unwrap(); + let text = format!("{metadata:?}"); + + assert!(text.contains("qwen2moe")); + assert!(text.contains("qwen2moe.expert_count")); + assert!(text.contains("60")); + assert!(text.contains("qwen2moe.expert_feed_forward_length")); + assert!(text.contains("1408")); + assert!(text.contains("qwen2moe.expert_shared_feed_forward_length")); + assert!(text.contains("5632")); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn rejects_newer_qwen_variants_that_need_specific_native_templates() { + for model_type in [ + "qwen3_next", + "qwen3next", + "qwen3_vl", + "qwen3vl_moe", + "qwen3.5", + "qwen3_5_moe", + "qwen35", + "qwen35moe", + ] { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + fs::write( + root.join("config.json"), + format!( + r#"{{ + "model_type": "{model_type}", + "vocab_size": 151936, + "max_position_embeddings": 32768, + "hidden_size": 2048, + "intermediate_size": 5632, + "num_hidden_layers": 24, + "num_attention_heads": 16 + }}"#, + ), + ) + .unwrap(); + + let err = metadata_from_hf_config(&root, 3).unwrap_err(); + let message = err.to_string(); + assert!( + message.contains("Qwen3.5/Qwen3Next/Qwen3VL-specific"), + "{model_type}: {message}" + ); + assert!( + message.contains("external convert_hf_to_gguf.py backend"), + "{model_type}: {message}" + ); + fs::remove_dir_all(root).unwrap(); + } + } + + fn unique_temp_dir() -> PathBuf { + static NEXT_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let id = NEXT_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + std::env::temp_dir().join(format!("skippy-gguf-template-{nanos}-{id}")) + } +} diff --git a/crates/skippy-quantize/src/gguf_writer.rs b/crates/skippy-quantize/src/gguf_writer.rs new file mode 100644 index 0000000000..8ca89a1342 --- /dev/null +++ b/crates/skippy-quantize/src/gguf_writer.rs @@ -0,0 +1,945 @@ +use std::collections::{BTreeMap, btree_map::Entry}; +use std::fs::{self, File}; +use std::io::{Seek, Write}; +use std::path::Path; + +use anyhow::{Context, Result, ensure}; +use serde::Serialize; + +use crate::float_convert::{FloatDType, convert_float_chunk, target_dtype_for_tensor}; +use crate::hf_checkpoint::{SafetensorFile, SafetensorTensorInfo, open_safetensor_files}; +use crate::tensor_map::{ + TensorNameMap, hf_layer_id, is_mtp_source_tensor, is_shared_mtp_context_tensor, +}; +use crate::types::ConvertOutputType; + +const GGUF_MAGIC: &[u8; 4] = b"GGUF"; +const GGUF_VERSION: u32 = 3; +const GGUF_ALIGNMENT: u64 = 32; +const GGUF_TYPE_BOOL: u32 = 7; +const GGUF_TYPE_UINT32: u32 = 4; +const GGUF_TYPE_INT32: u32 = 5; +const GGUF_TYPE_FLOAT32: u32 = 6; +const GGUF_TYPE_STRING: u32 = 8; +const GGUF_TYPE_ARRAY: u32 = 9; +const GGUF_TYPE_UINT16: u32 = 2; +const GGUF_TYPE_UINT64: u32 = 10; +const GGML_TYPE_F32: u32 = 0; +const GGML_TYPE_F16: u32 = 1; +const GGML_TYPE_BF16: u32 = 30; + +#[derive(Debug, Clone)] +pub(crate) struct RawGgufWriteOptions { + pub(crate) buffer_size: usize, + pub(crate) metadata: Option>, + pub(crate) tensor_name_map: TensorNameMap, + pub(crate) split: Option, + pub(crate) output_type: Option, + pub(crate) tensor_selection: TensorSelection, +} + +#[derive(Debug, Clone, Copy, Default)] +pub(crate) enum TensorSelection { + #[default] + All, + ExcludeMtp { + layer_start: u32, + }, + MtpOnly { + layer_start: u32, + }, +} + +#[derive(Debug, Clone, Copy)] +pub(crate) struct GgufSplit { + pub(crate) split_index: u32, + pub(crate) split_count: u32, +} + +pub(crate) fn write_raw_safetensors_gguf( + source: &Path, + output: &Path, + options: RawGgufWriteOptions, +) -> Result<()> { + let PreparedGgufWrite { + files, + tensors, + metadata, + } = prepare_raw_safetensors_gguf(source, &options)?; + if let Some(parent) = output.parent() { + fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?; + } + let mut writer = + File::create(output).with_context(|| format!("create {}", output.display()))?; + write_header_and_tensor_table(&mut writer, &metadata, &tensors)?; + stream_tensor_data(&mut writer, &files, &tensors, options.buffer_size) +} + +pub(crate) fn validate_raw_safetensors_gguf( + source: &Path, + options: RawGgufWriteOptions, +) -> Result { + let PreparedGgufWrite { + tensors, metadata, .. + } = prepare_raw_safetensors_gguf(source, &options)?; + Ok(RawGgufValidation { + selected_tensor_count: tensors.len(), + selected_tensor_bytes: tensors.iter().map(|tensor| tensor.byte_len).sum(), + metadata_count: metadata.len(), + output_type: options.output_type.map(|kind| kind.as_arg().to_string()), + }) +} + +struct PreparedGgufWrite { + files: Vec, + tensors: Vec, + metadata: Vec, +} + +#[derive(Debug, Serialize)] +pub(crate) struct RawGgufValidation { + pub(crate) selected_tensor_count: usize, + pub(crate) selected_tensor_bytes: u64, + pub(crate) metadata_count: usize, + #[serde(skip_serializing_if = "Option::is_none")] + pub(crate) output_type: Option, +} + +fn prepare_raw_safetensors_gguf( + source: &Path, + options: &RawGgufWriteOptions, +) -> Result { + ensure!( + options.buffer_size > 0, + "buffer_size must be greater than zero" + ); + let files = open_safetensor_files(source)?; + ensure!( + !files.is_empty(), + "no safetensors files found under {}", + source.display() + ); + let tensors = collect_tensor_sources( + &files, + options.tensor_name_map, + options.output_type, + options.tensor_selection, + )?; + ensure!( + !tensors.is_empty(), + "no tensors found under {}", + source.display() + ); + let total_tensor_count = tensors.len(); + let mut tensors = select_split_tensors(tensors, options.split)?; + assign_gguf_offsets(&mut tensors)?; + let metadata = options + .metadata + .clone() + .unwrap_or_else(|| raw_metadata(source, total_tensor_count)); + let metadata = split_metadata(metadata, options.split, total_tensor_count)?; + Ok(PreparedGgufWrite { + files, + tensors, + metadata, + }) +} + +fn select_split_tensors( + tensors: Vec, + split: Option, +) -> Result> { + let Some(split) = split else { + return Ok(tensors); + }; + split.validate()?; + let total_tensors = tensors.len(); + ensure!( + usize::try_from(split.split_count).is_ok_and(|count| count <= total_tensors), + "split_count {} cannot exceed tensor count {}", + split.split_count, + total_tensors + ); + let split_index = + usize::try_from(split.split_index).context("split_index does not fit usize")?; + let boundaries = byte_balanced_split_boundaries(&tensors, split)?; + let start = boundaries[split_index - 1]; + let end = boundaries[split_index]; + ensure!( + start < end, + "split {} of {} would contain no tensors", + split.split_index, + split.split_count + ); + Ok(tensors + .into_iter() + .enumerate() + .filter_map(|(index, tensor)| (start <= index && index < end).then_some(tensor)) + .collect()) +} + +fn byte_balanced_split_boundaries( + tensors: &[TensorSource], + split: GgufSplit, +) -> Result> { + split.validate()?; + let split_count = + usize::try_from(split.split_count).context("split_count does not fit usize")?; + ensure!( + split_count <= tensors.len(), + "split_count {} cannot exceed tensor count {}", + split.split_count, + tensors.len() + ); + let total_bytes = tensors + .iter() + .try_fold(0_u128, |acc, tensor| { + acc.checked_add(tensor.byte_len as u128) + }) + .context("split tensor byte total overflow")?; + let mut boundaries = vec![0_usize]; + let mut accumulated = 0_u128; + for (index, tensor) in tensors.iter().enumerate() { + accumulated = accumulated + .checked_add(tensor.byte_len as u128) + .context("split tensor byte total overflow")?; + let remaining_tensors = tensors.len() - (index + 1); + let remaining_splits = split_count - boundaries.len(); + if boundaries.len() < split_count && remaining_tensors >= remaining_splits { + let target = total_bytes + .checked_mul(boundaries.len() as u128) + .context("split target byte overflow")? + / split_count as u128; + if accumulated >= target { + boundaries.push(index + 1); + } + } + } + while boundaries.len() < split_count { + let next = boundaries.last().copied().unwrap_or(0) + 1; + boundaries.push(next); + } + boundaries.push(tensors.len()); + Ok(boundaries) +} + +fn assign_gguf_offsets(tensors: &mut [TensorSource]) -> Result<()> { + let mut offset = 0_u64; + for tensor in tensors { + offset = align_to(offset, GGUF_ALIGNMENT); + tensor.gguf_offset = offset; + offset = offset + .checked_add(tensor.byte_len) + .with_context(|| format!("GGUF data offset overflow after {}", tensor.name))?; + } + Ok(()) +} + +fn split_metadata( + mut metadata: Vec, + split: Option, + total_tensor_count: usize, +) -> Result> { + let Some(split) = split else { + return Ok(metadata); + }; + split.validate()?; + metadata.push(GgufKv::u16( + "split.no", + u16::try_from(split.split_index - 1).context("split index does not fit uint16")?, + )); + metadata.push(GgufKv::u16( + "split.count", + u16::try_from(split.split_count).context("split count does not fit uint16")?, + )); + metadata.push(GgufKv::i32( + "split.tensors.count", + i32::try_from(total_tensor_count).context("tensor count does not fit int32")?, + )); + Ok(metadata) +} + +impl GgufSplit { + fn validate(self) -> Result<()> { + ensure!( + self.split_count > 0, + "split_count must be greater than zero" + ); + ensure!( + self.split_index > 0, + "split_index is 1-based and cannot be zero" + ); + ensure!( + self.split_index <= self.split_count, + "split_index {} exceeds split_count {}", + self.split_index, + self.split_count + ); + ensure!( + u16::try_from(self.split_count).is_ok(), + "split_count {} exceeds GGUF uint16 split metadata", + self.split_count + ); + Ok(()) + } +} + +fn collect_tensor_sources( + files: &[SafetensorFile], + tensor_name_map: TensorNameMap, + output_type: Option, + tensor_selection: TensorSelection, +) -> Result> { + let mut tensors = Vec::new(); + let mut expert_groups = BTreeMap::::new(); + for (file_index, file) in files.iter().enumerate() { + for tensor in file.tensors().values() { + if !tensor_selection.includes(tensor.name())? { + continue; + } + if matches!( + tensor_name_map, + TensorNameMap::HfToGguf | TensorNameMap::HfToGgufWithMtp { .. } + ) && let Some(expert) = ExpertSourceTensor::parse(tensor.name())? + { + match expert_groups.entry(expert.group_key()) { + Entry::Vacant(entry) => { + entry + .insert(ExpertGroup::new(expert.group_key(), tensor, output_type)?) + .push(file_index, tensor, expert.expert_id)?; + } + Entry::Occupied(mut entry) => { + entry.get_mut().push(file_index, tensor, expert.expert_id)?; + } + } + continue; + } + tensors.push(TensorSource::from_safetensor( + file_index, + tensor, + tensor_name_map, + output_type, + )?); + } + } + for group in expert_groups.into_values() { + tensors.push(group.into_tensor_source()?); + } + tensors.sort_by(|a, b| a.name.cmp(&b.name)); + Ok(tensors) +} + +impl TensorSelection { + fn includes(self, name: &str) -> Result { + let is_mtp = match self { + Self::All => return Ok(true), + Self::ExcludeMtp { layer_start } | Self::MtpOnly { layer_start } => { + is_mtp_source_tensor(name) + || hf_layer_id(name)?.is_some_and(|layer| layer >= layer_start) + } + }; + match self { + Self::All => Ok(true), + Self::ExcludeMtp { .. } => Ok(!is_mtp), + Self::MtpOnly { .. } => Ok(is_mtp || is_shared_mtp_context_tensor(name)), + } + } +} + +struct TensorSource { + segments: Vec, + name: String, + dims: Vec, + ggml_type: u32, + byte_len: u64, + gguf_offset: u64, +} + +impl TensorSource { + fn from_safetensor( + file_index: usize, + tensor: &SafetensorTensorInfo, + tensor_name_map: TensorNameMap, + output_type: Option, + ) -> Result { + let source_dtype = FloatDType::from_safetensor(tensor.dtype()).with_context(|| { + format!("unsupported dtype {} for {}", tensor.dtype(), tensor.name()) + })?; + let target_dtype = target_dtype_for_tensor(source_dtype, output_type, tensor.shape())?; + let name = tensor_name_map.map_tensor_name(tensor.name())?; + let element_count = tensor_element_count(tensor)?; + Ok(Self { + segments: vec![TensorSegment { + file_index, + source_name: tensor.name().to_string(), + source_dtype, + target_dtype, + element_count, + source_byte_len: tensor.byte_len(), + target_byte_len: tensor_byte_len(element_count, target_dtype)?, + }], + name, + dims: tensor.shape().iter().rev().copied().collect(), + ggml_type: ggml_type_for_dtype(target_dtype), + byte_len: tensor_byte_len(element_count, target_dtype)?, + gguf_offset: 0, + }) + } +} + +struct TensorSegment { + file_index: usize, + source_name: String, + source_dtype: FloatDType, + target_dtype: FloatDType, + element_count: u64, + source_byte_len: u64, + target_byte_len: u64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +struct ExpertGroupKey { + layer: u32, + projection: ExpertProjection, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +enum ExpertProjection { + Down, + Gate, + Up, +} + +impl ExpertProjection { + fn gguf_name(self, layer: u32) -> String { + match self { + Self::Down => format!("blk.{layer}.ffn_down_exps.weight"), + Self::Gate => format!("blk.{layer}.ffn_gate_exps.weight"), + Self::Up => format!("blk.{layer}.ffn_up_exps.weight"), + } + } +} + +#[derive(Debug, Clone, Copy)] +struct ExpertSourceTensor { + layer: u32, + expert_id: u32, + projection: ExpertProjection, +} + +impl ExpertSourceTensor { + fn parse(name: &str) -> Result> { + let Some(rest) = name.strip_prefix("model.layers.") else { + return Ok(None); + }; + let Some((layer, suffix)) = rest.split_once('.') else { + return Ok(None); + }; + let Some(expert_suffix) = suffix.strip_prefix("mlp.experts.") else { + return Ok(None); + }; + let Some((expert_id, projection_suffix)) = expert_suffix.split_once('.') else { + return Ok(None); + }; + let layer = layer + .parse::() + .with_context(|| format!("parse expert layer id in {name}"))?; + let expert_id = expert_id + .parse::() + .with_context(|| format!("parse expert id in {name}"))?; + let projection = match projection_suffix { + "down_proj.weight" => ExpertProjection::Down, + "gate_proj.weight" => ExpertProjection::Gate, + "up_proj.weight" => ExpertProjection::Up, + _ => return Ok(None), + }; + Ok(Some(Self { + layer, + expert_id, + projection, + })) + } + + fn group_key(self) -> ExpertGroupKey { + ExpertGroupKey { + layer: self.layer, + projection: self.projection, + } + } +} + +struct ExpertGroup { + key: ExpertGroupKey, + source_dtype: FloatDType, + target_dtype: FloatDType, + shape: Vec, + source_byte_len_per_expert: u64, + target_byte_len_per_expert: u64, + experts: BTreeMap, +} + +impl ExpertGroup { + fn new( + key: ExpertGroupKey, + tensor: &SafetensorTensorInfo, + output_type: Option, + ) -> Result { + let source_dtype = FloatDType::from_safetensor(tensor.dtype()).with_context(|| { + format!("unsupported dtype {} for {}", tensor.dtype(), tensor.name()) + })?; + let target_dtype = target_dtype_for_tensor(source_dtype, output_type, tensor.shape())?; + let element_count = tensor_element_count(tensor)?; + Ok(Self { + key, + source_dtype, + target_dtype, + shape: tensor.shape().to_vec(), + source_byte_len_per_expert: tensor.byte_len(), + target_byte_len_per_expert: tensor_byte_len(element_count, target_dtype)?, + experts: BTreeMap::new(), + }) + } + + fn push( + &mut self, + file_index: usize, + tensor: &SafetensorTensorInfo, + expert_id: u32, + ) -> Result<()> { + ensure!( + FloatDType::from_safetensor(tensor.dtype()) == Some(self.source_dtype), + "expert tensor {} dtype {} does not match group dtype {:?}", + tensor.name(), + tensor.dtype(), + self.source_dtype + ); + ensure!( + tensor.shape() == self.shape, + "expert tensor {} shape {:?} does not match group shape {:?}", + tensor.name(), + tensor.shape(), + self.shape + ); + ensure!( + tensor.byte_len() == self.source_byte_len_per_expert, + "expert tensor {} byte length {} does not match group byte length {}", + tensor.name(), + tensor.byte_len(), + self.source_byte_len_per_expert + ); + let element_count = tensor_element_count(tensor)?; + let previous = self.experts.insert( + expert_id, + TensorSegment { + file_index, + source_name: tensor.name().to_string(), + source_dtype: self.source_dtype, + target_dtype: self.target_dtype, + element_count, + source_byte_len: tensor.byte_len(), + target_byte_len: tensor_byte_len(element_count, self.target_dtype)?, + }, + ); + ensure!( + previous.is_none(), + "duplicate expert tensor id {expert_id} for {}", + self.key.projection.gguf_name(self.key.layer) + ); + Ok(()) + } + + fn into_tensor_source(self) -> Result { + ensure!( + !self.experts.is_empty(), + "expert group {} has no tensors", + self.key.projection.gguf_name(self.key.layer) + ); + for (expected, actual) in self.experts.keys().copied().enumerate() { + ensure!( + expected as u32 == actual, + "expert group {} is missing expert id {}", + self.key.projection.gguf_name(self.key.layer), + expected + ); + } + let expert_count = self.experts.len() as u64; + let mut dims = self.shape.iter().rev().copied().collect::>(); + dims.push(expert_count); + let byte_len = self + .target_byte_len_per_expert + .checked_mul(expert_count) + .with_context(|| { + format!( + "expert group {} byte length overflow", + self.key.projection.gguf_name(self.key.layer) + ) + })?; + Ok(TensorSource { + segments: self.experts.into_values().collect(), + name: self.key.projection.gguf_name(self.key.layer), + dims, + ggml_type: ggml_type_for_dtype(self.target_dtype), + byte_len, + gguf_offset: 0, + }) + } +} + +fn tensor_element_count(tensor: &SafetensorTensorInfo) -> Result { + tensor.shape().iter().try_fold(1_u64, |acc, dim| { + acc.checked_mul(*dim) + .with_context(|| format!("tensor {} element count overflow", tensor.name())) + }) +} + +fn tensor_byte_len(element_count: u64, dtype: FloatDType) -> Result { + element_count + .checked_mul(dtype.byte_size()) + .context("target tensor byte length overflow") +} + +fn ggml_type_for_dtype(dtype: FloatDType) -> u32 { + match dtype { + FloatDType::F32 => GGML_TYPE_F32, + FloatDType::F16 => GGML_TYPE_F16, + FloatDType::Bf16 => GGML_TYPE_BF16, + } +} + +fn raw_metadata(source: &Path, tensor_count: usize) -> Vec { + vec![ + GgufKv::string("general.architecture", "raw-safetensors"), + GgufKv::string( + "general.name", + source + .file_name() + .and_then(|value| value.to_str()) + .unwrap_or("checkpoint"), + ), + GgufKv::bool("skippy.convert.raw_safetensors", true), + GgufKv::u64("skippy.convert.tensor_count", tensor_count as u64), + ] +} + +#[derive(Debug, Clone)] +pub(crate) enum GgufKv { + ArrayF32 { key: String, value: Vec }, + ArrayI32 { key: String, value: Vec }, + ArrayString { key: String, value: Vec }, + Bool { key: String, value: bool }, + F32 { key: String, value: f32 }, + I32 { key: String, value: i32 }, + String { key: String, value: String }, + U16 { key: String, value: u16 }, + U32 { key: String, value: u32 }, + U64 { key: String, value: u64 }, +} + +impl GgufKv { + pub(crate) fn array_f32(key: &str, value: Vec) -> Self { + Self::ArrayF32 { + key: key.to_string(), + value, + } + } + + pub(crate) fn array_i32(key: &str, value: Vec) -> Self { + Self::ArrayI32 { + key: key.to_string(), + value, + } + } + + pub(crate) fn array_string(key: &str, value: Vec) -> Self { + Self::ArrayString { + key: key.to_string(), + value, + } + } + + pub(crate) fn bool(key: &str, value: bool) -> Self { + Self::Bool { + key: key.to_string(), + value, + } + } + + pub(crate) fn f32(key: &str, value: f32) -> Self { + Self::F32 { + key: key.to_string(), + value, + } + } + + pub(crate) fn i32(key: &str, value: i32) -> Self { + Self::I32 { + key: key.to_string(), + value, + } + } + + pub(crate) fn string(key: &str, value: &str) -> Self { + Self::String { + key: key.to_string(), + value: value.to_string(), + } + } + + pub(crate) fn u16(key: &str, value: u16) -> Self { + Self::U16 { + key: key.to_string(), + value, + } + } + + pub(crate) fn u32(key: &str, value: u32) -> Self { + Self::U32 { + key: key.to_string(), + value, + } + } + + pub(crate) fn u64(key: &str, value: u64) -> Self { + Self::U64 { + key: key.to_string(), + value, + } + } +} + +fn write_header_and_tensor_table( + writer: &mut W, + metadata: &[GgufKv], + tensors: &[TensorSource], +) -> Result<()> { + writer.write_all(GGUF_MAGIC)?; + write_u32(writer, GGUF_VERSION)?; + write_u64(writer, tensors.len() as u64)?; + write_u64(writer, metadata.len() as u64)?; + for kv in metadata { + write_kv(writer, kv)?; + } + for tensor in tensors { + write_string(writer, &tensor.name)?; + write_u32(writer, tensor.dims.len() as u32)?; + for dim in &tensor.dims { + write_u64(writer, *dim)?; + } + write_u32(writer, tensor.ggml_type)?; + write_u64(writer, tensor.gguf_offset)?; + } + Ok(()) +} + +fn write_kv(writer: &mut W, kv: &GgufKv) -> Result<()> { + match kv { + GgufKv::ArrayF32 { key, value } => { + write_array_header(writer, key, GGUF_TYPE_FLOAT32, value.len())?; + for item in value { + writer.write_all(&item.to_le_bytes())?; + } + } + GgufKv::ArrayI32 { key, value } => { + write_array_header(writer, key, GGUF_TYPE_INT32, value.len())?; + for item in value { + writer.write_all(&item.to_le_bytes())?; + } + } + GgufKv::ArrayString { key, value } => { + write_array_header(writer, key, GGUF_TYPE_STRING, value.len())?; + for item in value { + write_string(writer, item)?; + } + } + GgufKv::Bool { key, value } => { + write_string(writer, key)?; + write_u32(writer, GGUF_TYPE_BOOL)?; + writer.write_all(&[*value as u8])?; + } + GgufKv::F32 { key, value } => { + write_string(writer, key)?; + write_u32(writer, GGUF_TYPE_FLOAT32)?; + writer.write_all(&value.to_le_bytes())?; + } + GgufKv::I32 { key, value } => { + write_string(writer, key)?; + write_u32(writer, GGUF_TYPE_INT32)?; + writer.write_all(&value.to_le_bytes())?; + } + GgufKv::String { key, value } => { + write_string(writer, key)?; + write_u32(writer, GGUF_TYPE_STRING)?; + write_string(writer, value)?; + } + GgufKv::U16 { key, value } => { + write_string(writer, key)?; + write_u32(writer, GGUF_TYPE_UINT16)?; + writer.write_all(&value.to_le_bytes())?; + } + GgufKv::U32 { key, value } => { + write_string(writer, key)?; + write_u32(writer, GGUF_TYPE_UINT32)?; + write_u32(writer, *value)?; + } + GgufKv::U64 { key, value } => { + write_string(writer, key)?; + write_u32(writer, GGUF_TYPE_UINT64)?; + write_u64(writer, *value)?; + } + } + Ok(()) +} + +fn write_array_header( + writer: &mut W, + key: &str, + element_type: u32, + len: usize, +) -> Result<()> { + ensure!(!key.is_empty(), "GGUF metadata key cannot be empty"); + ensure!( + len > 0, + "GGUF array metadata {key:?} cannot be empty because llama.cpp rejects empty arrays" + ); + write_string(writer, key)?; + write_u32(writer, GGUF_TYPE_ARRAY)?; + write_u32(writer, element_type)?; + write_u64(writer, len as u64) +} + +fn stream_tensor_data( + writer: &mut File, + files: &[SafetensorFile], + tensors: &[TensorSource], + buffer_size: usize, +) -> Result<()> { + pad_writer_to_alignment(writer, GGUF_ALIGNMENT)?; + let data_start = writer.stream_position()?; + for tensor in tensors { + let expected_position = data_start + tensor.gguf_offset; + pad_writer_to_position(writer, expected_position)?; + let mut copied = 0_u64; + for segment in &tensor.segments { + let segment_copied = + stream_segment(writer, &files[segment.file_index], segment, buffer_size)?; + ensure!( + segment_copied == segment.target_byte_len, + "copied {} bytes for {}, expected {}", + segment_copied, + segment.source_name, + segment.target_byte_len + ); + copied += segment_copied; + } + ensure!( + copied == tensor.byte_len, + "copied {} bytes for {}, expected {}", + copied, + tensor.name, + tensor.byte_len + ); + } + Ok(()) +} + +fn stream_segment( + writer: &mut File, + file: &SafetensorFile, + segment: &TensorSegment, + buffer_size: usize, +) -> Result { + if segment.source_dtype == segment.target_dtype { + let copied = file.stream_tensor(&segment.source_name, writer, buffer_size)?; + ensure!( + copied == segment.source_byte_len, + "read {} bytes for {}, expected {}", + copied, + segment.source_name, + segment.source_byte_len + ); + return Ok(copied); + } + + let source_element_size = usize::try_from(segment.source_dtype.byte_size()) + .context("source dtype byte size does not fit usize")?; + let chunk_size = aligned_chunk_size(buffer_size, source_element_size); + let mut output_bytes = 0_u64; + let mut source_bytes = 0_u64; + file.stream_tensor_chunks(&segment.source_name, chunk_size, |chunk| { + ensure!( + chunk.len() % source_element_size == 0, + "chunk for {} split an element boundary", + segment.source_name + ); + source_bytes += chunk.len() as u64; + output_bytes += + convert_float_chunk(chunk, segment.source_dtype, segment.target_dtype, writer)?; + Ok(()) + })?; + ensure!( + source_bytes == segment.source_byte_len, + "read {} bytes for {}, expected {}", + source_bytes, + segment.source_name, + segment.source_byte_len + ); + ensure!( + source_bytes / segment.source_dtype.byte_size() == segment.element_count, + "read element count mismatch for {}", + segment.source_name + ); + Ok(output_bytes) +} + +fn aligned_chunk_size(buffer_size: usize, element_size: usize) -> usize { + let aligned = buffer_size - (buffer_size % element_size); + aligned.max(element_size) +} + +fn pad_writer_to_alignment(writer: &mut File, alignment: u64) -> Result<()> { + let position = writer.stream_position()?; + pad_writer_to_position(writer, align_to(position, alignment)) +} + +fn pad_writer_to_position(writer: &mut File, position: u64) -> Result<()> { + let current = writer.stream_position()?; + ensure!( + current <= position, + "writer is past requested output position {position}" + ); + let mut remaining = position - current; + let zeros = [0_u8; 4096]; + while remaining > 0 { + let write_len = zeros.len().min(remaining as usize); + writer.write_all(&zeros[..write_len])?; + remaining -= write_len as u64; + } + Ok(()) +} + +fn align_to(value: u64, alignment: u64) -> u64 { + if alignment <= 1 { + return value; + } + value.div_ceil(alignment) * alignment +} + +fn write_string(writer: &mut W, value: &str) -> Result<()> { + write_u64(writer, value.len() as u64)?; + writer.write_all(value.as_bytes())?; + Ok(()) +} + +fn write_u32(writer: &mut W, value: u32) -> Result<()> { + writer.write_all(&value.to_le_bytes())?; + Ok(()) +} + +fn write_u64(writer: &mut W, value: u64) -> Result<()> { + writer.write_all(&value.to_le_bytes())?; + Ok(()) +} + +#[cfg(test)] +#[path = "gguf_writer_tests.rs"] +mod tests; diff --git a/crates/skippy-quantize/src/gguf_writer_tests.rs b/crates/skippy-quantize/src/gguf_writer_tests.rs new file mode 100644 index 0000000000..8244f301bc --- /dev/null +++ b/crates/skippy-quantize/src/gguf_writer_tests.rs @@ -0,0 +1,1208 @@ +use std::io::Read; +use std::path::PathBuf; + +use crate::gguf_template::metadata_from_hf_config; +use crate::tensor_map::TensorNameMap; + +use super::*; + +#[test] +fn writes_raw_gguf_from_safetensors_with_streamed_payloads() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + write_safetensor( + &root.join("model.safetensors"), + &[ + ("b.weight", "BF16", &[2], &[9, 8, 7, 6]), + ("a.weight", "F32", &[1], &[1, 2, 3, 4]), + ], + ); + let output = root.join("raw.gguf"); + + write_raw_safetensors_gguf( + &root, + &output, + RawGgufWriteOptions { + buffer_size: 2, + metadata: None, + tensor_name_map: TensorNameMap::Raw, + split: None, + output_type: None, + tensor_selection: TensorSelection::All, + }, + ) + .unwrap(); + + let bytes = fs::read(&output).unwrap(); + assert_eq!(&bytes[..4], GGUF_MAGIC); + let parsed = parse_test_gguf(&bytes); + assert_eq!(parsed.tensor_count, 2); + assert_eq!(parsed.metadata_count, 4); + assert_eq!(parsed.tensors[0].name, "a.weight"); + assert_eq!(parsed.tensors[0].ggml_type, GGML_TYPE_F32); + assert_eq!( + &bytes[parsed.tensors[0].absolute_offset..parsed.tensors[0].absolute_offset + 4], + &[1, 2, 3, 4] + ); + assert_eq!(parsed.tensors[1].name, "b.weight"); + assert_eq!(parsed.tensors[1].ggml_type, GGML_TYPE_BF16); + assert_eq!( + &bytes[parsed.tensors[1].absolute_offset..parsed.tensors[1].absolute_offset + 4], + &[9, 8, 7, 6] + ); + fs::remove_dir_all(root).unwrap(); +} + +#[test] +fn writes_mapped_hf_tensor_names_when_requested() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + write_safetensor( + &root.join("model.safetensors"), + &[( + "model.layers.0.input_layernorm.weight", + "F32", + &[1], + &[1, 2, 3, 4], + )], + ); + let output = root.join("mapped.gguf"); + + write_raw_safetensors_gguf( + &root, + &output, + RawGgufWriteOptions { + buffer_size: 2, + metadata: None, + tensor_name_map: TensorNameMap::HfToGguf, + split: None, + output_type: None, + tensor_selection: TensorSelection::All, + }, + ) + .unwrap(); + + let bytes = fs::read(&output).unwrap(); + let parsed = parse_test_gguf(&bytes); + assert_eq!(parsed.tensors[0].name, "blk.0.attn_norm.weight"); + fs::remove_dir_all(root).unwrap(); +} + +#[test] +fn excludes_mtp_source_tensors_before_hf_name_mapping() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + write_safetensor( + &root.join("model.safetensors"), + &[ + ( + "model.layers.0.input_layernorm.weight", + "F32", + &[1], + &[1, 2, 3, 4], + ), + ( + "model.layers.1.input_layernorm.weight", + "F32", + &[1], + &[5, 6, 7, 8], + ), + ( + "model.layers.1.eh_proj.weight", + "F32", + &[1], + &[9, 10, 11, 12], + ), + ("mtp.fc.weight", "F32", &[1], &[13, 14, 15, 16]), + ], + ); + let output = root.join("no-mtp.gguf"); + + write_raw_safetensors_gguf( + &root, + &output, + RawGgufWriteOptions { + buffer_size: 2, + metadata: None, + tensor_name_map: TensorNameMap::HfToGguf, + split: None, + output_type: None, + tensor_selection: TensorSelection::ExcludeMtp { layer_start: 1 }, + }, + ) + .unwrap(); + + let bytes = fs::read(&output).unwrap(); + let parsed = parse_test_gguf(&bytes); + assert_eq!(parsed.tensor_count, 1); + assert_eq!(parsed.tensors[0].name, "blk.0.attn_norm.weight"); + fs::remove_dir_all(root).unwrap(); +} + +#[test] +fn writes_mtp_only_tensors_with_shared_context() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + write_safetensor( + &root.join("model.safetensors"), + &[ + ("model.embed_tokens.weight", "F32", &[1], &[1, 0, 0, 0]), + ("lm_head.weight", "F32", &[1], &[2, 0, 0, 0]), + ( + "model.layers.0.input_layernorm.weight", + "F32", + &[1], + &[3, 0, 0, 0], + ), + ( + "model.layers.1.input_layernorm.weight", + "F32", + &[1], + &[4, 0, 0, 0], + ), + ("model.layers.1.eh_proj.weight", "F32", &[1], &[5, 0, 0, 0]), + ], + ); + let output = root.join("mtp-only.gguf"); + + write_raw_safetensors_gguf( + &root, + &output, + RawGgufWriteOptions { + buffer_size: 4, + metadata: None, + tensor_name_map: TensorNameMap::HfToGguf, + split: None, + output_type: None, + tensor_selection: TensorSelection::MtpOnly { layer_start: 1 }, + }, + ) + .unwrap(); + + let bytes = fs::read(&output).unwrap(); + let parsed = parse_test_gguf(&bytes); + let names = parsed + .tensors + .iter() + .map(|tensor| tensor.name.as_str()) + .collect::>(); + assert_eq!( + names, + vec![ + "blk.1.attn_norm.weight", + "blk.1.nextn.eh_proj.weight", + "output.weight", + "token_embd.weight", + ] + ); + fs::remove_dir_all(root).unwrap(); +} + +#[test] +fn writes_qwen_style_mtp_only_tensors_with_shared_context() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + write_safetensor( + &root.join("model.safetensors"), + &[ + ("embed_tokens.weight", "F32", &[1], &[1, 0, 0, 0]), + ("norm.weight", "F32", &[1], &[2, 0, 0, 0]), + ( + "model.layers.0.input_layernorm.weight", + "F32", + &[1], + &[3, 0, 0, 0], + ), + ("mtp.fc.weight", "F32", &[1], &[4, 0, 0, 0]), + ("model.mtp.norm.weight", "F32", &[1], &[5, 0, 0, 0]), + ( + "mtp.layers.1.self_attn.q_proj.weight", + "F32", + &[1], + &[6, 0, 0, 0], + ), + ], + ); + let output = root.join("qwen-mtp-only.gguf"); + + write_raw_safetensors_gguf( + &root, + &output, + RawGgufWriteOptions { + buffer_size: 4, + metadata: None, + tensor_name_map: TensorNameMap::HfToGgufWithMtp { layer_start: 32 }, + split: None, + output_type: None, + tensor_selection: TensorSelection::MtpOnly { layer_start: 32 }, + }, + ) + .unwrap(); + + let bytes = fs::read(&output).unwrap(); + let parsed = parse_test_gguf(&bytes); + let names = parsed + .tensors + .iter() + .map(|tensor| tensor.name.as_str()) + .collect::>(); + assert_eq!( + names, + vec![ + "blk.32.nextn.eh_proj.weight", + "blk.32.nextn.shared_head_norm.weight", + "blk.33.attn_q.weight", + "output_norm.weight", + "token_embd.weight", + ] + ); + fs::remove_dir_all(root).unwrap(); +} + +#[test] +fn validates_qwen_dense_native_conversion_fixture() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + write_qwen_config_and_tokenizer(&root); + write_safetensor( + &root.join("model.safetensors"), + &[ + ("model.embed_tokens.weight", "F32", &[1], &[1, 0, 0, 0]), + ( + "model.layers.0.input_layernorm.weight", + "F32", + &[1], + &[2, 0, 0, 0], + ), + ( + "model.layers.0.self_attn.q_proj.weight", + "F32", + &[1], + &[3, 0, 0, 0], + ), + ( + "model.layers.0.self_attn.k_proj.weight", + "F32", + &[1], + &[4, 0, 0, 0], + ), + ( + "model.layers.0.self_attn.v_proj.weight", + "F32", + &[1], + &[5, 0, 0, 0], + ), + ( + "model.layers.0.self_attn.o_proj.weight", + "F32", + &[1], + &[6, 0, 0, 0], + ), + ( + "model.layers.0.self_attn.q_norm.weight", + "F32", + &[1], + &[7, 0, 0, 0], + ), + ( + "model.layers.0.self_attn.k_norm.weight", + "F32", + &[1], + &[8, 0, 0, 0], + ), + ( + "model.layers.0.post_attention_layernorm.weight", + "F32", + &[1], + &[9, 0, 0, 0], + ), + ( + "model.layers.0.mlp.gate_proj.weight", + "F32", + &[1], + &[10, 0, 0, 0], + ), + ( + "model.layers.0.mlp.up_proj.weight", + "F32", + &[1], + &[11, 0, 0, 0], + ), + ( + "model.layers.0.mlp.down_proj.weight", + "F32", + &[1], + &[12, 0, 0, 0], + ), + ("model.norm.weight", "F32", &[1], &[13, 0, 0, 0]), + ("lm_head.weight", "F32", &[1], &[14, 0, 0, 0]), + ], + ); + let metadata = metadata_from_hf_config(&root, 14).unwrap(); + let validation = validate_raw_safetensors_gguf( + &root, + RawGgufWriteOptions { + buffer_size: 4, + metadata: Some(metadata.clone()), + tensor_name_map: TensorNameMap::HfToGguf, + split: None, + output_type: Some(ConvertOutputType::Bf16), + tensor_selection: TensorSelection::All, + }, + ) + .unwrap(); + assert_eq!(validation.selected_tensor_count, 14); + + let output = root.join("qwen-native.gguf"); + write_raw_safetensors_gguf( + &root, + &output, + RawGgufWriteOptions { + buffer_size: 4, + metadata: Some(metadata), + tensor_name_map: TensorNameMap::HfToGguf, + split: None, + output_type: Some(ConvertOutputType::Bf16), + tensor_selection: TensorSelection::All, + }, + ) + .unwrap(); + let bytes = fs::read(&output).unwrap(); + let parsed = parse_test_gguf(&bytes); + assert!(parsed.metadata_count > 10); + let attn_k = parsed.tensor("blk.0.attn_k.weight"); + assert_eq!(attn_k.ggml_type, GGML_TYPE_F32); + fs::remove_dir_all(root).unwrap(); +} + +#[test] +fn writes_glm_dsa_indexer_tensors_with_hf_name_mapping() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + write_safetensor( + &root.join("model.safetensors"), + &[ + ( + "model.layers.0.self_attn.indexer.k_norm.weight", + "F32", + &[1], + &[1, 0, 0, 0], + ), + ( + "model.layers.0.self_attn.indexer.k_norm.bias", + "F32", + &[1], + &[2, 0, 0, 0], + ), + ( + "model.layers.0.self_attn.indexer.weights_proj.weight", + "F32", + &[1], + &[3, 0, 0, 0], + ), + ( + "model.layers.0.self_attn.indexer.wk.weight", + "F32", + &[1], + &[4, 0, 0, 0], + ), + ( + "model.layers.0.self_attn.indexer.wq_b.weight", + "F32", + &[1], + &[5, 0, 0, 0], + ), + ], + ); + + let output = root.join("glm-dsa-indexer.gguf"); + write_raw_safetensors_gguf( + &root, + &output, + RawGgufWriteOptions { + buffer_size: 4, + metadata: None, + tensor_name_map: TensorNameMap::HfToGguf, + split: None, + output_type: Some(ConvertOutputType::Bf16), + tensor_selection: TensorSelection::All, + }, + ) + .unwrap(); + + let bytes = fs::read(&output).unwrap(); + let parsed = parse_test_gguf(&bytes); + assert_eq!(parsed.tensor_count, 5); + parsed.tensor("blk.0.indexer.k_norm.weight"); + parsed.tensor("blk.0.indexer.k_norm.bias"); + parsed.tensor("blk.0.indexer.proj.weight"); + parsed.tensor("blk.0.indexer.attn_k.weight"); + parsed.tensor("blk.0.indexer.attn_q_b.weight"); + fs::remove_dir_all(root).unwrap(); +} + +#[test] +fn validates_qwen2_moe_native_conversion_fixture() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + write_qwen2_moe_config_and_tokenizer(&root); + write_safetensor( + &root.join("model.safetensors"), + &[ + ("model.embed_tokens.weight", "F32", &[1], &[1, 0, 0, 0]), + ( + "model.layers.0.mlp.shared_expert_gate", + "F32", + &[1], + &[2, 0, 0, 0], + ), + ( + "model.layers.0.mlp.shared_expert.gate_proj.weight", + "F32", + &[1], + &[3, 0, 0, 0], + ), + ( + "model.layers.0.mlp.shared_expert.down_proj.weight", + "F32", + &[1], + &[4, 0, 0, 0], + ), + ( + "model.layers.0.mlp.shared_expert.up_proj.weight", + "F32", + &[1], + &[5, 0, 0, 0], + ), + ( + "model.layers.0.mlp.experts.0.gate_proj.weight", + "BF16", + &[2], + &[6, 7, 8, 9], + ), + ( + "model.layers.0.mlp.experts.1.gate_proj.weight", + "BF16", + &[2], + &[10, 11, 12, 13], + ), + ], + ); + let metadata = metadata_from_hf_config(&root, 7).unwrap(); + let validation = validate_raw_safetensors_gguf( + &root, + RawGgufWriteOptions { + buffer_size: 3, + metadata: Some(metadata.clone()), + tensor_name_map: TensorNameMap::HfToGguf, + split: None, + output_type: None, + tensor_selection: TensorSelection::All, + }, + ) + .unwrap(); + assert_eq!(validation.selected_tensor_count, 6); + + let output = root.join("qwen2-moe-native.gguf"); + write_raw_safetensors_gguf( + &root, + &output, + RawGgufWriteOptions { + buffer_size: 3, + metadata: Some(metadata), + tensor_name_map: TensorNameMap::HfToGguf, + split: None, + output_type: None, + tensor_selection: TensorSelection::All, + }, + ) + .unwrap(); + let bytes = fs::read(&output).unwrap(); + let parsed = parse_test_gguf(&bytes); + + assert_eq!( + parsed.tensor("blk.0.ffn_gate_inp_shexp.weight").ggml_type, + GGML_TYPE_F32 + ); + assert_eq!( + parsed.tensor("blk.0.ffn_gate_shexp.weight").ggml_type, + GGML_TYPE_F32 + ); + let merged_experts = parsed.tensor("blk.0.ffn_gate_exps.weight"); + assert_eq!(merged_experts.dims, vec![2, 2]); + assert_eq!(merged_experts.ggml_type, GGML_TYPE_BF16); + assert_eq!( + &bytes[merged_experts.absolute_offset..merged_experts.absolute_offset + 8], + &[6, 7, 8, 9, 10, 11, 12, 13] + ); + fs::remove_dir_all(root).unwrap(); +} + +#[test] +fn validates_qwen3_moe_native_conversion_fixture() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + write_qwen3_moe_config_and_tokenizer(&root); + write_safetensor( + &root.join("model.safetensors"), + &[ + ("model.embed_tokens.weight", "F32", &[1], &[1, 0, 0, 0]), + ( + "model.layers.0.input_layernorm.weight", + "F32", + &[1], + &[2, 0, 0, 0], + ), + ( + "model.layers.0.self_attn.q_norm.weight", + "F32", + &[1], + &[3, 0, 0, 0], + ), + ( + "model.layers.0.self_attn.k_norm.weight", + "F32", + &[1], + &[4, 0, 0, 0], + ), + ("model.layers.0.mlp.gate.weight", "F32", &[1], &[5, 0, 0, 0]), + ( + "model.layers.0.mlp.experts.0.down_proj.weight", + "BF16", + &[2], + &[6, 7, 8, 9], + ), + ( + "model.layers.0.mlp.experts.1.down_proj.weight", + "BF16", + &[2], + &[10, 11, 12, 13], + ), + ("model.norm.weight", "F32", &[1], &[14, 0, 0, 0]), + ], + ); + let metadata = metadata_from_hf_config(&root, 8).unwrap(); + let output = root.join("qwen3-moe-native.gguf"); + + write_raw_safetensors_gguf( + &root, + &output, + RawGgufWriteOptions { + buffer_size: 3, + metadata: Some(metadata), + tensor_name_map: TensorNameMap::HfToGguf, + split: None, + output_type: None, + tensor_selection: TensorSelection::All, + }, + ) + .unwrap(); + let bytes = fs::read(&output).unwrap(); + let parsed = parse_test_gguf(&bytes); + + assert_eq!( + parsed.tensor("blk.0.attn_q_norm.weight").ggml_type, + GGML_TYPE_F32 + ); + assert_eq!( + parsed.tensor("blk.0.ffn_gate_inp.weight").ggml_type, + GGML_TYPE_F32 + ); + let merged_experts = parsed.tensor("blk.0.ffn_down_exps.weight"); + assert_eq!(merged_experts.dims, vec![2, 2]); + assert_eq!(merged_experts.ggml_type, GGML_TYPE_BF16); + assert_eq!( + &bytes[merged_experts.absolute_offset..merged_experts.absolute_offset + 8], + &[6, 7, 8, 9, 10, 11, 12, 13] + ); + fs::remove_dir_all(root).unwrap(); +} + +#[test] +fn validates_llama_dense_native_conversion_fixture() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + write_llama_config_and_tokenizer(&root); + write_dense_hf_safetensor(&root); + let metadata = metadata_from_hf_config(&root, 14).unwrap(); + let validation = validate_raw_safetensors_gguf( + &root, + RawGgufWriteOptions { + buffer_size: 4, + metadata: Some(metadata.clone()), + tensor_name_map: TensorNameMap::HfToGguf, + split: None, + output_type: Some(ConvertOutputType::Bf16), + tensor_selection: TensorSelection::All, + }, + ) + .unwrap(); + assert_eq!(validation.selected_tensor_count, 14); + + let output = root.join("llama-native.gguf"); + write_raw_safetensors_gguf( + &root, + &output, + RawGgufWriteOptions { + buffer_size: 4, + metadata: Some(metadata), + tensor_name_map: TensorNameMap::HfToGguf, + split: None, + output_type: Some(ConvertOutputType::Bf16), + tensor_selection: TensorSelection::All, + }, + ) + .unwrap(); + let bytes = fs::read(&output).unwrap(); + let parsed = parse_test_gguf(&bytes); + assert!(parsed.metadata_count > 10); + assert_eq!( + parsed.tensor("blk.0.attn_q.weight").ggml_type, + GGML_TYPE_F32 + ); + fs::remove_dir_all(root).unwrap(); +} + +#[test] +fn streams_expert_tensors_as_merged_gguf_tensor() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + write_safetensor( + &root.join("model.safetensors"), + &[ + ( + "model.layers.1.mlp.experts.1.gate_proj.weight", + "BF16", + &[2, 2], + &[5, 6, 7, 8, 9, 10, 11, 12], + ), + ( + "model.layers.1.mlp.experts.0.gate_proj.weight", + "BF16", + &[2, 2], + &[1, 2, 3, 4, 13, 14, 15, 16], + ), + ], + ); + let output = root.join("experts.gguf"); + + write_raw_safetensors_gguf( + &root, + &output, + RawGgufWriteOptions { + buffer_size: 3, + metadata: None, + tensor_name_map: TensorNameMap::HfToGguf, + split: None, + output_type: None, + tensor_selection: TensorSelection::All, + }, + ) + .unwrap(); + + let bytes = fs::read(&output).unwrap(); + let parsed = parse_test_gguf(&bytes); + assert_eq!(parsed.tensor_count, 1); + assert_eq!(parsed.tensors[0].name, "blk.1.ffn_gate_exps.weight"); + assert_eq!(parsed.tensors[0].dims, vec![2, 2, 2]); + assert_eq!( + &bytes[parsed.tensors[0].absolute_offset..parsed.tensors[0].absolute_offset + 16], + &[1, 2, 3, 4, 13, 14, 15, 16, 5, 6, 7, 8, 9, 10, 11, 12] + ); + fs::remove_dir_all(root).unwrap(); +} + +#[test] +fn writes_only_selected_split_with_split_metadata() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + write_safetensor( + &root.join("model.safetensors"), + &[ + ("a.weight", "F32", &[1], &[1, 0, 0, 0]), + ("b.weight", "F32", &[1], &[2, 0, 0, 0]), + ("c.weight", "F32", &[1], &[3, 0, 0, 0]), + ("d.weight", "F32", &[1], &[4, 0, 0, 0]), + ], + ); + let output = root.join("split.gguf"); + + write_raw_safetensors_gguf( + &root, + &output, + RawGgufWriteOptions { + buffer_size: 2, + metadata: Some(vec![GgufKv::array_string( + "tokenizer.ggml.tokens", + vec!["a".to_string()], + )]), + tensor_name_map: TensorNameMap::Raw, + split: Some(GgufSplit { + split_index: 2, + split_count: 2, + }), + output_type: None, + tensor_selection: TensorSelection::All, + }, + ) + .unwrap(); + + let bytes = fs::read(&output).unwrap(); + let parsed = parse_test_gguf(&bytes); + assert_eq!(parsed.tensor_count, 2); + assert_eq!(parsed.metadata_count, 4); + assert_eq!(parsed.tensors[0].name, "c.weight"); + assert_eq!(parsed.tensors[1].name, "d.weight"); + assert_eq!(parsed.tensors[0].absolute_offset, parsed.data_start); + assert_eq!( + &bytes[parsed.tensors[0].absolute_offset..parsed.tensors[0].absolute_offset + 4], + &[3, 0, 0, 0] + ); + assert_eq!( + &bytes[parsed.tensors[1].absolute_offset..parsed.tensors[1].absolute_offset + 4], + &[4, 0, 0, 0] + ); + fs::remove_dir_all(root).unwrap(); +} + +#[test] +fn native_splits_are_byte_balanced_not_tensor_count_balanced() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + write_safetensor( + &root.join("model.safetensors"), + &[ + ("a.weight", "F32", &[64], &[1; 256]), + ("b.weight", "F32", &[1], &[2, 0, 0, 0]), + ("c.weight", "F32", &[1], &[3, 0, 0, 0]), + ("d.weight", "F32", &[1], &[4, 0, 0, 0]), + ], + ); + let output = root.join("split.gguf"); + + write_raw_safetensors_gguf( + &root, + &output, + RawGgufWriteOptions { + buffer_size: 64, + metadata: None, + tensor_name_map: TensorNameMap::Raw, + split: Some(GgufSplit { + split_index: 1, + split_count: 2, + }), + output_type: None, + tensor_selection: TensorSelection::All, + }, + ) + .unwrap(); + + let bytes = fs::read(&output).unwrap(); + let parsed = parse_test_gguf(&bytes); + assert_eq!(parsed.tensor_count, 1); + assert_eq!(parsed.tensors[0].name, "a.weight"); + assert_eq!(parsed.tensors[0].absolute_offset, parsed.data_start); + fs::remove_dir_all(root).unwrap(); +} + +#[test] +fn keeps_rank_one_f32_tensor_as_f32_for_bf16_output() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + write_safetensor( + &root.join("model.safetensors"), + &[("a.weight", "F32", &[2], &[0, 0, 0x80, 0x3f, 0, 0, 0, 0x40])], + ); + let output = root.join("bf16.gguf"); + + write_raw_safetensors_gguf( + &root, + &output, + RawGgufWriteOptions { + buffer_size: 4, + metadata: None, + tensor_name_map: TensorNameMap::Raw, + split: None, + output_type: Some(ConvertOutputType::Bf16), + tensor_selection: TensorSelection::All, + }, + ) + .unwrap(); + + let bytes = fs::read(&output).unwrap(); + let parsed = parse_test_gguf(&bytes); + assert_eq!(parsed.tensors[0].ggml_type, GGML_TYPE_F32); + assert_eq!( + &bytes[parsed.tensors[0].absolute_offset..parsed.tensors[0].absolute_offset + 8], + &[0, 0, 0x80, 0x3f, 0, 0, 0, 0x40] + ); + fs::remove_dir_all(root).unwrap(); +} + +struct ParsedGguf { + tensor_count: u64, + metadata_count: u64, + data_start: usize, + tensors: Vec, +} + +impl ParsedGguf { + fn tensor(&self, name: &str) -> &ParsedTensor { + self.tensors + .iter() + .find(|tensor| tensor.name == name) + .unwrap_or_else(|| panic!("missing tensor {name}")) + } +} + +struct ParsedTensor { + name: String, + dims: Vec, + ggml_type: u32, + absolute_offset: usize, +} + +fn parse_test_gguf(bytes: &[u8]) -> ParsedGguf { + let mut cursor = std::io::Cursor::new(bytes); + let mut magic = [0_u8; 4]; + cursor.read_exact(&mut magic).unwrap(); + assert_eq!(&magic, GGUF_MAGIC); + assert_eq!(read_u32(&mut cursor), GGUF_VERSION); + let tensor_count = read_u64(&mut cursor); + let metadata_count = read_u64(&mut cursor); + for _ in 0..metadata_count { + let _key = read_string(&mut cursor); + let value_type = read_u32(&mut cursor); + match value_type { + GGUF_TYPE_BOOL => { + let mut value = [0_u8; 1]; + cursor.read_exact(&mut value).unwrap(); + } + GGUF_TYPE_UINT16 => { + let mut value = [0_u8; 2]; + cursor.read_exact(&mut value).unwrap(); + } + GGUF_TYPE_INT32 => { + let _ = read_u32(&mut cursor); + } + GGUF_TYPE_STRING => { + let _ = read_string(&mut cursor); + } + GGUF_TYPE_UINT32 => { + let _ = read_u32(&mut cursor); + } + GGUF_TYPE_FLOAT32 => { + let _ = read_u32(&mut cursor); + } + GGUF_TYPE_UINT64 => { + let _ = read_u64(&mut cursor); + } + GGUF_TYPE_ARRAY => skip_array(&mut cursor), + other => panic!("unexpected metadata type {other}"), + } + } + let mut tensors = Vec::new(); + for _ in 0..tensor_count { + let name = read_string(&mut cursor); + let dim_count = read_u32(&mut cursor); + let dims = (0..dim_count) + .map(|_| read_u64(&mut cursor)) + .collect::>(); + let ggml_type = read_u32(&mut cursor); + let relative_offset = read_u64(&mut cursor); + tensors.push((name, dims, ggml_type, relative_offset)); + } + let data_start = align_to(cursor.position(), GGUF_ALIGNMENT) as usize; + ParsedGguf { + tensor_count, + metadata_count, + data_start, + tensors: tensors + .into_iter() + .map(|(name, dims, ggml_type, relative_offset)| ParsedTensor { + name, + dims, + ggml_type, + absolute_offset: data_start + relative_offset as usize, + }) + .collect(), + } +} + +fn read_string(cursor: &mut std::io::Cursor<&[u8]>) -> String { + let len = read_u64(cursor); + let mut bytes = vec![0_u8; len as usize]; + cursor.read_exact(&mut bytes).unwrap(); + String::from_utf8(bytes).unwrap() +} + +fn read_u32(cursor: &mut std::io::Cursor<&[u8]>) -> u32 { + let mut bytes = [0_u8; 4]; + cursor.read_exact(&mut bytes).unwrap(); + u32::from_le_bytes(bytes) +} + +fn read_u64(cursor: &mut std::io::Cursor<&[u8]>) -> u64 { + let mut bytes = [0_u8; 8]; + cursor.read_exact(&mut bytes).unwrap(); + u64::from_le_bytes(bytes) +} + +fn skip_array(cursor: &mut std::io::Cursor<&[u8]>) { + let element_type = read_u32(cursor); + let len = read_u64(cursor); + for _ in 0..len { + match element_type { + GGUF_TYPE_STRING => { + let _ = read_string(cursor); + } + GGUF_TYPE_INT32 | GGUF_TYPE_FLOAT32 | GGUF_TYPE_UINT32 => { + let _ = read_u32(cursor); + } + other => panic!("unexpected test array element type {other}"), + } + } +} + +fn unique_temp_dir() -> PathBuf { + static NEXT_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let id = NEXT_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + std::env::temp_dir().join(format!("skippy-gguf-writer-{nanos}-{id}")) +} + +fn write_safetensor(path: &Path, tensors: &[(&str, &str, &[u64], &[u8])]) { + let mut offset = 0_u64; + let mut entries = serde_json::Map::new(); + for (name, dtype, shape, bytes) in tensors { + let end = offset + bytes.len() as u64; + entries.insert( + (*name).to_string(), + serde_json::json!({ + "dtype": dtype, + "shape": shape, + "data_offsets": [offset, end], + }), + ); + offset = end; + } + let header = serde_json::Value::Object(entries).to_string(); + let mut bytes = Vec::new(); + bytes.extend_from_slice(&(header.len() as u64).to_le_bytes()); + bytes.extend_from_slice(header.as_bytes()); + for (_, _, _, tensor_bytes) in tensors { + bytes.extend_from_slice(tensor_bytes); + } + fs::write(path, bytes).unwrap(); +} + +fn write_qwen_config_and_tokenizer(root: &Path) { + fs::write( + root.join("config.json"), + r#"{ + "model_type": "qwen3", + "vocab_size": 4, + "max_position_embeddings": 128, + "hidden_size": 4, + "intermediate_size": 8, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 2, + "rope_theta": 1000000, + "rms_norm_eps": 1e-6 + }"#, + ) + .unwrap(); + fs::write( + root.join("tokenizer.json"), + r#"{ + "model": { + "type": "BPE", + "vocab": {"a": 0, "b": 1, "<|endoftext|>": 2, "<|im_end|>": 3}, + "merges": ["a b"] + }, + "decoder": {"type": "ByteLevel"}, + "added_tokens": [ + {"id": 2, "content": "<|endoftext|>", "special": true}, + {"id": 3, "content": "<|im_end|>", "special": true} + ] + }"#, + ) + .unwrap(); + fs::write( + root.join("tokenizer_config.json"), + r#"{"eos_token": "<|im_end|>", "pad_token": "<|endoftext|>", "add_bos_token": false}"#, + ) + .unwrap(); +} + +fn write_qwen2_moe_config_and_tokenizer(root: &Path) { + write_qwen_config_and_tokenizer(root); + fs::write( + root.join("config.json"), + r#"{ + "model_type": "qwen2_moe", + "vocab_size": 4, + "max_position_embeddings": 128, + "hidden_size": 4, + "intermediate_size": 8, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 2, + "num_experts": 2, + "num_experts_per_tok": 1, + "moe_intermediate_size": 2, + "shared_expert_intermediate_size": 8, + "rope_theta": 1000000, + "rms_norm_eps": 1e-6 + }"#, + ) + .unwrap(); +} + +fn write_qwen3_moe_config_and_tokenizer(root: &Path) { + write_qwen_config_and_tokenizer(root); + fs::write( + root.join("config.json"), + r#"{ + "model_type": "qwen3_moe", + "vocab_size": 4, + "max_position_embeddings": 128, + "hidden_size": 4, + "intermediate_size": 8, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 2, + "num_experts": 2, + "num_experts_per_tok": 1, + "moe_intermediate_size": 2, + "rope_theta": 1000000, + "rms_norm_eps": 1e-6 + }"#, + ) + .unwrap(); +} + +fn write_llama_config_and_tokenizer(root: &Path) { + fs::write( + root.join("config.json"), + r#"{ + "model_type": "llama", + "vocab_size": 4, + "max_position_embeddings": 128, + "hidden_size": 4, + "intermediate_size": 8, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 2, + "rope_theta": 500000, + "rms_norm_eps": 1e-5 + }"#, + ) + .unwrap(); + fs::write( + root.join("tokenizer.json"), + r#"{ + "model": { + "type": "BPE", + "vocab": {"a": 0, "b": 1, "<|end_of_text|>": 2, "<|start_header_id|>": 3}, + "merges": ["a b"] + }, + "decoder": {"type": "ByteLevel"}, + "added_tokens": [ + {"id": 2, "content": "<|end_of_text|>", "special": true}, + {"id": 3, "content": "<|start_header_id|>", "special": true} + ] + }"#, + ) + .unwrap(); + fs::write( + root.join("tokenizer_config.json"), + r#"{"eos_token": "<|end_of_text|>", "add_bos_token": true}"#, + ) + .unwrap(); +} + +fn write_dense_hf_safetensor(root: &Path) { + write_safetensor( + &root.join("model.safetensors"), + &[ + ("model.embed_tokens.weight", "F32", &[1], &[1, 0, 0, 0]), + ( + "model.layers.0.input_layernorm.weight", + "F32", + &[1], + &[2, 0, 0, 0], + ), + ( + "model.layers.0.self_attn.q_proj.weight", + "F32", + &[1], + &[3, 0, 0, 0], + ), + ( + "model.layers.0.self_attn.k_proj.weight", + "F32", + &[1], + &[4, 0, 0, 0], + ), + ( + "model.layers.0.self_attn.v_proj.weight", + "F32", + &[1], + &[5, 0, 0, 0], + ), + ( + "model.layers.0.self_attn.o_proj.weight", + "F32", + &[1], + &[6, 0, 0, 0], + ), + ( + "model.layers.0.self_attn.q_norm.weight", + "F32", + &[1], + &[7, 0, 0, 0], + ), + ( + "model.layers.0.self_attn.k_norm.weight", + "F32", + &[1], + &[8, 0, 0, 0], + ), + ( + "model.layers.0.post_attention_layernorm.weight", + "F32", + &[1], + &[9, 0, 0, 0], + ), + ( + "model.layers.0.mlp.gate_proj.weight", + "F32", + &[1], + &[10, 0, 0, 0], + ), + ( + "model.layers.0.mlp.up_proj.weight", + "F32", + &[1], + &[11, 0, 0, 0], + ), + ( + "model.layers.0.mlp.down_proj.weight", + "F32", + &[1], + &[12, 0, 0, 0], + ), + ("model.norm.weight", "F32", &[1], &[13, 0, 0, 0]), + ("lm_head.weight", "F32", &[1], &[14, 0, 0, 0]), + ], + ); +} diff --git a/crates/skippy-quantize/src/hf_checkpoint.rs b/crates/skippy-quantize/src/hf_checkpoint.rs new file mode 100644 index 0000000000..e898523f90 --- /dev/null +++ b/crates/skippy-quantize/src/hf_checkpoint.rs @@ -0,0 +1,731 @@ +use std::collections::{BTreeMap, BTreeSet}; +use std::fs::{self, File}; +use std::io::{Read, Seek, SeekFrom, Write}; +use std::ops::Range; +use std::path::{Path, PathBuf}; + +use anyhow::{Context, Result, anyhow, ensure}; +use serde::{Deserialize, Serialize}; + +use crate::memory_budget::MemorySize; +use crate::types::ConvertOutputType; + +#[derive(Debug, Serialize)] +pub(crate) struct HfCheckpointPlan { + pub(crate) source: PathBuf, + pub(crate) safetensor_count: usize, + pub(crate) tensor_count: usize, + pub(crate) total_tensor_bytes: u64, + pub(crate) largest_tensor_bytes: u64, + pub(crate) source_windows: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub(crate) stream_verification: Option, +} + +#[derive(Debug, Serialize)] +pub(crate) struct HfSourceWindow { + pub(crate) index: u32, + pub(crate) files: Vec, + pub(crate) tensor_count: usize, + pub(crate) total_tensor_bytes: u64, + pub(crate) largest_tensor_bytes: u64, +} + +#[derive(Debug, Serialize)] +pub(crate) struct HfStreamVerification { + pub(crate) safetensor_count: usize, + pub(crate) tensor_count: usize, + pub(crate) streamed_bytes: u64, + pub(crate) buffer_size: usize, +} + +#[derive(Debug)] +struct SafetensorSummary { + path: PathBuf, + tensor_count: usize, + total_tensor_bytes: u64, + largest_tensor_bytes: u64, +} + +#[derive(Debug, Deserialize)] +struct SafetensorTensor { + dtype: String, + shape: Vec, + data_offsets: [u64; 2], +} + +#[derive(Debug)] +pub(crate) struct SafetensorFile { + path: PathBuf, + data_start: u64, + tensors: BTreeMap, +} + +impl SafetensorFile { + pub(crate) fn open(path: &Path) -> Result { + let (data_start, raw_tensors) = read_safetensor_header(path)?; + let file_len = fs::metadata(path) + .with_context(|| format!("stat {}", path.display()))? + .len(); + let mut tensors = BTreeMap::new(); + for (name, tensor) in raw_tensors { + tensors.insert( + name.clone(), + SafetensorTensorInfo::from_raw(name, tensor, data_start, file_len)?, + ); + } + Ok(Self { + path: path.to_path_buf(), + data_start, + tensors, + }) + } + + pub(crate) fn path(&self) -> &Path { + &self.path + } + + pub(crate) fn data_start(&self) -> u64 { + self.data_start + } + + pub(crate) fn tensors(&self) -> &BTreeMap { + &self.tensors + } + + pub(crate) fn stream_tensor( + &self, + name: &str, + writer: &mut W, + buffer_size: usize, + ) -> Result { + let tensor = self + .tensors + .get(name) + .with_context(|| format!("tensor {name} not found in {}", self.path.display()))?; + stream_file_range( + &self.path, + tensor.absolute_data_range(), + writer, + buffer_size, + ) + } + + pub(crate) fn stream_tensor_chunks( + &self, + name: &str, + buffer_size: usize, + mut on_chunk: F, + ) -> Result + where + F: FnMut(&[u8]) -> Result<()>, + { + let tensor = self + .tensors + .get(name) + .with_context(|| format!("tensor {name} not found in {}", self.path.display()))?; + stream_file_range_chunks( + &self.path, + tensor.absolute_data_range(), + buffer_size, + |chunk| on_chunk(chunk), + ) + } +} + +#[derive(Debug)] +pub(crate) struct SafetensorTensorInfo { + name: String, + dtype: String, + shape: Vec, + relative_data_offsets: [u64; 2], + absolute_data_start: u64, + byte_len: u64, +} + +impl SafetensorTensorInfo { + fn from_raw( + name: String, + tensor: SafetensorTensor, + data_start: u64, + file_len: u64, + ) -> Result { + let relative_start = tensor.data_offsets[0]; + let relative_end = tensor.data_offsets[1]; + let byte_len = relative_end + .checked_sub(relative_start) + .with_context(|| format!("invalid data_offsets for tensor {name}"))?; + let shape_bytes = tensor_shape_bytes(&tensor) + .with_context(|| format!("validate shape for tensor {name}"))?; + ensure!( + byte_len == shape_bytes, + "tensor {name} byte length {byte_len} does not match dtype/shape byte length {shape_bytes}" + ); + let absolute_data_start = data_start + .checked_add(relative_start) + .with_context(|| format!("absolute data offset overflow for tensor {name}"))?; + let absolute_data_end = absolute_data_start + .checked_add(byte_len) + .with_context(|| format!("absolute data end overflow for tensor {name}"))?; + ensure!( + absolute_data_end <= file_len, + "tensor {name} extends past end of safetensors file" + ); + Ok(Self { + name, + dtype: tensor.dtype, + shape: tensor.shape, + relative_data_offsets: tensor.data_offsets, + absolute_data_start, + byte_len, + }) + } + + pub(crate) fn name(&self) -> &str { + &self.name + } + + pub(crate) fn dtype(&self) -> &str { + &self.dtype + } + + pub(crate) fn shape(&self) -> &[u64] { + &self.shape + } + + pub(crate) fn relative_data_offsets(&self) -> [u64; 2] { + self.relative_data_offsets + } + + pub(crate) fn byte_len(&self) -> u64 { + self.byte_len + } + + fn absolute_data_range(&self) -> Range { + self.absolute_data_start..self.absolute_data_start + self.byte_len + } +} + +pub(crate) fn inspect_hf_checkpoint( + source: &Path, + max_memory: Option, + staging_fraction: f64, +) -> Result { + ensure!( + staging_fraction > 0.0 && staging_fraction <= 1.0, + "--staging-fraction must be in the range (0, 1]" + ); + let safetensors = discover_safetensors(source)?; + ensure!( + !safetensors.is_empty(), + "no safetensors files found under {}", + source.display() + ); + let mut summaries = safetensors + .iter() + .map(|path| summarize_safetensor(path)) + .collect::>>()?; + summaries.sort_by(|a, b| a.path.cmp(&b.path)); + let tensor_count = summaries.iter().map(|summary| summary.tensor_count).sum(); + let total_tensor_bytes = summaries + .iter() + .map(|summary| summary.total_tensor_bytes) + .sum(); + let largest_tensor_bytes = summaries + .iter() + .map(|summary| summary.largest_tensor_bytes) + .max() + .unwrap_or(0); + let source_windows = plan_source_windows(&summaries, max_memory, staging_fraction)?; + Ok(HfCheckpointPlan { + source: source.to_path_buf(), + safetensor_count: summaries.len(), + tensor_count, + total_tensor_bytes, + largest_tensor_bytes, + source_windows, + stream_verification: None, + }) +} + +pub(crate) fn verify_hf_checkpoint_tensor_streams( + source: &Path, + buffer_size: usize, +) -> Result { + let safetensors = discover_safetensors(source)?; + let mut sink = std::io::sink(); + let mut tensor_count = 0_usize; + let mut streamed_bytes = 0_u64; + for path in &safetensors { + let safetensor = SafetensorFile::open(path)?; + ensure!( + safetensor.path().is_file(), + "safetensor path is not a file: {}", + safetensor.path().display() + ); + ensure!( + safetensor.data_start() >= 8, + "invalid safetensors data start in {}", + safetensor.path().display() + ); + for tensor in safetensor.tensors().values() { + ensure!( + !tensor.name().is_empty(), + "safetensor tensor has empty name" + ); + ensure!( + dtype_size(tensor.dtype()).is_some(), + "unsupported safetensors dtype {}", + tensor.dtype() + ); + let offsets = tensor.relative_data_offsets(); + ensure!( + offsets[0] <= offsets[1], + "invalid safetensors offsets for {}", + tensor.name() + ); + let _rank = tensor.shape().len(); + streamed_bytes += safetensor.stream_tensor(tensor.name(), &mut sink, buffer_size)?; + tensor_count += 1; + } + } + Ok(HfStreamVerification { + safetensor_count: safetensors.len(), + tensor_count, + streamed_bytes, + buffer_size, + }) +} + +pub(crate) fn open_safetensor_files(source: &Path) -> Result> { + discover_safetensors(source)? + .iter() + .map(|path| SafetensorFile::open(path)) + .collect() +} + +pub(crate) fn resolve_auto_output_type( + source: &Path, + requested: ConvertOutputType, +) -> Result { + if requested != ConvertOutputType::Auto { + return Ok(requested); + } + for safetensor in open_safetensor_files(source)? { + for tensor in safetensor.tensors().values() { + if tensor.shape().len() < 2 { + continue; + } + match tensor.dtype() { + "BF16" => return Ok(ConvertOutputType::Bf16), + "F16" => return Ok(ConvertOutputType::F16), + _ => {} + } + } + } + Ok(ConvertOutputType::F16) +} + +fn discover_safetensors(source: &Path) -> Result> { + ensure!( + source.is_dir(), + "HF checkpoint source must be a directory: {}", + source.display() + ); + let mut indexed = discover_indexed_safetensors(source)?; + if !indexed.is_empty() { + return Ok(indexed); + } + indexed = fs::read_dir(source) + .with_context(|| format!("read checkpoint directory {}", source.display()))? + .map(|entry| entry.map(|entry| entry.path())) + .collect::>>()? + .into_iter() + .filter(|path| path.extension().is_some_and(|ext| ext == "safetensors")) + .collect(); + indexed.sort(); + Ok(indexed) +} + +fn discover_indexed_safetensors(source: &Path) -> Result> { + let index_path = source.join("model.safetensors.index.json"); + if !index_path.is_file() { + return Ok(Vec::new()); + } + let index: SafetensorIndex = serde_json::from_slice( + &fs::read(&index_path).with_context(|| format!("read {}", index_path.display()))?, + ) + .with_context(|| format!("parse {}", index_path.display()))?; + let mut files = index + .weight_map + .values() + .map(|name| source.join(name)) + .collect::>() + .into_iter() + .collect::>(); + files.sort(); + Ok(files) +} + +#[derive(Debug, Deserialize)] +struct SafetensorIndex { + weight_map: BTreeMap, +} + +fn summarize_safetensor(path: &Path) -> Result { + let safetensor = SafetensorFile::open(path)?; + let tensor_count = safetensor.tensors.len(); + let total_tensor_bytes = safetensor + .tensors + .values() + .map(SafetensorTensorInfo::byte_len) + .sum(); + let largest_tensor_bytes = safetensor + .tensors + .values() + .map(SafetensorTensorInfo::byte_len) + .max() + .unwrap_or(0); + Ok(SafetensorSummary { + path: path.to_path_buf(), + tensor_count, + total_tensor_bytes, + largest_tensor_bytes, + }) +} + +fn read_safetensor_header(path: &Path) -> Result<(u64, BTreeMap)> { + let mut file = File::open(path).with_context(|| format!("open {}", path.display()))?; + let mut len_bytes = [0_u8; 8]; + file.read_exact(&mut len_bytes) + .with_context(|| format!("read safetensors header length from {}", path.display()))?; + let header_len = u64::from_le_bytes(len_bytes); + ensure!( + header_len <= 256 * 1024 * 1024, + "safetensors header is unexpectedly large in {}: {header_len} bytes", + path.display() + ); + let mut header = vec![0_u8; header_len as usize]; + file.read_exact(&mut header) + .with_context(|| format!("read safetensors header from {}", path.display()))?; + let raw: BTreeMap = serde_json::from_slice(&header) + .with_context(|| format!("parse safetensors header {}", path.display()))?; + let mut tensors = BTreeMap::new(); + for (name, value) in raw { + if name == "__metadata__" { + continue; + } + tensors.insert(name, serde_json::from_value(value)?); + } + let data_start = 8_u64 + .checked_add(header_len) + .with_context(|| format!("safetensors data start overflow in {}", path.display()))?; + Ok((data_start, tensors)) +} + +fn tensor_shape_bytes(tensor: &SafetensorTensor) -> Result { + let element_size = dtype_size(&tensor.dtype) + .ok_or_else(|| anyhow!("unsupported safetensors dtype {}", tensor.dtype))?; + let elements = tensor.shape.iter().try_fold(1_u64, |acc, dim| { + acc.checked_mul(*dim).context("tensor shape overflow") + })?; + elements + .checked_mul(element_size) + .context("tensor byte size overflow") +} + +fn dtype_size(dtype: &str) -> Option { + match dtype { + "BOOL" | "I8" | "U8" | "F8_E4M3" | "F8_E5M2" => Some(1), + "I16" | "U16" | "F16" | "BF16" => Some(2), + "I32" | "U32" | "F32" => Some(4), + "I64" | "U64" | "F64" => Some(8), + _ => None, + } +} + +fn plan_source_windows( + summaries: &[SafetensorSummary], + max_memory: Option, + staging_fraction: f64, +) -> Result> { + let budget = max_memory + .map(|memory| ((memory.bytes() as f64) * staging_fraction).floor() as u64) + .unwrap_or(u64::MAX) + .max(1); + let mut windows = Vec::new(); + let mut current = SourceWindowBuilder::new(1); + for summary in summaries { + if !current.is_empty() && current.total_tensor_bytes + summary.total_tensor_bytes > budget { + windows.push(current.finish()); + current = SourceWindowBuilder::new(windows.len() as u32 + 1); + } + current.push(summary); + } + if !current.is_empty() { + windows.push(current.finish()); + } + Ok(windows) +} + +struct SourceWindowBuilder { + index: u32, + files: Vec, + tensor_count: usize, + total_tensor_bytes: u64, + largest_tensor_bytes: u64, +} + +impl SourceWindowBuilder { + fn new(index: u32) -> Self { + Self { + index, + files: Vec::new(), + tensor_count: 0, + total_tensor_bytes: 0, + largest_tensor_bytes: 0, + } + } + + fn is_empty(&self) -> bool { + self.files.is_empty() + } + + fn push(&mut self, summary: &SafetensorSummary) { + self.files.push(summary.path.clone()); + self.tensor_count += summary.tensor_count; + self.total_tensor_bytes += summary.total_tensor_bytes; + self.largest_tensor_bytes = self.largest_tensor_bytes.max(summary.largest_tensor_bytes); + } + + fn finish(self) -> HfSourceWindow { + HfSourceWindow { + index: self.index, + files: self.files, + tensor_count: self.tensor_count, + total_tensor_bytes: self.total_tensor_bytes, + largest_tensor_bytes: self.largest_tensor_bytes, + } + } +} + +fn stream_file_range( + path: &Path, + range: Range, + writer: &mut W, + buffer_size: usize, +) -> Result { + stream_file_range_chunks(path, range, buffer_size, |chunk| { + writer.write_all(chunk).context("write tensor bytes") + }) +} + +fn stream_file_range_chunks( + path: &Path, + range: Range, + buffer_size: usize, + mut on_chunk: F, +) -> Result +where + F: FnMut(&[u8]) -> Result<()>, +{ + ensure!(buffer_size > 0, "buffer_size must be greater than zero"); + let mut file = File::open(path).with_context(|| format!("open {}", path.display()))?; + file.seek(SeekFrom::Start(range.start)) + .with_context(|| format!("seek {}", path.display()))?; + let mut remaining = range.end - range.start; + let mut copied = 0_u64; + let mut buffer = vec![0_u8; buffer_size]; + while remaining > 0 { + let read_len = buffer.len().min(remaining as usize); + file.read_exact(&mut buffer[..read_len]) + .with_context(|| format!("read tensor bytes from {}", path.display()))?; + on_chunk(&buffer[..read_len])?; + remaining -= read_len as u64; + copied += read_len as u64; + } + Ok(copied) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicU64, Ordering}; + + static TEMP_DIR_COUNTER: AtomicU64 = AtomicU64::new(0); + + #[test] + fn plans_unindexed_safetensors_under_memory_budget() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + write_safetensor( + &root.join("model-00001-of-00002.safetensors"), + &[("a.weight", "F32", &[2], &[1, 2, 3, 4, 5, 6, 7, 8])], + ); + write_safetensor( + &root.join("model-00002-of-00002.safetensors"), + &[("b.weight", "BF16", &[4], &[1, 2, 3, 4, 5, 6, 7, 8])], + ); + + let plan = + inspect_hf_checkpoint(&root, Some(MemorySize::from_bytes_for_tests(12)), 1.0).unwrap(); + + assert_eq!(plan.safetensor_count, 2); + assert_eq!(plan.tensor_count, 2); + assert_eq!(plan.total_tensor_bytes, 16); + assert_eq!(plan.source_windows.len(), 2); + assert_eq!(plan.source_windows[0].total_tensor_bytes, 8); + assert_eq!(plan.source_windows[1].total_tensor_bytes, 8); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn uses_index_weight_map_when_present() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + write_safetensor( + &root.join("shard-b.safetensors"), + &[("b.weight", "F32", &[1], &[1, 2, 3, 4])], + ); + write_safetensor( + &root.join("shard-a.safetensors"), + &[("a.weight", "F32", &[1], &[1, 2, 3, 4])], + ); + fs::write( + root.join("model.safetensors.index.json"), + r#"{"metadata":{},"weight_map":{"a.weight":"shard-a.safetensors","b.weight":"shard-b.safetensors"}}"#, + ) + .unwrap(); + + let plan = inspect_hf_checkpoint(&root, None, 1.0).unwrap(); + + assert_eq!(plan.safetensor_count, 2); + assert_eq!(plan.source_windows.len(), 1); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn streams_tensor_bytes_without_reading_neighbor_tensors() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + let path = root.join("model.safetensors"); + write_safetensor( + &path, + &[ + ("a.weight", "U8", &[4], &[1, 2, 3, 4]), + ("b.weight", "U8", &[3], &[9, 8, 7]), + ], + ); + + let safetensor = SafetensorFile::open(&path).unwrap(); + let tensor = safetensor.tensors().get("b.weight").unwrap(); + let mut output = Vec::new(); + let copied = safetensor + .stream_tensor("b.weight", &mut output, 2) + .unwrap(); + + assert_eq!(safetensor.path(), path); + assert!(safetensor.data_start() > 8); + assert_eq!(tensor.name(), "b.weight"); + assert_eq!(tensor.dtype(), "U8"); + assert_eq!(tensor.shape(), &[3]); + assert_eq!(tensor.relative_data_offsets(), [4, 7]); + assert_eq!(tensor.byte_len(), 3); + assert_eq!(copied, 3); + assert_eq!(output, vec![9, 8, 7]); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn resolves_auto_output_type_from_first_rank_two_float_tensor() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + write_safetensor( + &root.join("model.safetensors"), + &[ + ("a.bias", "BF16", &[4], &[1, 2, 3, 4, 5, 6, 7, 8]), + ("b.weight", "F16", &[2, 2], &[1, 2, 3, 4, 5, 6, 7, 8]), + ], + ); + + let output_type = resolve_auto_output_type(&root, ConvertOutputType::Auto).unwrap(); + + assert_eq!(output_type, ConvertOutputType::F16); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn resolves_auto_output_type_to_bf16_when_rank_two_bf16_appears_first() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + write_safetensor( + &root.join("model.safetensors"), + &[("a.weight", "BF16", &[2, 2], &[1, 2, 3, 4, 5, 6, 7, 8])], + ); + + let output_type = resolve_auto_output_type(&root, ConvertOutputType::Auto).unwrap(); + + assert_eq!(output_type, ConvertOutputType::Bf16); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn resolves_auto_output_type_to_f16_when_checkpoint_has_no_float_matrix() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + write_safetensor( + &root.join("model.safetensors"), + &[ + ("a.bias", "BF16", &[4], &[1, 2, 3, 4, 5, 6, 7, 8]), + ( + "b.count", + "I32", + &[2, 2], + &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], + ), + ], + ); + + let output_type = resolve_auto_output_type(&root, ConvertOutputType::Auto).unwrap(); + + assert_eq!(output_type, ConvertOutputType::F16); + fs::remove_dir_all(root).unwrap(); + } + + fn unique_temp_dir() -> PathBuf { + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let counter = TEMP_DIR_COUNTER.fetch_add(1, Ordering::Relaxed); + std::env::temp_dir().join(format!( + "skippy-hf-checkpoint-{}-{nanos}-{counter}", + std::process::id() + )) + } + + fn write_safetensor(path: &Path, tensors: &[(&str, &str, &[u64], &[u8])]) { + let mut offset = 0_u64; + let mut entries = serde_json::Map::new(); + for (name, dtype, shape, bytes) in tensors { + let end = offset + bytes.len() as u64; + entries.insert( + (*name).to_string(), + serde_json::json!({ + "dtype": dtype, + "shape": shape, + "data_offsets": [offset, end], + }), + ); + offset = end; + } + let header = serde_json::Value::Object(entries).to_string(); + let mut bytes = Vec::new(); + bytes.extend_from_slice(&(header.len() as u64).to_le_bytes()); + bytes.extend_from_slice(header.as_bytes()); + for (_, _, _, tensor_bytes) in tensors { + bytes.extend_from_slice(tensor_bytes); + } + fs::write(path, bytes).unwrap(); + } +} diff --git a/crates/skippy-quantize/src/imatrix.rs b/crates/skippy-quantize/src/imatrix.rs new file mode 100644 index 0000000000..1f83d7574e --- /dev/null +++ b/crates/skippy-quantize/src/imatrix.rs @@ -0,0 +1,655 @@ +use std::collections::BTreeMap; +use std::ffi::CString; +use std::fs; +use std::io::{Cursor, Read}; +use std::path::{Path, PathBuf}; + +use anyhow::{Context, Result, anyhow, ensure}; + +const GGUF_MAGIC: &[u8; 4] = b"GGUF"; +const GGUF_ALIGNMENT_DEFAULT: u64 = 32; +const GGML_TYPE_F32: u32 = 0; +const TENSOR_SUMS_SUFFIX: &str = ".in_sum2"; +const TENSOR_COUNTS_SUFFIX: &str = ".counts"; + +const GGUF_TYPE_UINT32: u32 = 4; +const GGUF_TYPE_FLOAT32: u32 = 6; +const GGUF_TYPE_BOOL: u32 = 7; +const GGUF_TYPE_STRING: u32 = 8; +const GGUF_TYPE_ARRAY: u32 = 9; +const GGUF_TYPE_UINT64: u32 = 10; +const GGUF_TYPE_INT64: u32 = 11; +const GGUF_TYPE_FLOAT64: u32 = 12; + +const KV_GENERAL_ALIGNMENT: &str = "general.alignment"; +const KV_IMATRIX_DATASETS: &str = "imatrix.datasets"; +const KV_IMATRIX_CHUNK_COUNT: &str = "imatrix.chunk_count"; + +pub(crate) struct NativeImatrix { + _names: Vec, + _values: Vec>, + entries: Vec, + source_path: PathBuf, + dataset: Option, + chunk_count: i32, +} + +impl NativeImatrix { + pub(crate) fn load( + path: &Path, + include_weights: &[String], + exclude_weights: &[String], + ) -> Result { + let bytes = + fs::read(path).with_context(|| format!("read imatrix file {}", path.display()))?; + let loaded = if bytes.starts_with(GGUF_MAGIC) { + load_gguf_imatrix(&bytes, path)? + } else { + load_legacy_imatrix(&bytes, path)? + }; + Self::from_loaded(path, loaded, include_weights, exclude_weights) + } + + pub(crate) fn as_ptr(&self) -> *const llama_quant_ffi::LlamaModelImatrixData { + self.entries.as_ptr() + } + + pub(crate) fn source_path(&self) -> &Path { + &self.source_path + } + + pub(crate) fn dataset(&self) -> Option<&str> { + self.dataset.as_deref() + } + + pub(crate) fn chunk_count(&self) -> i32 { + self.chunk_count + } + + pub(crate) fn entry_count(&self) -> usize { + self.entries.len().saturating_sub(1) + } + + fn from_loaded( + path: &Path, + loaded: LoadedImatrix, + include_weights: &[String], + exclude_weights: &[String], + ) -> Result { + let mut selected = loaded + .entries + .into_iter() + .filter(|entry| include_exclude_match(&entry.name, include_weights, exclude_weights)) + .collect::>(); + selected.sort_by(|a, b| a.name.cmp(&b.name)); + ensure!( + !selected.is_empty(), + "imatrix filters removed all entries from {}", + path.display() + ); + + let mut names = Vec::with_capacity(selected.len()); + let mut values = Vec::with_capacity(selected.len()); + for entry in selected { + names.push(CString::new(entry.name)?); + values.push(entry.values); + } + let mut entries = names + .iter() + .zip(values.iter()) + .map(|(name, value)| llama_quant_ffi::LlamaModelImatrixData { + name: name.as_ptr(), + data: value.as_ptr(), + size: value.len(), + }) + .collect::>(); + entries.push(llama_quant_ffi::LlamaModelImatrixData { + name: std::ptr::null(), + data: std::ptr::null(), + size: 0, + }); + Ok(Self { + _names: names, + _values: values, + entries, + source_path: path.to_path_buf(), + dataset: loaded.dataset, + chunk_count: loaded.chunk_count, + }) + } +} + +struct LoadedImatrix { + entries: Vec, + dataset: Option, + chunk_count: i32, +} + +struct ImatrixEntry { + name: String, + values: Vec, +} + +fn load_legacy_imatrix(bytes: &[u8], path: &Path) -> Result { + let mut reader = Cursor::new(bytes.to_vec()); + let entry_count = read_i32(&mut reader) + .with_context(|| format!("read imatrix entry count from {}", path.display()))?; + ensure!( + entry_count > 0, + "imatrix file has no entries: {}", + path.display() + ); + let mut entries = Vec::with_capacity(entry_count as usize); + for index in 0..entry_count { + entries.push(read_legacy_imatrix_entry(&mut reader, index)?); + } + let (chunk_count, dataset) = read_legacy_imatrix_trailer(&mut reader)?; + Ok(LoadedImatrix { + entries, + dataset, + chunk_count, + }) +} + +fn read_legacy_imatrix_entry(reader: &mut Cursor>, index: i32) -> Result { + let name_len = read_i32(reader).with_context(|| format!("read imatrix name length {index}"))?; + ensure!(name_len > 0, "imatrix entry {index} has empty name"); + let mut name_bytes = vec![0_u8; name_len as usize]; + reader + .read_exact(&mut name_bytes) + .with_context(|| format!("read imatrix name {index}"))?; + let name = String::from_utf8(name_bytes).with_context(|| format!("imatrix name {index}"))?; + let ncall = read_i32(reader).with_context(|| format!("read imatrix ncall for {name}"))?; + let value_count = + read_i32(reader).with_context(|| format!("read imatrix value count for {name}"))?; + ensure!(value_count > 0, "imatrix entry {name} has no values"); + let mut values = (0..value_count) + .map(|_| read_f32(reader)) + .collect::>>()?; + if ncall > 0 { + let denom = ncall as f32; + for value in &mut values { + *value /= denom; + } + } + Ok(ImatrixEntry { name, values }) +} + +fn read_legacy_imatrix_trailer(reader: &mut Cursor>) -> Result<(i32, Option)> { + if reader.position() as usize >= reader.get_ref().len() { + return Ok((0, None)); + } + let chunk_count = read_i32(reader)?; + if reader.position() as usize >= reader.get_ref().len() { + return Ok((chunk_count, None)); + } + let len = read_i32(reader)?; + if len <= 0 { + return Ok((chunk_count, None)); + } + let mut dataset = vec![0_u8; len as usize]; + reader.read_exact(&mut dataset)?; + Ok((chunk_count, Some(String::from_utf8(dataset)?))) +} + +fn load_gguf_imatrix(bytes: &[u8], path: &Path) -> Result { + let mut reader = Cursor::new(bytes.to_vec()); + let header = GgufHeader::read(&mut reader)?; + ensure!( + header.version >= 2, + "unsupported GGUF imatrix version {} in {}", + header.version, + path.display() + ); + let mut metadata = GgufMetadata::default(); + for _ in 0..header.metadata_count { + let key = read_gguf_string(&mut reader)?; + let value_type = read_u32(&mut reader)?; + metadata.read_value(&mut reader, &key, value_type)?; + } + + let mut tensors = Vec::with_capacity(header.tensor_count as usize); + for _ in 0..header.tensor_count { + tensors.push(GgufTensorInfo::read(&mut reader)?); + } + let data_start = align_to(reader.position(), metadata.alignment); + let entries = read_gguf_entries(bytes, &tensors, data_start)?; + ensure!( + !entries.is_empty(), + "GGUF imatrix has no paired tensors in {}", + path.display() + ); + Ok(LoadedImatrix { + entries, + dataset: metadata.datasets.into_iter().next(), + chunk_count: metadata.chunk_count.unwrap_or(0) as i32, + }) +} + +struct GgufHeader { + version: u32, + tensor_count: u64, + metadata_count: u64, +} + +impl GgufHeader { + fn read(reader: &mut Cursor>) -> Result { + let mut magic = [0_u8; 4]; + reader.read_exact(&mut magic)?; + ensure!(&magic == GGUF_MAGIC, "not a GGUF file"); + Ok(Self { + version: read_u32(reader)?, + tensor_count: read_u64(reader)?, + metadata_count: read_u64(reader)?, + }) + } +} + +#[derive(Default)] +struct GgufMetadata { + alignment: u64, + datasets: Vec, + chunk_count: Option, +} + +impl GgufMetadata { + fn read_value( + &mut self, + reader: &mut Cursor>, + key: &str, + value_type: u32, + ) -> Result<()> { + match (key, value_type) { + (KV_GENERAL_ALIGNMENT, GGUF_TYPE_UINT32) => { + self.alignment = read_u32(reader)? as u64; + } + (KV_GENERAL_ALIGNMENT, GGUF_TYPE_UINT64) => { + self.alignment = read_u64(reader)?; + } + (KV_IMATRIX_CHUNK_COUNT, GGUF_TYPE_UINT32) => { + self.chunk_count = Some(read_u32(reader)?); + } + (KV_IMATRIX_DATASETS, GGUF_TYPE_ARRAY) => { + self.datasets = read_string_array(reader)?; + } + _ => skip_gguf_value(reader, value_type)?, + } + if self.alignment == 0 { + self.alignment = GGUF_ALIGNMENT_DEFAULT; + } + Ok(()) + } +} + +struct GgufTensorInfo { + name: String, + dims: Vec, + tensor_type: u32, + offset: u64, +} + +impl GgufTensorInfo { + fn read(reader: &mut Cursor>) -> Result { + let name = read_gguf_string(reader)?; + let n_dims = read_u32(reader)?; + ensure!(n_dims > 0, "GGUF tensor {name} has no dimensions"); + let dims = (0..n_dims) + .map(|_| read_u64(reader)) + .collect::>>()?; + let tensor_type = read_u32(reader)?; + let offset = read_u64(reader)?; + Ok(Self { + name, + dims, + tensor_type, + offset, + }) + } + + fn element_count(&self) -> Result { + self.dims.iter().try_fold(1_usize, |acc, dim| { + acc.checked_mul(*dim as usize) + .with_context(|| format!("GGUF tensor {} is too large", self.name)) + }) + } +} + +fn read_gguf_entries( + bytes: &[u8], + tensors: &[GgufTensorInfo], + data_start: u64, +) -> Result> { + let mut sums = BTreeMap::>::new(); + let mut counts = BTreeMap::>::new(); + for tensor in tensors { + if tensor.tensor_type != GGML_TYPE_F32 { + continue; + } + let Some((base_name, kind)) = imatrix_tensor_name(&tensor.name) else { + continue; + }; + let values = read_gguf_f32_tensor(bytes, tensor, data_start)?; + match kind { + ImatrixTensorKind::Sums => { + sums.insert(base_name.to_string(), values); + } + ImatrixTensorKind::Counts => { + counts.insert(base_name.to_string(), values); + } + } + } + let mut entries = Vec::new(); + for (name, sum_values) in sums { + let count_values = counts + .remove(&name) + .with_context(|| format!("GGUF imatrix tensor {name} is missing counts"))?; + entries.push(ImatrixEntry { + name, + values: normalize_gguf_imatrix_values(&sum_values, &count_values)?, + }); + } + Ok(entries) +} + +enum ImatrixTensorKind { + Sums, + Counts, +} + +fn imatrix_tensor_name(name: &str) -> Option<(&str, ImatrixTensorKind)> { + if let Some(base) = name.strip_suffix(TENSOR_SUMS_SUFFIX) { + return Some((base, ImatrixTensorKind::Sums)); + } + name.strip_suffix(TENSOR_COUNTS_SUFFIX) + .map(|base| (base, ImatrixTensorKind::Counts)) +} + +fn normalize_gguf_imatrix_values(sums: &[f32], counts: &[f32]) -> Result> { + ensure!(!counts.is_empty(), "GGUF imatrix entry has no counts"); + ensure!( + sums.len().is_multiple_of(counts.len()), + "GGUF imatrix sums/counts shape mismatch" + ); + let values_per_count = sums.len() / counts.len(); + let mut output = vec![1.0_f32; sums.len()]; + for (count_index, count) in counts.iter().enumerate() { + if *count <= 0.0 { + continue; + } + let offset = count_index * values_per_count; + for i in 0..values_per_count { + output[offset + i] = sums[offset + i] / count; + } + } + Ok(output) +} + +fn read_gguf_f32_tensor( + bytes: &[u8], + tensor: &GgufTensorInfo, + data_start: u64, +) -> Result> { + let element_count = tensor.element_count()?; + let byte_len = element_count + .checked_mul(std::mem::size_of::()) + .with_context(|| format!("GGUF tensor {} is too large", tensor.name))?; + let offset = data_start + .checked_add(tensor.offset) + .with_context(|| format!("GGUF tensor {} offset overflow", tensor.name))? + as usize; + let end = offset + .checked_add(byte_len) + .with_context(|| format!("GGUF tensor {} byte range overflow", tensor.name))?; + ensure!( + end <= bytes.len(), + "GGUF tensor {} extends past end of file", + tensor.name + ); + bytes[offset..end] + .chunks_exact(4) + .map(|chunk| Ok(f32::from_le_bytes(chunk.try_into()?))) + .collect::>>() +} + +fn include_exclude_match( + name: &str, + include_weights: &[String], + exclude_weights: &[String], +) -> bool { + if !exclude_weights.is_empty() { + return !exclude_weights.iter().any(|filter| name.contains(filter)); + } + if !include_weights.is_empty() { + return include_weights.iter().any(|filter| name.contains(filter)); + } + true +} + +fn read_string_array(reader: &mut Cursor>) -> Result> { + let element_type = read_u32(reader)?; + ensure!( + element_type == GGUF_TYPE_STRING, + "expected GGUF string array, found type {element_type}" + ); + let len = read_u64(reader)?; + (0..len).map(|_| read_gguf_string(reader)).collect() +} + +fn skip_gguf_value(reader: &mut Cursor>, value_type: u32) -> Result<()> { + match value_type { + 0 | 1 => skip_bytes(reader, 1), + 2 | 3 => skip_bytes(reader, 2), + GGUF_TYPE_UINT32 | 5 | GGUF_TYPE_FLOAT32 => skip_bytes(reader, 4), + GGUF_TYPE_BOOL => skip_bytes(reader, 1), + GGUF_TYPE_STRING => { + let _ = read_gguf_string(reader)?; + Ok(()) + } + GGUF_TYPE_ARRAY => skip_gguf_array(reader), + GGUF_TYPE_UINT64 | GGUF_TYPE_INT64 | GGUF_TYPE_FLOAT64 => skip_bytes(reader, 8), + _ => Err(anyhow!("unsupported GGUF metadata type {value_type}")), + } +} + +fn skip_gguf_array(reader: &mut Cursor>) -> Result<()> { + let element_type = read_u32(reader)?; + let len = read_u64(reader)?; + for _ in 0..len { + skip_gguf_value(reader, element_type)?; + } + Ok(()) +} + +fn skip_bytes(reader: &mut Cursor>, len: u64) -> Result<()> { + let position = reader.position(); + let next = position + .checked_add(len) + .context("GGUF metadata offset overflow")?; + ensure!( + next <= reader.get_ref().len() as u64, + "GGUF metadata extends past end of file" + ); + reader.set_position(next); + Ok(()) +} + +fn read_gguf_string(reader: &mut Cursor>) -> Result { + let len = read_u64(reader)?; + let mut bytes = vec![0_u8; len as usize]; + reader.read_exact(&mut bytes)?; + String::from_utf8(bytes).context("GGUF string is not UTF-8") +} + +fn read_i32(reader: &mut Cursor>) -> Result { + let mut bytes = [0_u8; 4]; + reader.read_exact(&mut bytes)?; + Ok(i32::from_le_bytes(bytes)) +} + +fn read_u32(reader: &mut Cursor>) -> Result { + let mut bytes = [0_u8; 4]; + reader.read_exact(&mut bytes)?; + Ok(u32::from_le_bytes(bytes)) +} + +fn read_u64(reader: &mut Cursor>) -> Result { + let mut bytes = [0_u8; 8]; + reader.read_exact(&mut bytes)?; + Ok(u64::from_le_bytes(bytes)) +} + +fn read_f32(reader: &mut Cursor>) -> Result { + let mut bytes = [0_u8; 4]; + reader.read_exact(&mut bytes)?; + Ok(f32::from_le_bytes(bytes)) +} + +fn align_to(value: u64, alignment: u64) -> u64 { + if alignment <= 1 { + return value; + } + value.div_ceil(alignment) * alignment +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn loads_legacy_imatrix_with_include_filter() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + let imatrix_path = root.join("imatrix.dat"); + write_legacy_imatrix( + &imatrix_path, + &[ + ("blk.0.attn_q.weight", 2, &[2.0, 4.0]), + ("blk.0.ffn_down.weight", 1, &[9.0, 12.0]), + ], + ); + + let imatrix = + NativeImatrix::load(&imatrix_path, &["attn_q".to_string()], &Vec::new()).unwrap(); + + assert_eq!(imatrix.entry_count(), 1); + assert_eq!(imatrix._values[0], vec![1.0, 2.0]); + assert_eq!(imatrix.dataset(), Some("dataset.txt")); + assert_eq!(imatrix.chunk_count(), 3); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn loads_gguf_imatrix_with_normalized_counts() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + let imatrix_path = root.join("imatrix.gguf"); + write_gguf_imatrix(&imatrix_path); + + let imatrix = + NativeImatrix::load(&imatrix_path, &["attn_q".to_string()], &Vec::new()).unwrap(); + + assert_eq!(imatrix.entry_count(), 1); + assert_eq!(imatrix._values[0], vec![1.0, 2.0, 1.0, 1.0]); + assert_eq!(imatrix.dataset(), Some("calibration.txt")); + assert_eq!(imatrix.chunk_count(), 7); + fs::remove_dir_all(root).unwrap(); + } + + fn unique_temp_dir() -> PathBuf { + static NEXT_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let id = NEXT_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + std::env::temp_dir().join(format!("skippy-imatrix-{nanos}-{id}")) + } + + fn write_legacy_imatrix(path: &Path, entries: &[(&str, i32, &[f32])]) { + let mut bytes = Vec::new(); + bytes.extend_from_slice(&(entries.len() as i32).to_le_bytes()); + for (name, ncall, values) in entries { + bytes.extend_from_slice(&(name.len() as i32).to_le_bytes()); + bytes.extend_from_slice(name.as_bytes()); + bytes.extend_from_slice(&ncall.to_le_bytes()); + bytes.extend_from_slice(&(values.len() as i32).to_le_bytes()); + for value in *values { + bytes.extend_from_slice(&value.to_le_bytes()); + } + } + bytes.extend_from_slice(&3_i32.to_le_bytes()); + let dataset = "dataset.txt"; + bytes.extend_from_slice(&(dataset.len() as i32).to_le_bytes()); + bytes.extend_from_slice(dataset.as_bytes()); + fs::write(path, bytes).unwrap(); + } + + fn write_gguf_imatrix(path: &Path) { + let sums_name = "blk.0.attn_q.weight.in_sum2"; + let counts_name = "blk.0.attn_q.weight.counts"; + let mut bytes = Vec::new(); + bytes.extend_from_slice(GGUF_MAGIC); + bytes.extend_from_slice(&3_u32.to_le_bytes()); + bytes.extend_from_slice(&2_u64.to_le_bytes()); + bytes.extend_from_slice(&4_u64.to_le_bytes()); + write_gguf_kv_string(&mut bytes, "general.type", "imatrix"); + write_gguf_kv_u32(&mut bytes, KV_GENERAL_ALIGNMENT, 32); + write_gguf_kv_u32(&mut bytes, KV_IMATRIX_CHUNK_COUNT, 7); + write_gguf_kv_string_array(&mut bytes, KV_IMATRIX_DATASETS, &["calibration.txt"]); + write_gguf_tensor_info(&mut bytes, sums_name, &[2, 2], GGML_TYPE_F32, 0); + write_gguf_tensor_info(&mut bytes, counts_name, &[1, 2], GGML_TYPE_F32, 16); + while bytes.len() % 32 != 0 { + bytes.push(0); + } + for value in [2.0_f32, 4.0, 9.0, 11.0] { + bytes.extend_from_slice(&value.to_le_bytes()); + } + for value in [2.0_f32, 0.0] { + bytes.extend_from_slice(&value.to_le_bytes()); + } + fs::write(path, bytes).unwrap(); + } + + fn write_gguf_kv_string(bytes: &mut Vec, key: &str, value: &str) { + write_gguf_string(bytes, key); + bytes.extend_from_slice(&GGUF_TYPE_STRING.to_le_bytes()); + write_gguf_string(bytes, value); + } + + fn write_gguf_kv_u32(bytes: &mut Vec, key: &str, value: u32) { + write_gguf_string(bytes, key); + bytes.extend_from_slice(&GGUF_TYPE_UINT32.to_le_bytes()); + bytes.extend_from_slice(&value.to_le_bytes()); + } + + fn write_gguf_kv_string_array(bytes: &mut Vec, key: &str, values: &[&str]) { + write_gguf_string(bytes, key); + bytes.extend_from_slice(&GGUF_TYPE_ARRAY.to_le_bytes()); + bytes.extend_from_slice(&GGUF_TYPE_STRING.to_le_bytes()); + bytes.extend_from_slice(&(values.len() as u64).to_le_bytes()); + for value in values { + write_gguf_string(bytes, value); + } + } + + fn write_gguf_tensor_info( + bytes: &mut Vec, + name: &str, + dims: &[u64], + tensor_type: u32, + offset: u64, + ) { + write_gguf_string(bytes, name); + bytes.extend_from_slice(&(dims.len() as u32).to_le_bytes()); + for dim in dims { + bytes.extend_from_slice(&dim.to_le_bytes()); + } + bytes.extend_from_slice(&tensor_type.to_le_bytes()); + bytes.extend_from_slice(&offset.to_le_bytes()); + } + + fn write_gguf_string(bytes: &mut Vec, value: &str) { + bytes.extend_from_slice(&(value.len() as u64).to_le_bytes()); + bytes.extend_from_slice(value.as_bytes()); + } +} diff --git a/crates/skippy-quantize/src/llama_load.rs b/crates/skippy-quantize/src/llama_load.rs new file mode 100644 index 0000000000..51755c711b --- /dev/null +++ b/crates/skippy-quantize/src/llama_load.rs @@ -0,0 +1,219 @@ +use std::path::{Path, PathBuf}; +use std::process::Command; + +use anyhow::{Context, Result, ensure}; +use clap::Parser; +use serde::Serialize; + +use crate::output::{print_json_pretty, print_success}; +use crate::tool_paths::resolve_llama_cli; + +#[derive(Debug, Parser)] +pub(crate) struct ValidateLlamaLoadArgs { + #[arg(long)] + llama_cli: Option, + #[arg(long)] + check_tensors: bool, + #[arg(long)] + json: bool, + model: PathBuf, +} + +#[derive(Debug, Serialize)] +pub(crate) struct LlamaLoadReport { + pub(crate) model: PathBuf, + pub(crate) llama_cli: PathBuf, + pub(crate) command: Vec, + pub(crate) status_code: Option, + pub(crate) success: bool, + pub(crate) stdout_tail: String, + pub(crate) stderr_tail: String, +} + +pub(crate) fn run_validate_llama_load(args: ValidateLlamaLoadArgs) -> Result<()> { + let report = validate_llama_load( + &args.model, + args.llama_cli.as_deref(), + LlamaLoadOptions { + check_tensors: args.check_tensors, + }, + )?; + if args.json { + print_json_pretty(&report)?; + } else { + print_success(format!( + "llama load valid: model={} llama_cli={} status={}", + report.model.display(), + report.llama_cli.display(), + report + .status_code + .map_or_else(|| "signal".to_string(), |code| code.to_string()) + )); + } + Ok(()) +} + +#[derive(Debug, Clone, Copy)] +pub(crate) struct LlamaLoadOptions { + pub(crate) check_tensors: bool, +} + +pub(crate) fn validate_llama_load( + model: &Path, + llama_cli: Option<&Path>, + options: LlamaLoadOptions, +) -> Result { + ensure!(model.is_file(), "model does not exist: {}", model.display()); + let llama_cli = resolve_llama_cli(llama_cli).with_context( + || "llama-cli was not found; pass --llama-cli or set SKIPPY_QUANTIZE_LLAMA_CLI", + )?; + ensure!( + llama_cli.is_file(), + "llama-cli does not exist: {}", + llama_cli.display() + ); + let command = build_llama_load_command(&llama_cli, model, options)?; + let output = Command::new(&command[0]) + .args(&command[1..]) + .output() + .with_context(|| format!("run {}", command.join(" ")))?; + let report = LlamaLoadReport { + model: model.to_path_buf(), + llama_cli, + command, + status_code: output.status.code(), + success: output.status.success(), + stdout_tail: tail_lossy(&output.stdout, 16 * 1024), + stderr_tail: tail_lossy(&output.stderr, 16 * 1024), + }; + ensure!( + report.success, + "llama-cli failed to load model {} status={:?}\nstderr_tail:\n{}", + model.display(), + report.status_code, + report.stderr_tail + ); + Ok(report) +} + +pub(crate) fn build_llama_load_command( + llama_cli: &Path, + model: &Path, + options: LlamaLoadOptions, +) -> Result> { + if is_llama_simple(llama_cli) { + ensure!( + !options.check_tensors, + "--check-tensors requires llama-cli; llama-simple only supports a load smoke" + ); + return Ok(vec![ + llama_cli.display().to_string(), + "-m".to_string(), + model.display().to_string(), + "-n".to_string(), + "1".to_string(), + " ".to_string(), + ]); + } + let mut command = vec![ + llama_cli.display().to_string(), + "--model".to_string(), + model.display().to_string(), + "--n-predict".to_string(), + "0".to_string(), + "--prompt".to_string(), + String::new(), + "--no-conversation".to_string(), + ]; + if options.check_tensors { + command.push("--check-tensors".to_string()); + } + Ok(command) +} + +fn is_llama_simple(path: &Path) -> bool { + path.file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| name == "llama-simple") +} + +fn tail_lossy(bytes: &[u8], max_bytes: usize) -> String { + let start = bytes.len().saturating_sub(max_bytes); + String::from_utf8_lossy(&bytes[start..]).into_owned() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn builds_load_command_with_optional_tensor_check() { + let command = build_llama_load_command( + Path::new("/bin/llama-cli"), + Path::new("/models/model.gguf"), + LlamaLoadOptions { + check_tensors: true, + }, + ); + + assert_eq!( + command.unwrap(), + vec![ + "/bin/llama-cli", + "--model", + "/models/model.gguf", + "--n-predict", + "0", + "--prompt", + "", + "--no-conversation", + "--check-tensors", + ] + ); + } + + #[test] + fn builds_simple_load_command_for_older_llama_builds() { + let command = build_llama_load_command( + Path::new("/bin/llama-simple"), + Path::new("/models/model.gguf"), + LlamaLoadOptions { + check_tensors: false, + }, + ) + .unwrap(); + + assert_eq!( + command, + vec![ + "/bin/llama-simple", + "-m", + "/models/model.gguf", + "-n", + "1", + " " + ] + ); + } + + #[test] + fn rejects_tensor_check_with_simple_loader() { + let error = build_llama_load_command( + Path::new("/bin/llama-simple"), + Path::new("/models/model.gguf"), + LlamaLoadOptions { + check_tensors: true, + }, + ) + .unwrap_err() + .to_string(); + + assert!(error.contains("--check-tensors requires llama-cli")); + } + + #[test] + fn keeps_bounded_output_tail() { + assert_eq!(tail_lossy(b"abcdef", 3), "def"); + assert_eq!(tail_lossy(b"abc", 16), "abc"); + } +} diff --git a/crates/skippy-quantize/src/locking.rs b/crates/skippy-quantize/src/locking.rs new file mode 100644 index 0000000000..2209e39e82 --- /dev/null +++ b/crates/skippy-quantize/src/locking.rs @@ -0,0 +1,110 @@ +use std::fs::{self, File, OpenOptions}; +use std::path::{Path, PathBuf}; + +use anyhow::{Context, Result}; + +use crate::output::print_path_event; + +pub fn with_manifest_lock( + manifest_path: &Path, + action: impl FnOnce() -> Result, +) -> Result { + let _guard = ManifestLock::acquire(manifest_path)?; + action() +} + +struct ManifestLock { + path: PathBuf, + #[allow(dead_code)] + file: File, +} + +impl ManifestLock { + fn acquire(manifest_path: &Path) -> Result { + let path = lock_path(manifest_path); + if let Some(parent) = path.parent() { + fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?; + } + let file = OpenOptions::new() + .create(true) + .truncate(false) + .read(true) + .write(true) + .open(&path) + .with_context(|| format!("open manifest lock {}", path.display()))?; + lock_file(&file).with_context(|| format!("lock manifest {}", manifest_path.display()))?; + print_path_event("🔒", "Manifest lock acquired", &path); + Ok(Self { path, file }) + } +} + +impl Drop for ManifestLock { + fn drop(&mut self) { + unlock_file(&self.file); + print_path_event("🔓", "Manifest lock released", &self.path); + } +} + +fn lock_path(manifest_path: &Path) -> PathBuf { + let mut lock_name = manifest_path.as_os_str().to_os_string(); + lock_name.push(".lock"); + PathBuf::from(lock_name) +} + +#[cfg(unix)] +fn lock_file(file: &File) -> Result<()> { + use std::io; + use std::os::fd::AsRawFd; + + let ret = unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_EX | libc::LOCK_NB) }; + if ret == 0 { + return Ok(()); + } + let err = io::Error::last_os_error(); + if err.kind() == io::ErrorKind::WouldBlock { + anyhow::bail!("another skippy-quantize process holds this manifest lock"); + } + Err(err).context("flock failed") +} + +#[cfg(not(unix))] +fn lock_file(_file: &File) -> Result<()> { + Ok(()) +} + +#[cfg(unix)] +fn unlock_file(file: &File) { + use std::os::fd::AsRawFd; + + let _ = unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_UN) }; +} + +#[cfg(not(unix))] +fn unlock_file(_file: &File) {} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use crate::records::unix_timestamp_ms; + + use super::*; + + #[test] + fn manifest_lock_holds_during_action() { + let root = + std::env::temp_dir().join(format!("skippy-quantize-lock-test-{}", unix_timestamp_ms())); + let manifest = root.join("job.json"); + let calls = AtomicUsize::new(0); + + with_manifest_lock(&manifest, || { + calls.fetch_add(1, Ordering::SeqCst); + assert!(lock_path(&manifest).is_file()); + Ok(()) + }) + .unwrap(); + + assert_eq!(calls.load(Ordering::SeqCst), 1); + fs::remove_dir_all(root).unwrap(); + } +} diff --git a/crates/skippy-quantize/src/main.rs b/crates/skippy-quantize/src/main.rs new file mode 100644 index 0000000000..d14fbacfb6 --- /dev/null +++ b/crates/skippy-quantize/src/main.rs @@ -0,0 +1,1527 @@ +use std::fs; +use std::path::{Path, PathBuf}; +use std::process::Command as ProcessCommand; +use std::time::Instant; + +use anyhow::{Context, Result, ensure}; +use clap::{Parser, Subcommand}; + +mod artifacts; +mod backend; +mod command_reports; +mod direct_convert; +mod direct_quantize; +mod float_convert; +mod gguf_template; +mod gguf_writer; +mod hf_checkpoint; +mod imatrix; +mod llama_load; +mod locking; +mod manifest; +mod memory_budget; +mod native_convert; +mod native_quantize; +mod output; +mod plan_convert; +mod preflight; +mod quantize; +mod records; +mod residency; +mod splits; +mod tensor_map; +mod tokenizer_metadata; +mod tool_paths; +mod type_catalog; +mod types; +mod validation_commands; +mod verify; +mod verify_command; +mod window_loop; + +use artifacts::{clean_spooled_window, execution_root, publish_spooled_window}; +use backend::{ + BackendArgs, BackendKind, ensure_convert_backend, ensure_quant_backend, ensure_success, +}; +use command_reports::{ConvertWindowPlan, QuantWindowPlan}; +use direct_convert::{DirectConvertArgs, run_direct_convert}; +use direct_quantize::{DirectQuantizeArgs, run_direct_quantize}; +use hf_checkpoint::resolve_auto_output_type; +use llama_load::{ValidateLlamaLoadArgs, run_validate_llama_load}; +use locking::with_manifest_lock; +use manifest::{ + MANIFEST_VERSION, Manifest, ensure_manifest, manifest_progress, read_manifest, write_manifest, +}; +use memory_budget::{ + MemoryBudgetPlanInput, MemoryPolicy, MemorySize, effective_stream_buffer_bytes, + native_convert_stream_working_set_bytes, print_memory_budget_plan, +}; +use native_convert::{build_native_convert_command, run_native_convert}; +use native_quantize::{build_native_quantize_command, run_native_quantize}; +use output::{ + JsonEventConfig, JsonEventReporter, print_info, print_json_pretty, print_path_event, + print_success, print_warn, print_window, +}; +use plan_convert::{PlanConvertArgs, run_plan_convert}; +use preflight::run_job_preflight; +use records::{WindowRunRecordInput, unix_timestamp_ms, write_window_record}; +use residency::remove_dir_if_exists; +use splits::{ + SplitWindow, find_first_shard, next_missing_window_in_range, split_status, stage_source_window, + validate_split_window, +}; +use type_catalog::{TypeCatalogArgs, list_quants, list_tensor_types}; +use types::{ConvertOutputType, JobKind, QuantSpec}; +use validation_commands::{ + run_next_window as run_next_window_command, run_status as run_status_command, + validate_splits_command, validate_tensor_types, validate_tensor_types_command, +}; +use verify::{VerifyOnCompleteOptions, print_verify_on_complete}; +use verify_command::verify_job as run_verify_job; +use window_loop::run_window_loop; + +#[derive(Debug, Parser)] +#[command(name = "skippy-quantize")] +#[command(about = "Resumable GGUF conversion and quantization for Skippy workflows")] +struct Args { + #[command(subcommand)] + command: Command, +} + +#[derive(Debug, Subcommand)] +enum Command { + Backends(BackendArgs), + ListQuants(TypeCatalogArgs), + ListTensorTypes(TypeCatalogArgs), + InitQuant(InitQuantArgs), + InitConvert(InitConvertArgs), + Convert(DirectConvertArgs), + PlanConvert(PlanConvertArgs), + Quantize(DirectQuantizeArgs), + QuantizeLayerPackage(QuantizeLayerPackageArgs), + ConvertJob(ConvertJobArgs), + QuantJob(QuantJobArgs), + Status(StatusArgs), + NextWindow(NextWindowArgs), + RunConvert(RunConvertArgs), + RunConvertWindow(RunConvertWindowArgs), + RunQuant(RunQuantArgs), + RunQuantWindow(RunQuantWindowArgs), + VerifyJob(VerifyJobArgs), + ValidateLlamaLoad(ValidateLlamaLoadArgs), + ValidateTensorTypes(ValidateTensorTypesArgs), + ValidateSplits(ValidateSplitsArgs), +} + +#[derive(Debug, Parser)] +struct InitQuantArgs { + #[arg(long)] + source: PathBuf, + #[arg(long)] + source_prefix: String, + #[arg(long)] + target: PathBuf, + #[arg(long)] + target_prefix: String, + #[arg(long)] + output_basename: String, + #[arg(long)] + quant: QuantSpec, + #[arg(long)] + tensor_type_file: Option, + #[arg(long, default_value_t = 1)] + window_size: u32, + #[arg(long)] + manifest: PathBuf, +} + +#[derive(Debug, Parser)] +struct InitConvertArgs { + #[arg(long)] + source: PathBuf, + #[arg(long)] + target: PathBuf, + #[arg(long)] + target_prefix: String, + #[arg(long)] + output_basename: String, + #[arg(long, value_enum, default_value_t = ConvertOutputType::Bf16)] + output_type: ConvertOutputType, + #[arg(long)] + expected_splits: u32, + #[arg(long, default_value_t = 1)] + window_size: u32, + #[arg(long)] + manifest: PathBuf, +} + +#[derive(Debug, Parser)] +struct ConvertJobArgs { + #[command(flatten)] + init: InitConvertArgs, + #[command(flatten)] + run: ConvertJobRunArgs, +} + +#[derive(Debug, Parser)] +struct ConvertJobRunArgs { + #[command(flatten)] + runner: ConvertRunnerArgs, + #[arg(long)] + max_windows: Option, + #[arg(long)] + preflight_only: bool, + #[arg(long = "no-verify-on-complete", action = clap::ArgAction::SetFalse, default_value_t = true)] + verify_on_complete: bool, + #[command(flatten)] + verify_load: VerifyLoadArgs, + #[arg(long)] + json: bool, +} + +#[derive(Debug, Parser)] +struct QuantJobArgs { + #[command(flatten)] + init: InitQuantArgs, + #[command(flatten)] + run: QuantJobRunArgs, +} + +#[derive(Debug, Parser)] +struct QuantJobRunArgs { + #[command(flatten)] + runner: QuantRunnerArgs, + #[arg(long)] + max_windows: Option, + #[arg(long)] + preflight_only: bool, + #[arg(long = "no-verify-on-complete", action = clap::ArgAction::SetFalse, default_value_t = true)] + verify_on_complete: bool, + #[command(flatten)] + verify_load: VerifyLoadArgs, + #[arg(long)] + json: bool, +} + +#[derive(Debug, Parser)] +struct QuantizeLayerPackageArgs { + #[command(flatten)] + init: InitQuantArgs, + #[command(flatten)] + runner: QuantRunnerArgs, + #[arg(long)] + package_dir: PathBuf, + #[arg(long)] + package_model_id: String, + #[arg(long)] + package_source_repo: String, + #[arg(long)] + package_source_revision: String, + #[arg(long)] + package_source_file: Option, + #[arg(long, default_value = "target/release/skippy-model-package")] + skippy_model_package_bin: PathBuf, + #[arg(long)] + stages: Option, + #[arg(long)] + keep_quant: bool, + #[arg(long)] + replace_package: bool, + #[arg(long)] + json: bool, +} + +#[derive(Debug, Parser, Clone)] +pub(crate) struct VerifyLoadArgs { + #[arg(long = "verify-llama-load")] + llama_load: bool, + #[arg(long = "verify-llama-cli")] + llama_cli: Option, + #[arg(long = "verify-check-tensors")] + check_tensors: bool, +} + +impl VerifyLoadArgs { + pub(crate) fn options(&self, enabled: bool) -> VerifyOnCompleteOptions<'_> { + VerifyOnCompleteOptions { + enabled, + llama_load: self.llama_load, + llama_cli: self.llama_cli.as_deref(), + check_tensors: self.check_tensors, + } + } +} + +#[derive(Debug, Parser, Clone)] +struct ConvertRunnerArgs { + #[arg(long, value_enum, default_value_t = BackendKind::NativeRust)] + backend: BackendKind, + #[arg(long, default_value = "0")] + split_max_size: String, + #[arg(long)] + split_max_tensors: Option, + #[arg(long)] + skip_output_shards_before: Option, + #[arg(long)] + stop_output_shards_after: Option, + #[arg(long)] + remote: bool, + #[arg(long)] + vocab_only: bool, + #[arg(long)] + bigendian: bool, + #[arg(long)] + verbose: bool, + #[arg(long)] + dry_run: bool, + #[arg(long)] + use_temp_file: bool, + #[arg(long)] + no_lazy: bool, + #[arg(long)] + model_name: Option, + #[arg(long)] + no_tensor_first_split: bool, + #[arg(long)] + metadata: Option, + #[arg(long)] + print_supported_models: bool, + #[arg(long)] + mmproj: bool, + #[arg(long)] + mtp: bool, + #[arg(long)] + no_mtp: bool, + #[arg(long)] + mistral_format: bool, + #[arg(long)] + disable_mistral_community_chat_template: bool, + #[arg(long)] + sentence_transformers_dense_modules: bool, + #[arg(long)] + fuse_gate_up_exps: bool, + #[arg(long)] + fp8_as_q8: bool, + #[arg(long)] + target_model_dir: Option, + #[arg(long)] + spool_dir: Option, + #[arg(long)] + keep_spool: bool, + #[arg(long)] + watchdog_seconds: Option, + #[arg(long)] + max_memory: Option, + #[arg(long, value_enum, default_value_t = MemoryPolicy::Hard)] + memory_policy: MemoryPolicy, + #[arg(long, default_value_t = 8 * 1024 * 1024)] + stream_buffer_bytes: usize, + #[arg(long)] + print_only: bool, + #[arg(long)] + record_dir: Option, + #[arg(long)] + json_event_file: Option, + #[arg(long, default_value_t = 120)] + json_event_interval_seconds: u64, + #[arg(long, default_value_t = 8)] + json_event_window: usize, +} + +#[derive(Debug, Parser)] +struct StatusArgs { + #[arg(long)] + manifest: PathBuf, + #[arg(long)] + json: bool, +} + +#[derive(Debug, Parser)] +struct NextWindowArgs { + #[arg(long)] + manifest: PathBuf, + #[arg(long)] + json: bool, +} + +#[derive(Debug, Parser, Clone)] +struct RunConvertWindowArgs { + #[arg(long)] + manifest: PathBuf, + #[command(flatten)] + runner: ConvertRunnerArgs, + #[arg(long)] + json: bool, +} + +#[derive(Debug, Parser)] +struct RunConvertArgs { + #[command(flatten)] + window: RunConvertWindowArgs, + #[arg(long)] + max_windows: Option, +} + +#[derive(Debug, Parser, Clone)] +struct RunQuantWindowArgs { + #[arg(long)] + manifest: PathBuf, + #[command(flatten)] + runner: QuantRunnerArgs, + #[arg(long)] + json: bool, +} + +#[derive(Debug, Parser, Clone)] +struct QuantRunnerArgs { + #[arg(long, value_enum, default_value_t = BackendKind::LlamaApi)] + backend: BackendKind, + /// Optional dynamic llama.cpp runtime libraries for development builds. + /// + /// The normal skippy-quantize build statically links the pinned llama.cpp + /// quantization ABI and does not require this flag. + #[arg(long = "native-runtime-library", value_name = "PATH")] + native_runtime_libraries: Vec, + #[arg(long, default_value = "/tmp/skippy-quantize-work")] + work_dir: PathBuf, + #[arg(long)] + print_only: bool, + #[arg(long)] + dry_run: bool, + #[arg(long)] + allow_requantize: bool, + #[arg(long)] + pure: bool, + #[arg(long)] + imatrix: Option, + #[arg(long)] + include_weights: Vec, + #[arg(long)] + exclude_weights: Vec, + #[arg(long)] + output_tensor_type: Option, + #[arg(long)] + token_embedding_type: Option, + #[arg(long)] + tensor_type: Vec, + #[arg(long)] + prune_layers: Option, + #[arg(long)] + override_kv: Vec, + #[arg(long)] + nthreads: Option, + #[arg(long)] + leave_output_tensor: bool, + #[arg(long)] + no_stage_source: bool, + #[arg(long)] + keep_staged_source: bool, + #[arg(long)] + spool_dir: Option, + #[arg(long)] + keep_spool: bool, + #[arg(long)] + watchdog_seconds: Option, + #[arg(long)] + max_memory: Option, + #[arg(long, value_enum, default_value_t = MemoryPolicy::Hard)] + memory_policy: MemoryPolicy, + #[arg(long)] + record_dir: Option, + #[arg(long)] + json_event_file: Option, + #[arg(long, default_value_t = 120)] + json_event_interval_seconds: u64, + #[arg(long, default_value_t = 8)] + json_event_window: usize, +} + +#[derive(Debug, Parser)] +struct RunQuantArgs { + #[command(flatten)] + window: RunQuantWindowArgs, + #[arg(skip)] + window_override: Option, + #[arg(long)] + max_windows: Option, +} + +#[derive(Debug, Parser)] +struct ValidateTensorTypesArgs { + file: PathBuf, + #[arg(long)] + json: bool, +} + +#[derive(Debug, Parser)] +struct VerifyJobArgs { + #[arg(long)] + manifest: PathBuf, + #[arg(long)] + llama_load: bool, + #[arg(long)] + llama_cli: Option, + #[arg(long)] + check_tensors: bool, + #[arg(long)] + json: bool, +} + +#[derive(Debug, Parser)] +struct ValidateSplitsArgs { + #[arg(long)] + root: PathBuf, + #[arg(long)] + prefix: String, + #[arg(long)] + expected_splits: Option, + #[arg(long)] + basename: Option, + #[arg(long)] + json: bool, +} + +fn main() -> Result<()> { + match Args::parse().command { + Command::Backends(args) => backend::run_backends(args), + Command::ListQuants(args) => list_quants(args), + Command::ListTensorTypes(args) => list_tensor_types(args), + Command::InitQuant(args) => init_quant(args), + Command::InitConvert(args) => init_convert(args), + Command::Convert(args) => run_direct_convert(args), + Command::PlanConvert(args) => run_plan_convert(args), + Command::Quantize(args) => run_direct_quantize(args), + Command::QuantizeLayerPackage(args) => quantize_layer_package(args), + Command::ConvertJob(args) => convert_job(args), + Command::QuantJob(args) => quant_job(args), + Command::Status(args) => run_status_command(&args.manifest, args.json), + Command::NextWindow(args) => run_next_window_command(&args.manifest, args.json), + Command::RunConvert(args) => run_convert(args), + Command::RunConvertWindow(args) => run_convert_window(args), + Command::RunQuant(args) => run_quant(args), + Command::RunQuantWindow(args) => run_quant_window(args), + Command::VerifyJob(args) => run_verify_job( + &args.manifest, + args.llama_load, + args.llama_cli.as_deref(), + args.check_tensors, + args.json, + ), + Command::ValidateLlamaLoad(args) => run_validate_llama_load(args), + Command::ValidateTensorTypes(args) => validate_tensor_types_command(&args.file, args.json), + Command::ValidateSplits(args) => validate_splits_command( + &args.root, + &args.prefix, + args.expected_splits, + args.basename.as_deref(), + args.json, + ), + } +} + +pub(crate) fn prepare_convert_runner(runner: ConvertRunnerArgs) -> Result { + ensure_convert_backend(runner.backend)?; + ensure!( + !(runner.mtp && runner.no_mtp), + "--mtp and --no-mtp are mutually exclusive" + ); + ensure!( + runner.stream_buffer_bytes > 0, + "--stream-buffer-bytes must be greater than zero" + ); + if runner.backend == BackendKind::NativeRust { + ensure_native_convert_runner_supported(&runner)?; + } + Ok(runner) +} + +fn ensure_native_convert_runner_supported(runner: &ConvertRunnerArgs) -> Result<()> { + ensure!( + !runner.remote, + "--remote is not supported by the native converter" + ); + ensure!( + !runner.vocab_only, + "--vocab-only is not supported by the native converter" + ); + ensure!( + !runner.bigendian, + "--bigendian is not supported by the native converter" + ); + ensure!( + !runner.use_temp_file, + "--use-temp-file is not supported by the native converter" + ); + ensure!( + !runner.no_lazy, + "--no-lazy is not supported by the native converter" + ); + ensure!( + runner.model_name.is_none(), + "--model-name is not supported by the native converter" + ); + ensure!( + !runner.no_tensor_first_split, + "--no-tensor-first-split is not supported by the native converter" + ); + ensure!( + runner.metadata.is_none(), + "--metadata is not supported by the native converter" + ); + ensure!( + !runner.print_supported_models, + "--print-supported-models is not supported by the native converter" + ); + ensure!( + !runner.mmproj, + "--mmproj is not supported by the native converter" + ); + ensure!( + !runner.mistral_format, + "--mistral-format is not supported by the native converter" + ); + ensure!( + !runner.disable_mistral_community_chat_template, + "--disable-mistral-community-chat-template is not supported by the native converter" + ); + ensure!( + !runner.sentence_transformers_dense_modules, + "--sentence-transformers-dense-modules is not supported by the native converter" + ); + ensure!( + !runner.fuse_gate_up_exps, + "--fuse-gate-up-exps is not supported by the native converter" + ); + ensure!( + !runner.fp8_as_q8, + "--fp8-as-q8 is not supported by the native converter" + ); + ensure!( + runner.target_model_dir.is_none(), + "--target-model-dir is not supported by the native converter" + ); + Ok(()) +} + +impl ConvertRunnerArgs { + fn has_upstream_shard_controls(&self) -> bool { + self.skip_output_shards_before.is_some() || self.stop_output_shards_after.is_some() + } +} + +pub(crate) fn prepare_quant_runner(runner: QuantRunnerArgs) -> Result { + ensure_quant_backend(runner.backend)?; + Ok(runner) +} + +pub(crate) fn quant_backend_path(runner: &QuantRunnerArgs) -> Option<&Path> { + match runner.backend { + BackendKind::LlamaApi => runner + .native_runtime_libraries + .first() + .map(PathBuf::as_path), + BackendKind::NativeRust => None, + BackendKind::SkippyAbi => runner + .native_runtime_libraries + .first() + .map(PathBuf::as_path), + } +} + +fn init_quant(args: InitQuantArgs) -> Result<()> { + let manifest = quant_manifest_from_args(&args)?; + write_manifest(&args.manifest, &manifest) +} + +fn quant_job(args: QuantJobArgs) -> Result<()> { + let manifest = quant_manifest_from_args(&args.init)?; + let manifest_path = args.init.manifest.clone(); + let runner = prepare_quant_runner(args.run.runner)?; + if args.run.preflight_only { + return run_job_preflight( + &manifest_path, + &manifest, + Some((&args.init.source, &args.init.source_prefix)), + None, + runner.backend, + quant_backend_path(&runner), + args.run.json, + ); + } + if runner.dry_run { + return run_quant_window_once_with_manifest( + &RunQuantWindowArgs { + manifest: manifest_path, + runner, + json: args.run.json, + }, + &manifest, + None, + ) + .map(|_| ()); + } + let verify_options = args.run.verify_load.options(args.run.verify_on_complete); + with_manifest_lock(&manifest_path, || { + ensure_manifest(&manifest_path, &manifest)?; + run_quant_unlocked(RunQuantArgs { + window: RunQuantWindowArgs { + manifest: manifest_path.clone(), + runner, + json: args.run.json, + }, + window_override: None, + max_windows: args.run.max_windows, + })?; + print_verify_on_complete(&manifest_path, verify_options) + }) +} + +fn quantize_layer_package(args: QuantizeLayerPackageArgs) -> Result<()> { + ensure!( + args.init.window_size == 1, + "quantize-layer-package currently requires --window-size 1" + ); + ensure!( + !args.runner.dry_run, + "quantize-layer-package does not support --dry-run; use quant-job --preflight-only first" + ); + ensure!( + !args.runner.print_only, + "quantize-layer-package does not support --print-only; use quant-job --preflight-only first" + ); + ensure!( + args.skippy_model_package_bin.is_file(), + "missing skippy-model-package binary {}; build it with `cargo build --release --locked -p skippy-model-package` or pass --skippy-model-package-bin", + args.skippy_model_package_bin.display() + ); + if args.package_dir.exists() { + ensure!( + args.replace_package, + "package dir already exists: {}; pass --replace-package to overwrite it", + args.package_dir.display() + ); + fs::remove_dir_all(&args.package_dir) + .with_context(|| format!("remove package dir {}", args.package_dir.display()))?; + } + + let manifest = quant_manifest_from_args(&args.init)?; + let manifest_path = args.init.manifest.clone(); + let runner = prepare_quant_runner(args.runner.clone())?; + with_manifest_lock(&manifest_path, || { + ensure_manifest(&manifest_path, &manifest)?; + let hook = write_layer_package_quant_hook(&args, &runner)?; + write_and_preflight_layer_package(&args, &manifest, &hook)?; + if !args.keep_quant { + remove_dir_if_exists(&manifest.target)?; + print_path_event("🧹", "Cleaned quant scratch", &manifest.target); + } + Ok(()) + }) +} + +fn write_and_preflight_layer_package( + args: &QuantizeLayerPackageArgs, + manifest: &Manifest, + hook: &Path, +) -> Result<()> { + let first_source_shard = find_first_shard( + &manifest.source, + manifest + .source_prefix + .as_deref() + .context("quantize manifest is missing source_prefix")?, + )?; + run_skippy_model_package_write(args, &first_source_shard, hook)?; + run_skippy_model_package_preflight(args) +} + +fn write_layer_package_quant_hook( + args: &QuantizeLayerPackageArgs, + runner: &QuantRunnerArgs, +) -> Result { + let hook_dir = runner.work_dir.join("layer-package-hook"); + fs::create_dir_all(&hook_dir) + .with_context(|| format!("create hook directory {}", hook_dir.display()))?; + fs::create_dir_all(&args.init.target) + .with_context(|| format!("create quant scratch {}", args.init.target.display()))?; + let hook = hook_dir.join("quantize-package-artifact.sh"); + let current_exe = std::env::current_exe().context("resolve current skippy-quantize path")?; + let hook_work_dir = args.init.target.join("hook-work"); + let hook_spool_dir = args.init.target.join("hook-spool"); + let hook_record_dir = args.init.target.join("hook-records"); + let hook_status_file = args.init.target.join("hook-status.json"); + let mut script = String::new(); + script.push_str("#!/usr/bin/env bash\nset -euo pipefail\n"); + script.push_str("case \"${SKIPPY_PACKAGE_ARTIFACT_RELATIVE_PATH:-}\" in\n"); + script.push_str(" shared/metadata.gguf) exit 0 ;;\n"); + script.push_str("esac\n"); + script.push_str("artifact=\"${SKIPPY_PACKAGE_ARTIFACT_PATH:?}\"\n"); + script.push_str("tmp=\"${artifact}.quant-tmp.gguf\"\n"); + script.push_str("single_source=\"${artifact}.quant-src\"\n"); + script.push_str("rm -rf \"$single_source\" \"$tmp\"\n"); + script.push_str("mkdir -p \"$single_source\"\n"); + script.push_str("ln -s \"$artifact\" \"$single_source/model.gguf\"\n"); + script.push_str(&format!( + "{} quantize --backend {} --source-prefix '' --target-prefix '' --work-dir {} --spool-dir {} --record-dir {} --json-event-file {} --json-event-interval-seconds {} --json-event-window {} --no-verify-on-complete", + shell_quote(¤t_exe), + runner.backend.as_str(), + shell_quote(&hook_work_dir), + shell_quote(&hook_spool_dir), + shell_quote(&hook_record_dir), + shell_quote(&hook_status_file), + runner.json_event_interval_seconds, + runner.json_event_window, + )); + if let Some(watchdog_seconds) = runner.watchdog_seconds { + script.push_str(&format!(" --watchdog-seconds {watchdog_seconds}")); + } + if let Some(nthreads) = runner.nthreads { + script.push_str(&format!(" --nthreads {nthreads}")); + } + append_optional_tensor_type_file(&mut script, args.init.tensor_type_file.as_deref()); + append_override_kv_args(&mut script, &runner.override_kv); + if runner.allow_requantize { + script.push_str(" --allow-requantize"); + } + if runner.pure { + script.push_str(" --pure"); + } + if runner.leave_output_tensor { + script.push_str(" --leave-output-tensor"); + } + for library in &runner.native_runtime_libraries { + script.push_str(&format!( + " --native-runtime-library {}", + shell_quote(library) + )); + } + script.push_str(" \"$single_source/model.gguf\" \"$tmp\" "); + script.push_str(&shell_quote(args.init.quant.output_name())); + script.push('\n'); + script.push_str("rm -rf \"$single_source\"\n"); + script.push_str("if [[ ! -f \"$tmp\" && -f \"${tmp%.gguf}-00001-of-00001.gguf\" ]]; then\n"); + script.push_str(" tmp=\"${tmp%.gguf}-00001-of-00001.gguf\"\n"); + script.push_str("fi\n"); + script.push_str("mv \"$tmp\" \"$artifact\"\n"); + fs::write(&hook, script).with_context(|| format!("write hook {}", hook.display()))?; + make_executable(&hook)?; + Ok(hook) +} + +fn append_optional_tensor_type_file(script: &mut String, tensor_type_file: Option<&Path>) { + if let Some(path) = tensor_type_file { + script.push_str(&format!(" --tensor-type-file {}", shell_quote(path))); + } +} + +fn append_override_kv_args(script: &mut String, overrides: &[String]) { + for override_kv in overrides { + script.push_str(&format!(" --override-kv {}", shell_quote_str(override_kv))); + } +} + +fn run_skippy_model_package_write( + args: &QuantizeLayerPackageArgs, + first_source_shard: &Path, + hook: &Path, +) -> Result<()> { + print_path_event("📦", "Writing layer package", &args.package_dir); + let source_file = args.package_source_file.clone().unwrap_or_else(|| { + let file_name = first_source_shard + .file_name() + .and_then(|value| value.to_str()) + .unwrap_or("model.gguf"); + format!("{}/{}", args.init.source_prefix, file_name) + }); + let status = ProcessCommand::new(&args.skippy_model_package_bin) + .arg("write-package") + .arg(first_source_shard) + .arg("--out-dir") + .arg(&args.package_dir) + .arg("--after-artifact-command") + .arg(hook) + .arg("--model-id") + .arg(&args.package_model_id) + .arg("--source-repo") + .arg(&args.package_source_repo) + .arg("--source-revision") + .arg(&args.package_source_revision) + .arg("--source-file") + .arg(source_file) + .status() + .with_context(|| { + format!( + "run {} write-package", + args.skippy_model_package_bin.display() + ) + })?; + ensure!( + status.success(), + "{} write-package failed with status {status}", + args.skippy_model_package_bin.display() + ); + Ok(()) +} + +fn shell_quote(path: impl AsRef) -> String { + shell_quote_str(&path.as_ref().display().to_string()) +} + +fn shell_quote_str(raw: &str) -> String { + format!("'{}'", raw.replace('\'', "'\\''")) +} + +#[cfg(unix)] +fn make_executable(path: &Path) -> Result<()> { + use std::os::unix::fs::PermissionsExt; + + let mut permissions = fs::metadata(path) + .with_context(|| format!("read permissions {}", path.display()))? + .permissions(); + permissions.set_mode(0o755); + fs::set_permissions(path, permissions) + .with_context(|| format!("set executable permissions {}", path.display())) +} + +#[cfg(not(unix))] +fn make_executable(_path: &Path) -> Result<()> { + Ok(()) +} + +fn run_skippy_model_package_preflight(args: &QuantizeLayerPackageArgs) -> Result<()> { + print_path_event("✅", "Preflighting layer package", &args.package_dir); + let mut command = ProcessCommand::new(&args.skippy_model_package_bin); + command + .arg("preflight") + .arg(&args.package_dir) + .arg("--verify-sha256"); + if let Some(stages) = args.stages { + command.arg("--stages").arg(stages.to_string()); + } + let status = command + .status() + .with_context(|| format!("run {} preflight", args.skippy_model_package_bin.display()))?; + ensure!( + status.success(), + "{} preflight failed with status {status}", + args.skippy_model_package_bin.display() + ); + Ok(()) +} + +fn quant_manifest_from_args(args: &InitQuantArgs) -> Result { + ensure!( + args.window_size > 0, + "--window-size must be greater than zero" + ); + if let Some(path) = args.tensor_type_file.as_deref() { + validate_tensor_types(path)?; + } + args.quant + .validate_recipe_requirements(args.tensor_type_file.is_some()) + .map_err(anyhow::Error::msg)?; + + let source_status = split_status(&args.source, &args.source_prefix, None) + .with_context(|| format!("scan source {}", args.source.display()))?; + ensure!( + source_status.expected_splits > 0, + "source contains no split GGUF shards under prefix {:?}", + args.source_prefix + ); + ensure!( + source_status.complete, + "source split is incomplete: {}/{} shards present missing_ranges={:?}", + source_status.completed_count, + source_status.expected_splits, + source_status.missing_ranges + ); + + let manifest = Manifest { + schema_version: MANIFEST_VERSION, + kind: JobKind::QuantizeGguf, + source: args.source.clone(), + source_prefix: Some(args.source_prefix.clone()), + target: args.target.clone(), + target_prefix: args.target_prefix.clone(), + output_basename: args.output_basename.clone(), + expected_splits: source_status.expected_splits, + window_size: args.window_size, + quant: Some(args.quant.base_quant().as_llama_name().to_string()), + output_type: None, + tensor_type_file: args.tensor_type_file.clone(), + tensor_type_recipe: None, + }; + Ok(manifest) +} + +fn init_convert(args: InitConvertArgs) -> Result<()> { + let manifest = convert_manifest_from_args(&args)?; + write_manifest(&args.manifest, &manifest) +} + +fn convert_job(args: ConvertJobArgs) -> Result<()> { + let manifest = convert_manifest_from_args(&args.init)?; + let manifest_path = args.init.manifest.clone(); + let runner = prepare_convert_runner(args.run.runner)?; + if args.run.preflight_only { + return run_job_preflight( + &manifest_path, + &manifest, + None, + None, + runner.backend, + None, + args.run.json, + ); + } + if runner.dry_run { + return run_convert_window_once_with_manifest( + &RunConvertWindowArgs { + manifest: manifest_path, + runner, + json: args.run.json, + }, + &manifest, + ) + .map(|_| ()); + } + let verify_options = args.run.verify_load.options(args.run.verify_on_complete); + with_manifest_lock(&manifest_path, || { + ensure_manifest(&manifest_path, &manifest)?; + run_convert_unlocked(RunConvertArgs { + window: RunConvertWindowArgs { + manifest: manifest_path.clone(), + runner, + json: args.run.json, + }, + max_windows: args.run.max_windows, + })?; + print_verify_on_complete(&manifest_path, verify_options) + }) +} + +fn convert_manifest_from_args(args: &InitConvertArgs) -> Result { + ensure!( + args.expected_splits > 0, + "--expected-splits must be greater than zero" + ); + ensure!( + args.window_size > 0, + "--window-size must be greater than zero" + ); + let manifest = Manifest { + schema_version: MANIFEST_VERSION, + kind: JobKind::ConvertHf, + source: args.source.clone(), + source_prefix: None, + target: args.target.clone(), + target_prefix: args.target_prefix.clone(), + output_basename: args.output_basename.clone(), + expected_splits: args.expected_splits, + window_size: args.window_size, + quant: None, + output_type: Some(args.output_type), + tensor_type_file: None, + tensor_type_recipe: None, + }; + Ok(manifest) +} + +fn run_convert(args: RunConvertArgs) -> Result<()> { + let manifest_path = args.window.manifest.clone(); + with_manifest_lock(&manifest_path, || run_convert_unlocked(args)) +} + +pub(crate) fn run_convert_unlocked(args: RunConvertArgs) -> Result<()> { + ensure!( + !args.window.runner.print_only, + "run-convert does not support --print-only; use run-convert-window" + ); + if args.window.runner.dry_run { + return run_convert_window_once(&args.window).map(|_| ()); + } + run_window_loop("convert", args.max_windows, || { + run_convert_window_once(&args.window) + }) +} + +fn run_convert_window(args: RunConvertWindowArgs) -> Result<()> { + with_manifest_lock(&args.manifest, || { + run_convert_window_once(&args).map(|_| ()) + }) +} + +fn run_convert_window_once(args: &RunConvertWindowArgs) -> Result { + let manifest = read_manifest(&args.manifest)?; + run_convert_window_once_with_manifest(args, &manifest) +} + +pub(crate) fn run_convert_window_once_with_manifest( + args: &RunConvertWindowArgs, + manifest: &Manifest, +) -> Result { + ensure!( + manifest.kind == JobKind::ConvertHf, + "run-convert-window requires a convert manifest" + ); + let runner = prepare_convert_runner(args.runner.clone())?; + ensure!( + !runner.has_upstream_shard_controls(), + "run-convert-window owns shard selection; use direct convert passthrough for --skip-output-shards-before/--stop-output-shards-after" + ); + + let progress = manifest_progress(manifest)?; + let Some(window) = progress.next_window else { + if args.json { + print_json_pretty(&serde_json::json!({ + "event": "convert_windows_complete", + "completed": true, + }))?; + } else { + print_success("convert windows complete"); + } + return Ok(false); + }; + let event_reporter = + JsonEventReporter::start(convert_json_event_config(&runner), "convert", Some(window))?; + event_reporter.record("selected conversion window")?; + + let output_root = execution_root( + &manifest.target, + &manifest.target_prefix, + runner.spool_dir.as_deref(), + ); + let output_prefix = output_root.join(format!("{}.gguf", manifest.output_basename)); + let command = match runner.backend { + BackendKind::NativeRust => { + build_native_convert_command(&runner, manifest, &output_prefix, window) + } + BackendKind::LlamaApi | BackendKind::SkippyAbi => { + unreachable!("unsupported convert backend checked earlier") + } + }; + let plan = ConvertWindowPlan { + first_split: window.first_split, + last_split: window.last_split, + output_prefix, + command, + }; + if args.json { + print_json_pretty(&serde_json::json!({ + "event": "convert_window", + "plan": plan, + }))?; + } else { + print_window("convert window", window); + print_info(format!("Output prefix: {}", plan.output_prefix.display())); + print_info(format!("Command: {}", plan.command.join(" "))); + } + event_reporter.record("conversion plan ready")?; + let stream_buffer_bytes = (runner.backend == BackendKind::NativeRust) + .then(|| effective_stream_buffer_bytes(runner.stream_buffer_bytes, runner.max_memory)) + .transpose()?; + let native_output_type = if runner.backend == BackendKind::NativeRust { + manifest + .output_type + .map(|output_type| resolve_auto_output_type(&manifest.source, output_type)) + .transpose()? + } else { + manifest.output_type + }; + let estimated_stream_working_set_bytes = stream_buffer_bytes + .map(|buffer_size| native_convert_stream_working_set_bytes(buffer_size, native_output_type)) + .transpose()?; + print_memory_budget_plan(MemoryBudgetPlanInput { + kind: "convert", + backend: runner.backend.as_str(), + max_memory: runner.max_memory, + memory_policy: runner.memory_policy, + watchdog_seconds: runner.watchdog_seconds, + window, + stream_buffer_bytes, + estimated_stream_working_set_bytes, + llama_quantize_env_bytes: None, + json: args.json, + })?; + if runner.dry_run { + print_dry_run_complete(args.json, "convert")?; + event_reporter.record("dry run complete")?; + event_reporter.finish("dry_run")?; + return Ok(true); + } + if runner.print_only { + event_reporter.record("print only complete")?; + event_reporter.finish("planned")?; + return Ok(true); + } + event_reporter.set_phase("preparing")?; + fs::create_dir_all(&output_root) + .with_context(|| format!("create {}", output_root.display()))?; + clean_spooled_window( + runner.spool_dir.as_deref(), + &manifest.target_prefix, + &manifest.output_basename, + manifest.expected_splits, + window, + )?; + let started_unix_ms = unix_timestamp_ms(); + let started = Instant::now(); + event_reporter.set_phase("running")?; + event_reporter.record("native conversion started")?; + event_reporter.write_now()?; + let status = match runner.backend { + BackendKind::NativeRust => { + run_native_convert(&runner, manifest, window, &plan.output_prefix)? + } + BackendKind::LlamaApi | BackendKind::SkippyAbi => { + unreachable!("unsupported convert backend checked earlier") + } + }; + let duration_ms = started.elapsed().as_millis(); + write_window_record( + runner.record_dir.as_deref(), + WindowRunRecordInput { + schema_version: MANIFEST_VERSION, + kind: manifest.kind, + command: &plan.command, + output_prefix: &plan.output_prefix, + window, + status, + duration_ms, + started_unix_ms, + }, + )?; + ensure_success(status, &plan.command)?; + event_reporter.record("native conversion finished")?; + if !runner.dry_run { + event_reporter.set_phase("publishing")?; + publish_spooled_window( + runner.spool_dir.as_deref(), + &manifest.target, + &manifest.target_prefix, + &manifest.output_basename, + manifest.expected_splits, + window, + args.runner.keep_spool, + )?; + event_reporter.record("conversion window published")?; + } + event_reporter.finish("complete")?; + Ok(true) +} + +fn run_quant(args: RunQuantArgs) -> Result<()> { + let manifest_path = args.window.manifest.clone(); + with_manifest_lock(&manifest_path, || run_quant_unlocked(args)) +} + +pub(crate) fn run_quant_unlocked(args: RunQuantArgs) -> Result<()> { + ensure!( + !args.window.runner.print_only, + "run-quant does not support --print-only; use run-quant-window" + ); + if args.window.runner.dry_run { + return run_quant_window_once(&args.window, args.window_override).map(|_| ()); + } + run_window_loop("quant", args.max_windows, || { + run_quant_window_once(&args.window, args.window_override) + }) +} + +fn run_quant_window(args: RunQuantWindowArgs) -> Result<()> { + with_manifest_lock(&args.manifest, || { + run_quant_window_once(&args, None).map(|_| ()) + }) +} + +fn run_quant_window_once( + args: &RunQuantWindowArgs, + window_override: Option, +) -> Result { + let manifest = read_manifest(&args.manifest)?; + run_quant_window_once_with_manifest(args, &manifest, window_override) +} + +pub(crate) fn run_quant_window_once_with_manifest( + args: &RunQuantWindowArgs, + manifest: &Manifest, + window_override: Option, +) -> Result { + ensure!( + manifest.kind == JobKind::QuantizeGguf, + "run-quant-window requires a quantize manifest" + ); + let runner = prepare_quant_runner(args.runner.clone())?; + + let progress = manifest_progress(manifest)?; + let window = if let Some(requested) = window_override { + validate_split_window(requested, manifest.expected_splits)?; + let Some(window) = next_missing_window_in_range(&progress.missing_ranges, requested) else { + if args.json { + print_json_pretty(&serde_json::json!({ + "event": "quant_requested_window_complete", + "window": requested, + }))?; + } else { + print_success(format!( + "requested quant window {} is already complete", + output::format_window(requested) + )); + } + return Ok(false); + }; + window + } else if let Some(window) = progress.next_window { + window + } else { + if args.json { + print_json_pretty(&serde_json::json!({ + "event": "quant_windows_complete", + "completed": true, + }))?; + } else { + print_success("quant windows complete"); + } + return Ok(false); + }; + + let source_prefix = manifest + .source_prefix + .as_deref() + .context("quantize manifest is missing source_prefix")?; + let first_source_shard = find_first_shard(&manifest.source, source_prefix)?; + let stage_path = runner.work_dir.join("source-window"); + let event_reporter = + JsonEventReporter::start(quant_json_event_config(&runner), "quant", Some(window))?; + event_reporter.record("selected quantization window")?; + let staged_first_shard = if runner.no_stage_source { + event_reporter.record("source staging skipped")?; + first_source_shard + } else if runner.dry_run || runner.print_only { + event_reporter.record("source staging planned")?; + planned_staged_first_shard( + &stage_path, + source_prefix, + &first_source_shard, + manifest.expected_splits, + )? + } else { + event_reporter.set_phase("staging")?; + event_reporter.record("source staging started")?; + event_reporter.write_now()?; + stage_source_window( + &manifest.source, + source_prefix, + &first_source_shard, + &stage_path, + window, + manifest.expected_splits, + ) + .inspect(|_| { + let _ = event_reporter.record("source staging finished"); + })? + }; + + let output_root = execution_root( + &manifest.target, + &manifest.target_prefix, + runner.spool_dir.as_deref(), + ); + let output_prefix = output_root.join(format!("{}.gguf", manifest.output_basename)); + let command = match runner.backend { + BackendKind::LlamaApi | BackendKind::SkippyAbi => build_native_quantize_command( + &runner, + manifest, + &staged_first_shard, + &output_prefix, + window, + )?, + BackendKind::NativeRust => { + unreachable!("unsupported quant backend checked earlier") + } + }; + let plan = QuantWindowPlan { + first_split: window.first_split, + last_split: window.last_split, + staged_first_shard, + output_prefix, + command, + }; + if args.json { + print_json_pretty(&serde_json::json!({ + "event": "quant_window", + "plan": plan, + }))?; + } else { + print_window("quant window", window); + print_info(format!( + "Staged first shard: {}", + plan.staged_first_shard.display() + )); + print_info(format!("Output prefix: {}", plan.output_prefix.display())); + print_info(format!("Command: {}", plan.command.join(" "))); + } + event_reporter.record("quantization plan ready")?; + print_memory_budget_plan(MemoryBudgetPlanInput { + kind: "quant", + backend: runner.backend.as_str(), + max_memory: runner.max_memory, + memory_policy: runner.memory_policy, + watchdog_seconds: runner.watchdog_seconds, + window, + stream_buffer_bytes: None, + estimated_stream_working_set_bytes: None, + llama_quantize_env_bytes: runner.max_memory.map(MemorySize::bytes), + json: args.json, + })?; + + if runner.dry_run { + print_dry_run_complete(args.json, "quant")?; + event_reporter.record("dry run complete")?; + event_reporter.finish("dry_run")?; + return Ok(true); + } + if runner.print_only { + event_reporter.record("print only complete")?; + event_reporter.finish("planned")?; + return Ok(true); + } + event_reporter.set_phase("preparing")?; + fs::create_dir_all(&output_root) + .with_context(|| format!("create {}", output_root.display()))?; + clean_spooled_window( + runner.spool_dir.as_deref(), + &manifest.target_prefix, + &manifest.output_basename, + manifest.expected_splits, + window, + )?; + let started_unix_ms = unix_timestamp_ms(); + let started = Instant::now(); + event_reporter.set_phase("running")?; + event_reporter.record("native quantization started")?; + event_reporter.write_now()?; + let status = match runner.backend { + BackendKind::LlamaApi | BackendKind::SkippyAbi => run_native_quantize( + &runner, + manifest, + &plan.staged_first_shard, + &plan.output_prefix, + window, + )?, + BackendKind::NativeRust => { + unreachable!("unsupported quant backend checked earlier") + } + }; + let duration_ms = started.elapsed().as_millis(); + write_window_record( + runner.record_dir.as_deref(), + WindowRunRecordInput { + schema_version: MANIFEST_VERSION, + kind: manifest.kind, + command: &plan.command, + output_prefix: &plan.output_prefix, + window, + status, + duration_ms, + started_unix_ms, + }, + )?; + ensure_success(status, &plan.command)?; + event_reporter.record("native quantization finished")?; + if !runner.dry_run { + event_reporter.set_phase("publishing")?; + publish_spooled_window( + runner.spool_dir.as_deref(), + &manifest.target, + &manifest.target_prefix, + &manifest.output_basename, + manifest.expected_splits, + window, + args.runner.keep_spool, + )?; + event_reporter.record("quantization window published")?; + } + if !runner.no_stage_source && !runner.keep_staged_source { + event_reporter.set_phase("cleanup")?; + remove_dir_if_exists(&stage_path)?; + print_path_event("🧹", "Cleaned staged source", &stage_path); + event_reporter.record("staged source cleaned")?; + } + event_reporter.finish("complete")?; + Ok(true) +} + +fn convert_json_event_config(runner: &ConvertRunnerArgs) -> JsonEventConfig { + JsonEventConfig { + file: runner.json_event_file.clone(), + interval_seconds: runner.json_event_interval_seconds, + window_size: runner.json_event_window, + } +} + +fn quant_json_event_config(runner: &QuantRunnerArgs) -> JsonEventConfig { + JsonEventConfig { + file: runner.json_event_file.clone(), + interval_seconds: runner.json_event_interval_seconds, + window_size: runner.json_event_window, + } +} + +fn planned_staged_first_shard( + stage_path: &Path, + source_prefix: &str, + first_source_shard: &Path, + total: u32, +) -> Result { + Ok(stage_path + .join(source_prefix) + .join(splits::shard_name_for(first_source_shard, 1, total)?)) +} + +fn print_dry_run_complete(json: bool, kind: &str) -> Result<()> { + if json { + print_json_pretty(&serde_json::json!({ + "event": "dry_run", + "kind": kind, + "executed": false, + }))?; + } else { + print_warn(format!( + "{kind} dry run: no files were written, cleaned, recorded, or published" + )); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::path::Path; + + use super::{append_optional_tensor_type_file, append_override_kv_args}; + + #[test] + fn layer_package_hook_forwards_tensor_type_file() { + let mut script = String::new(); + + append_optional_tensor_type_file( + &mut script, + Some(Path::new("/tmp/recipes/glm 5.2 q2.tensor-types.txt")), + ); + + assert!(script.contains("--tensor-type-file")); + assert!(script.contains("'/tmp/recipes/glm 5.2 q2.tensor-types.txt'")); + } + + #[test] + fn layer_package_hook_forwards_metadata_overrides() { + let mut script = String::new(); + + append_override_kv_args( + &mut script, + &["glm-dsa.attention.indexer.head_count=int:32".to_string()], + ); + + assert!(script.contains("--override-kv 'glm-dsa.attention.indexer.head_count=int:32'")); + } +} diff --git a/crates/skippy-quantize/src/manifest.rs b/crates/skippy-quantize/src/manifest.rs new file mode 100644 index 0000000000..c32df73895 --- /dev/null +++ b/crates/skippy-quantize/src/manifest.rs @@ -0,0 +1,78 @@ +use std::fs; +use std::path::{Path, PathBuf}; + +use anyhow::{Context, Result, ensure}; +use serde::{Deserialize, Serialize}; + +use crate::output::print_path_event; +use crate::splits::{Progress, next_missing_window, split_status_for_basename}; +use crate::types::{ConvertOutputType, JobKind}; + +pub const MANIFEST_VERSION: u32 = 1; + +#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct Manifest { + pub schema_version: u32, + pub kind: JobKind, + pub source: PathBuf, + pub source_prefix: Option, + pub target: PathBuf, + pub target_prefix: String, + pub output_basename: String, + pub expected_splits: u32, + pub window_size: u32, + pub quant: Option, + pub output_type: Option, + pub tensor_type_file: Option, + #[serde(default)] + pub tensor_type_recipe: Option, +} + +pub fn ensure_manifest(path: &Path, manifest: &Manifest) -> Result<()> { + if path.exists() { + let existing = read_manifest(path)?; + ensure!( + existing == *manifest, + "existing manifest does not match requested job: {}", + path.display() + ); + print_path_event("📄", "Resuming manifest", path); + return Ok(()); + } + write_manifest(path, manifest)?; + print_path_event("📄", "Created manifest", path); + Ok(()) +} + +pub fn read_manifest(path: &Path) -> Result { + let data = fs::read_to_string(path).with_context(|| format!("read {}", path.display()))?; + let manifest: Manifest = + serde_json::from_str(&data).with_context(|| format!("parse {}", path.display()))?; + ensure!( + manifest.schema_version == MANIFEST_VERSION, + "unsupported manifest schema version {}", + manifest.schema_version + ); + Ok(manifest) +} + +pub fn write_manifest(path: &Path, manifest: &Manifest) -> Result<()> { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?; + } + fs::write(path, serde_json::to_vec_pretty(manifest)?) + .with_context(|| format!("write {}", path.display())) +} + +pub fn manifest_progress(manifest: &Manifest) -> Result { + split_status_for_basename( + &manifest.target, + &manifest.target_prefix, + &manifest.output_basename, + manifest.expected_splits, + ) + .map(|mut progress| { + progress.next_window = next_missing_window(&progress.missing_ranges, manifest.window_size); + progress + }) +} diff --git a/crates/skippy-quantize/src/memory_budget.rs b/crates/skippy-quantize/src/memory_budget.rs new file mode 100644 index 0000000000..4e16f7f500 --- /dev/null +++ b/crates/skippy-quantize/src/memory_budget.rs @@ -0,0 +1,261 @@ +use std::str::FromStr; + +use anyhow::{Context, Result, ensure}; +use clap::ValueEnum; +use serde::Serialize; + +use crate::output::{format_bytes, print_info, print_json_pretty, print_warn}; +use crate::splits::SplitWindow; +use crate::types::ConvertOutputType; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, ValueEnum, Serialize)] +pub(crate) enum MemoryPolicy { + Advisory, + #[default] + Hard, +} + +impl MemoryPolicy { + pub(crate) fn is_hard(self) -> bool { + matches!(self, Self::Hard) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct MemorySize(u64); + +impl MemorySize { + pub(crate) fn bytes(self) -> u64 { + self.0 + } + + #[cfg(test)] + pub(crate) fn from_bytes_for_tests(bytes: u64) -> Self { + Self(bytes) + } +} + +impl FromStr for MemorySize { + type Err = String; + + fn from_str(raw: &str) -> Result { + let trimmed = raw.trim(); + if trimmed.is_empty() { + return Err("memory size is empty".to_string()); + } + let suffix_start = trimmed + .find(|ch: char| !ch.is_ascii_digit()) + .unwrap_or(trimmed.len()); + let (digits, suffix) = trimmed.split_at(suffix_start); + if digits.is_empty() { + return Err(format!("memory size {raw:?} is missing a number")); + } + let value = digits + .parse::() + .map_err(|err| format!("invalid memory size {raw:?}: {err}"))?; + let multiplier = match suffix.trim().to_ascii_lowercase().as_str() { + "" | "b" => 1, + "k" | "kb" | "kib" => 1024, + "m" | "mb" | "mib" => 1024 * 1024, + "g" | "gb" | "gib" => 1024 * 1024 * 1024, + "t" | "tb" | "tib" => 1024_u64.pow(4), + other => return Err(format!("unsupported memory size suffix {other:?}")), + }; + value + .checked_mul(multiplier) + .map(Self) + .ok_or_else(|| format!("memory size {raw:?} is too large")) + } +} + +#[derive(Debug, Serialize)] +struct MemoryBudgetPlan<'a> { + kind: &'a str, + backend: &'a str, + max_memory_bytes: Option, + memory_policy: MemoryPolicy, + watchdog_seconds: Option, + hard_limit: bool, + first_split: u32, + last_split: u32, + window_shards: u32, + stream_buffer_bytes: Option, + estimated_stream_working_set_bytes: Option, + llama_quantize_env_bytes: Option, +} + +pub(crate) struct MemoryBudgetPlanInput<'a> { + pub(crate) kind: &'a str, + pub(crate) backend: &'a str, + pub(crate) max_memory: Option, + pub(crate) memory_policy: MemoryPolicy, + pub(crate) watchdog_seconds: Option, + pub(crate) window: SplitWindow, + pub(crate) stream_buffer_bytes: Option, + pub(crate) estimated_stream_working_set_bytes: Option, + pub(crate) llama_quantize_env_bytes: Option, + pub(crate) json: bool, +} + +pub(crate) fn print_memory_budget_plan(input: MemoryBudgetPlanInput<'_>) -> Result<()> { + if input.max_memory.is_none() && input.watchdog_seconds.is_none() { + return Ok(()); + } + let plan = MemoryBudgetPlan { + kind: input.kind, + backend: input.backend, + max_memory_bytes: input.max_memory.map(MemorySize::bytes), + memory_policy: input.memory_policy, + watchdog_seconds: input.watchdog_seconds, + hard_limit: input.memory_policy.is_hard() && input.max_memory.is_some(), + first_split: input.window.first_split, + last_split: input.window.last_split, + window_shards: input + .window + .last_split + .saturating_sub(input.window.first_split) + .saturating_add(1), + stream_buffer_bytes: input.stream_buffer_bytes, + estimated_stream_working_set_bytes: input.estimated_stream_working_set_bytes, + llama_quantize_env_bytes: input.llama_quantize_env_bytes, + }; + if input.json { + print_json_pretty(&serde_json::json!({ + "event": format!("{}_memory_budget", input.kind), + "plan": plan, + }))?; + } else if plan.hard_limit { + print_warn(format!( + "{} memory budget: hard cap {}", + input.kind, + format_bytes(input.max_memory.map(MemorySize::bytes).unwrap_or_default()) + )); + } else { + print_info(format!("{} memory budget configured", input.kind)); + } + Ok(()) +} + +pub(crate) fn effective_stream_buffer_bytes( + requested: usize, + max_memory: Option, +) -> Result { + if let Some(max_memory) = max_memory { + let budget_limited = bytes_to_usize(max_memory)? / 64; + return Ok(requested.min(budget_limited.max(1))); + } + Ok(requested) +} + +pub(crate) fn native_convert_stream_working_set_bytes( + stream_buffer_bytes: usize, + output_type: Option, +) -> Result { + let multiplier = match output_type { + None => 1, + Some(ConvertOutputType::F16 | ConvertOutputType::Bf16) => 2, + Some(ConvertOutputType::F32) => 3, + Some(other) => { + anyhow::bail!( + "native conversion does not support output type {}", + other.as_arg() + ); + } + }; + (stream_buffer_bytes as u64) + .checked_mul(multiplier) + .context("native conversion stream working-set estimate overflow") +} + +pub(crate) fn enforce_memory_budget( + label: &str, + estimated_bytes: u64, + max_memory: Option, + policy: MemoryPolicy, +) -> Result<()> { + let Some(max_memory) = max_memory else { + return Ok(()); + }; + if estimated_bytes <= max_memory.bytes() { + return Ok(()); + } + print_warn(format!( + "{label} estimated working set {} exceeds budget {} ({policy:?})", + format_bytes(estimated_bytes), + format_bytes(max_memory.bytes()) + )); + ensure!( + !policy.is_hard(), + "{label} estimated working set {} bytes exceeds --max-memory {} bytes", + estimated_bytes, + max_memory.bytes() + ); + Ok(()) +} + +fn bytes_to_usize(value: MemorySize) -> Result { + usize::try_from(value.bytes()).context("--max-memory does not fit usize on this platform") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_memory_sizes() { + assert_eq!("1".parse::().unwrap().bytes(), 1); + assert_eq!("2K".parse::().unwrap().bytes(), 2048); + assert_eq!( + "3MiB".parse::().unwrap().bytes(), + 3 * 1024 * 1024 + ); + assert_eq!( + "4G".parse::().unwrap().bytes(), + 4 * 1024 * 1024 * 1024 + ); + } + + #[test] + fn derives_effective_stream_buffer_from_budget() { + assert_eq!(effective_stream_buffer_bytes(1024, None).unwrap(), 1024); + assert_eq!( + effective_stream_buffer_bytes(1024, Some(MemorySize::from_bytes_for_tests(128))) + .unwrap(), + 2 + ); + assert_eq!( + effective_stream_buffer_bytes(1024, Some(MemorySize::from_bytes_for_tests(1))).unwrap(), + 1 + ); + } + + #[test] + fn estimates_native_convert_stream_working_set() { + assert_eq!( + native_convert_stream_working_set_bytes(1024, None).unwrap(), + 1024 + ); + assert_eq!( + native_convert_stream_working_set_bytes(1024, Some(ConvertOutputType::Bf16)).unwrap(), + 2048 + ); + assert_eq!( + native_convert_stream_working_set_bytes(1024, Some(ConvertOutputType::F32)).unwrap(), + 3072 + ); + assert!( + native_convert_stream_working_set_bytes(1024, Some(ConvertOutputType::Q8_0)).is_err() + ); + assert!( + native_convert_stream_working_set_bytes(1024, Some(ConvertOutputType::Auto)).is_err() + ); + } + + #[test] + fn hard_memory_policy_rejects_over_budget_working_set() { + let max_memory = Some(MemorySize::from_bytes_for_tests(100)); + assert!(enforce_memory_budget("test", 100, max_memory, MemoryPolicy::Hard).is_ok()); + assert!(enforce_memory_budget("test", 101, max_memory, MemoryPolicy::Hard).is_err()); + assert!(enforce_memory_budget("test", 101, max_memory, MemoryPolicy::Advisory).is_ok()); + } +} diff --git a/crates/skippy-quantize/src/native_convert.rs b/crates/skippy-quantize/src/native_convert.rs new file mode 100644 index 0000000000..454abd3348 --- /dev/null +++ b/crates/skippy-quantize/src/native_convert.rs @@ -0,0 +1,194 @@ +use std::path::{Path, PathBuf}; + +use anyhow::{Context, Result, ensure}; + +use crate::ConvertRunnerArgs; +use crate::backend::BackendRunStatus; +use crate::gguf_template::{ + MetadataOptions, metadata_from_hf_config_with_options, mtp_layer_start_from_hf_config, +}; +use crate::gguf_writer::{ + GgufSplit, RawGgufWriteOptions, TensorSelection, write_raw_safetensors_gguf, +}; +use crate::hf_checkpoint::{inspect_hf_checkpoint, resolve_auto_output_type}; +use crate::manifest::Manifest; +use crate::memory_budget::{ + effective_stream_buffer_bytes, enforce_memory_budget, native_convert_stream_working_set_bytes, +}; +use crate::output::{format_bytes, print_info}; +use crate::splits::SplitWindow; +use crate::tensor_map::TensorNameMap; +use crate::tokenizer_metadata::ensure_native_tokenizer_metadata_supported; + +pub(crate) fn build_native_convert_command( + runner: &ConvertRunnerArgs, + manifest: &Manifest, + output_prefix: &Path, + window: SplitWindow, +) -> Vec { + let mut command = vec![ + "skippy-quantize".to_string(), + "run-convert-window".to_string(), + "--backend".to_string(), + "native-rust".to_string(), + "--source".to_string(), + manifest.source.display().to_string(), + "--outfile".to_string(), + output_prefix.display().to_string(), + "--first-split".to_string(), + window.first_split.to_string(), + "--last-split".to_string(), + window.last_split.to_string(), + "--expected-splits".to_string(), + manifest.expected_splits.to_string(), + ]; + if runner.no_mtp { + command.push("--no-mtp".to_string()); + } + if runner.mtp { + command.push("--mtp".to_string()); + } + command +} + +pub(crate) fn run_native_convert( + runner: &ConvertRunnerArgs, + manifest: &Manifest, + window: SplitWindow, + output_prefix: &Path, +) -> Result { + if runner.dry_run { + return Ok(BackendRunStatus::from_code(0)); + } + ensure!( + window.first_split <= window.last_split, + "invalid native convert window {}..{}", + window.first_split, + window.last_split + ); + ensure!( + window.last_split <= manifest.expected_splits, + "native convert window ends after expected split count" + ); + let buffer_size = effective_stream_buffer_bytes(runner.stream_buffer_bytes, runner.max_memory)?; + let output_type = manifest + .output_type + .map(|output_type| resolve_auto_output_type(&manifest.source, output_type)) + .transpose()?; + let estimated_stream_working_set_bytes = + native_convert_stream_working_set_bytes(buffer_size, output_type)?; + enforce_memory_budget( + "native_convert_stream_buffers", + estimated_stream_working_set_bytes, + runner.max_memory, + runner.memory_policy, + )?; + let plan = inspect_hf_checkpoint(&manifest.source, runner.max_memory, 0.60)?; + ensure_native_tokenizer_metadata_supported(&manifest.source)?; + let mtp_layer_start = mtp_layer_start_from_hf_config(&manifest.source)?; + let tensor_selection = native_tensor_selection(runner, mtp_layer_start)?; + let tensor_name_map = native_tensor_name_map(mtp_layer_start); + for split_index in window.first_split..=window.last_split { + let output = output_shard_path(output_prefix, split_index, manifest.expected_splits)?; + print_info(format!( + "Writing native convert shard {}/{} -> {} (buffer {}, estimated working set {})", + split_index, + manifest.expected_splits, + output.display(), + format_bytes(buffer_size as u64), + format_bytes(estimated_stream_working_set_bytes) + )); + let metadata = metadata_from_hf_config_with_options( + &manifest.source, + plan.tensor_count, + MetadataOptions { + include_mtp: !runner.no_mtp, + }, + )?; + write_raw_safetensors_gguf( + &manifest.source, + &output, + RawGgufWriteOptions { + buffer_size, + metadata: Some(metadata), + tensor_name_map, + split: split_for(split_index, manifest.expected_splits), + output_type, + tensor_selection, + }, + )?; + } + Ok(BackendRunStatus::from_code(0)) +} + +fn native_tensor_selection( + runner: &ConvertRunnerArgs, + mtp_layer_start: Option, +) -> Result { + if runner.no_mtp { + let layer_start = mtp_layer_start + .context("--no-mtp requested but config.json does not declare MTP layers")?; + return Ok(TensorSelection::ExcludeMtp { layer_start }); + } + if runner.mtp { + let layer_start = mtp_layer_start + .context("--mtp requested but config.json does not declare MTP layers")?; + return Ok(TensorSelection::MtpOnly { layer_start }); + } + Ok(TensorSelection::All) +} + +fn native_tensor_name_map(mtp_layer_start: Option) -> TensorNameMap { + mtp_layer_start.map_or(TensorNameMap::HfToGguf, |layer_start| { + TensorNameMap::HfToGgufWithMtp { layer_start } + }) +} + +fn split_for(split_index: u32, split_count: u32) -> Option { + (split_count > 1).then_some(GgufSplit { + split_index, + split_count, + }) +} + +fn output_shard_path(output_prefix: &Path, split_index: u32, split_count: u32) -> Result { + ensure!(split_count > 0, "split_count must be greater than zero"); + ensure!(split_index > 0, "split_index must be greater than zero"); + ensure!( + split_index <= split_count, + "split_index {} exceeds split_count {}", + split_index, + split_count + ); + if split_count == 1 { + return Ok(output_prefix.to_path_buf()); + } + let file_name = output_prefix + .file_name() + .and_then(|name| name.to_str()) + .with_context(|| format!("invalid output prefix {}", output_prefix.display()))?; + let base = file_name.strip_suffix(".gguf").with_context(|| { + format!( + "output prefix must end in .gguf: {}", + output_prefix.display() + ) + })?; + Ok(output_prefix.with_file_name(format!("{base}-{split_index:05}-of-{split_count:05}.gguf"))) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn derives_split_output_path_from_prefix() { + assert_eq!( + output_shard_path(Path::new("/out/model.gguf"), 2, 7).unwrap(), + PathBuf::from("/out/model-00002-of-00007.gguf") + ); + assert_eq!( + output_shard_path(Path::new("/out/model.gguf"), 1, 1).unwrap(), + PathBuf::from("/out/model.gguf") + ); + } +} diff --git a/crates/skippy-quantize/src/native_quantize.rs b/crates/skippy-quantize/src/native_quantize.rs new file mode 100644 index 0000000000..7359ae2054 --- /dev/null +++ b/crates/skippy-quantize/src/native_quantize.rs @@ -0,0 +1,861 @@ +use std::ffi::{CString, c_char}; +use std::fs; +use std::path::{Path, PathBuf}; + +use anyhow::{Context, Result, anyhow, ensure}; + +use crate::QuantRunnerArgs; +use crate::backend::BackendRunStatus; +use crate::imatrix::NativeImatrix; +use crate::manifest::Manifest; +use crate::quantize::normalize_tensor_type_entry; +use crate::splits::SplitWindow; +use crate::types::{QuantType, TensorType}; + +const KV_QUANTIZE_IMATRIX_FILE: &str = "quantize.imatrix.file"; +const KV_QUANTIZE_IMATRIX_DATASET: &str = "quantize.imatrix.dataset"; +const KV_QUANTIZE_IMATRIX_N_ENTRIES: &str = "quantize.imatrix.entries_count"; +const KV_QUANTIZE_IMATRIX_N_CHUNKS: &str = "quantize.imatrix.chunks_count"; + +pub(crate) fn build_native_quantize_command( + args: &QuantRunnerArgs, + manifest: &Manifest, + staged_first_shard: &Path, + output_prefix: &Path, + window: SplitWindow, +) -> Result> { + ensure_native_quantize_supported(args)?; + ensure_full_split_window(manifest, window)?; + let quant = manifest + .quant + .as_deref() + .context("quantize manifest is missing quant type")?; + let mut command = vec![format!("{}-quantize", args.backend.as_str())]; + for library in &args.native_runtime_libraries { + command.push("--native-runtime-library".to_string()); + command.push(library.display().to_string()); + } + if args.allow_requantize { + command.push("--allow-requantize".to_string()); + } + if args.pure { + command.push("--pure".to_string()); + } + if args.dry_run { + command.push("--dry-run".to_string()); + } + if args.leave_output_tensor { + command.push("--leave-output-tensor".to_string()); + } + if let Some(imatrix) = args.imatrix.as_deref() { + command.push("--imatrix".to_string()); + command.push(imatrix.display().to_string()); + } + for include in &args.include_weights { + command.push("--include-weights".to_string()); + command.push(include.clone()); + } + for exclude in &args.exclude_weights { + command.push("--exclude-weights".to_string()); + command.push(exclude.clone()); + } + push_optional( + &mut command, + "--output-tensor-type", + &args.output_tensor_type, + ); + push_optional( + &mut command, + "--token-embedding-type", + &args.token_embedding_type, + ); + for entry in &args.tensor_type { + let normalized_entry = normalize_tensor_type_entry(entry)?; + command.push("--tensor-type".to_string()); + command.push(normalized_entry); + } + if let Some(tensor_type_file) = manifest.tensor_type_file.as_deref() { + command.push("--tensor-type-file".to_string()); + command.push(tensor_type_file.display().to_string()); + } + if let Some(prune_layers) = args.prune_layers.as_deref() { + command.push("--prune-layers".to_string()); + command.push(prune_layers.to_string()); + } + for override_kv in &args.override_kv { + command.push("--override-kv".to_string()); + command.push(override_kv.clone()); + } + command.extend([ + "--keep-split".to_string(), + staged_first_shard.display().to_string(), + llama_split_output_prefix(output_prefix) + .display() + .to_string(), + quant.to_string(), + ]); + if let Some(nthreads) = args.nthreads { + command.push(nthreads.to_string()); + } + Ok(command) +} + +pub(crate) fn run_native_quantize( + args: &QuantRunnerArgs, + manifest: &Manifest, + staged_first_shard: &Path, + output_prefix: &Path, + window: SplitWindow, +) -> Result { + ensure_native_quantize_supported(args)?; + ensure_full_split_window(manifest, window)?; + let native_inputs = NativeQuantizeInputs::build(args, manifest)?; + let mut params = unsafe { llama_quant_ffi::llama_model_quantize_default_params() }; + let quant = manifest + .quant + .as_deref() + .context("quantize manifest is missing quant type")? + .parse::() + .map_err(anyhow::Error::msg)?; + params.nthread = args.nthreads.map_or(0, |value| value as i32); + params.ftype = quant.as_llama_file_type(); + params.allow_requantize = args.allow_requantize; + params.quantize_output_tensor = !args.leave_output_tensor; + params.only_copy = quant == QuantType::Copy; + params.pure = args.pure; + params.keep_split = true; + params.dry_run = args.dry_run; + params.output_tensor_type = + optional_ggml_type(args.output_tensor_type.as_deref(), "--output-tensor-type")? + .unwrap_or(params.output_tensor_type); + params.token_embedding_type = optional_ggml_type( + args.token_embedding_type.as_deref(), + "--token-embedding-type", + )? + .unwrap_or(params.token_embedding_type); + params.tt_overrides = native_inputs.tensor_overrides_ptr(); + params.prune_layers = native_inputs.prune_layers_ptr(); + params.kv_overrides = native_inputs.kv_overrides_ptr(); + params.imatrix = native_inputs.imatrix_ptr(); + + let input = path_to_cstring(staged_first_shard)?; + let native_output_prefix = llama_split_output_prefix(output_prefix); + let output = path_to_cstring(&native_output_prefix)?; + let code = + unsafe { llama_quant_ffi::llama_model_quantize(input.as_ptr(), output.as_ptr(), ¶ms) }; + Ok(BackendRunStatus::from_code(code as i32)) +} + +fn ensure_native_quantize_supported(args: &QuantRunnerArgs) -> Result<()> { + load_llama_quant_runtime(&args.native_runtime_libraries)?; + ensure!( + llama_quant_ffi::native_runtime_loaded(), + "llama quant runtime is not linked; pass --native-runtime-library or build the standalone static target" + ); + ensure!( + args.include_weights.is_empty() || args.exclude_weights.is_empty(), + "--include-weights and --exclude-weights cannot be used together" + ); + ensure!( + args.max_memory.is_none(), + "--max-memory is not supported by the native llama quantize backend now that mesh-llm no longer patches llama-quantize memory chunking" + ); + Ok(()) +} + +fn ensure_full_split_window(manifest: &Manifest, window: SplitWindow) -> Result<()> { + ensure!( + window.first_split == 1 && window.last_split == manifest.expected_splits, + "native llama quantize backend no longer supports partial split windows after removing mesh-llm's patched llama-quantize split-window support; requested {}..{} of {}", + window.first_split, + window.last_split, + manifest.expected_splits + ); + Ok(()) +} + +fn load_llama_quant_runtime(libraries: &[PathBuf]) -> Result<()> { + if libraries.is_empty() || llama_quant_ffi::native_runtime_loaded() { + return Ok(()); + } + match unsafe { llama_quant_ffi::load_native_runtime_libraries(libraries) } { + Ok(()) | Err(llama_quant_ffi::NativeRuntimeLoadError::AlreadyLoaded) => Ok(()), + Err(error) => Err(anyhow!("load native llama quant runtime: {error}")), + } +} + +struct NativeQuantizeInputs { + _tensor_patterns: Vec, + tensor_overrides: Vec, + prune_layers: Vec, + kv_overrides: Vec, + imatrix: Option, +} + +impl NativeQuantizeInputs { + fn build(args: &QuantRunnerArgs, manifest: &Manifest) -> Result { + let (_tensor_patterns, tensor_overrides) = tensor_overrides(args, manifest)?; + let imatrix = args + .imatrix + .as_deref() + .map(|path| NativeImatrix::load(path, &args.include_weights, &args.exclude_weights)) + .transpose()?; + Ok(Self { + _tensor_patterns, + tensor_overrides, + prune_layers: prune_layers(args.prune_layers.as_deref())?, + kv_overrides: kv_overrides(&args.override_kv, imatrix.as_ref())?, + imatrix, + }) + } + + fn tensor_overrides_ptr(&self) -> *const llama_quant_ffi::LlamaModelTensorOverride { + if self.tensor_overrides.is_empty() { + std::ptr::null() + } else { + self.tensor_overrides.as_ptr() + } + } + + fn prune_layers_ptr(&self) -> *const i32 { + if self.prune_layers.is_empty() { + std::ptr::null() + } else { + self.prune_layers.as_ptr() + } + } + + fn kv_overrides_ptr(&self) -> *const llama_quant_ffi::LlamaModelKvOverride { + if self.kv_overrides.is_empty() { + std::ptr::null() + } else { + self.kv_overrides.as_ptr() + } + } + + fn imatrix_ptr(&self) -> *const llama_quant_ffi::LlamaModelImatrixData { + self.imatrix + .as_ref() + .map_or(std::ptr::null(), NativeImatrix::as_ptr) + } +} + +fn tensor_overrides( + args: &QuantRunnerArgs, + manifest: &Manifest, +) -> Result<(Vec, Vec)> { + let mut entries = args.tensor_type.clone(); + if let Some(path) = manifest.tensor_type_file.as_deref() { + let text = fs::read_to_string(path) + .with_context(|| format!("read tensor type file {}", path.display()))?; + entries.extend(text.split_whitespace().map(ToString::to_string)); + } + if entries.is_empty() { + return Ok((Vec::new(), Vec::new())); + } + + let mut patterns = Vec::with_capacity(entries.len()); + let mut overrides = Vec::with_capacity(entries.len() + 1); + for entry in entries { + let (raw_pattern, raw_type) = entry + .split_once('=') + .ok_or_else(|| anyhow!("malformed tensor type entry {entry:?}"))?; + ensure!( + !raw_pattern.is_empty(), + "tensor type entry has empty tensor name" + ); + let tensor_type = TensorType::parse(raw_type) + .with_context(|| format!("unsupported raw ggml tensor type {raw_type:?}"))? + .as_ggml_type() + .with_context(|| { + format!("tensor type override requires raw ggml_type, got {raw_type:?}") + })?; + patterns.push(CString::new(raw_pattern.to_ascii_lowercase())?); + overrides.push(llama_quant_ffi::LlamaModelTensorOverride { + pattern: patterns.last().expect("just pushed").as_ptr(), + tensor_type, + }); + } + overrides.push(llama_quant_ffi::LlamaModelTensorOverride { + pattern: std::ptr::null(), + tensor_type: llama_quant_ffi::GgmlType::Count, + }); + Ok((patterns, overrides)) +} + +fn prune_layers(raw: Option<&str>) -> Result> { + let Some(raw) = raw else { + return Ok(Vec::new()); + }; + let mut layers = raw + .split(',') + .map(|value| { + let layer = value + .parse::() + .with_context(|| format!("invalid layer id {value:?}"))?; + ensure!(layer >= 0, "invalid negative layer id {layer}"); + Ok(layer) + }) + .collect::>>()?; + layers.sort_unstable(); + layers.dedup(); + layers.push(-1); + Ok(layers) +} + +fn kv_overrides( + raw_overrides: &[String], + imatrix: Option<&NativeImatrix>, +) -> Result> { + if raw_overrides.is_empty() && imatrix.is_none() { + return Ok(Vec::new()); + } + let imatrix_extra = if imatrix.is_some() { 4 } else { 0 }; + let mut overrides = Vec::with_capacity(raw_overrides.len() + imatrix_extra + 1); + for raw in raw_overrides { + overrides.push(parse_kv_override(raw)?); + } + if let Some(imatrix) = imatrix { + overrides.push(string_kv_override( + KV_QUANTIZE_IMATRIX_FILE, + &imatrix.source_path().display().to_string(), + )?); + if let Some(dataset) = imatrix.dataset() { + overrides.push(string_kv_override(KV_QUANTIZE_IMATRIX_DATASET, dataset)?); + } + overrides.push(int_kv_override( + KV_QUANTIZE_IMATRIX_N_ENTRIES, + imatrix.entry_count() as i64, + )?); + if imatrix.chunk_count() > 0 { + overrides.push(int_kv_override( + KV_QUANTIZE_IMATRIX_N_CHUNKS, + imatrix.chunk_count() as i64, + )?); + } + } + overrides.push(llama_quant_ffi::LlamaModelKvOverride { + tag: llama_quant_ffi::LlamaModelKvOverrideType::Int, + key: [0; 128], + value: llama_quant_ffi::LlamaModelKvOverrideValue { val_i64: 0 }, + }); + Ok(overrides) +} + +fn int_kv_override(key: &str, value: i64) -> Result { + Ok(llama_quant_ffi::LlamaModelKvOverride { + tag: llama_quant_ffi::LlamaModelKvOverrideType::Int, + key: fixed_c_char_array(key, "KV override key")?, + value: llama_quant_ffi::LlamaModelKvOverrideValue { val_i64: value }, + }) +} + +fn string_kv_override(key: &str, value: &str) -> Result { + Ok(llama_quant_ffi::LlamaModelKvOverride { + tag: llama_quant_ffi::LlamaModelKvOverrideType::Str, + key: fixed_c_char_array(key, "KV override key")?, + value: llama_quant_ffi::LlamaModelKvOverrideValue { + val_str: fixed_c_char_array(value, "KV override string value")?, + }, + }) +} + +fn parse_kv_override(raw: &str) -> Result { + let (key, value) = raw + .split_once('=') + .ok_or_else(|| anyhow!("malformed KV override {raw:?}"))?; + ensure!(!key.is_empty(), "KV override has empty key"); + let key = fixed_c_char_array(key, "KV override key")?; + if let Some(rest) = value.strip_prefix("int:") { + return Ok(llama_quant_ffi::LlamaModelKvOverride { + tag: llama_quant_ffi::LlamaModelKvOverrideType::Int, + key, + value: llama_quant_ffi::LlamaModelKvOverrideValue { + val_i64: rest.parse::()?, + }, + }); + } + if let Some(rest) = value.strip_prefix("float:") { + return Ok(llama_quant_ffi::LlamaModelKvOverride { + tag: llama_quant_ffi::LlamaModelKvOverrideType::Float, + key, + value: llama_quant_ffi::LlamaModelKvOverrideValue { + val_f64: rest.parse::()?, + }, + }); + } + if let Some(rest) = value.strip_prefix("bool:") { + let val_bool = match rest { + "true" => true, + "false" => false, + _ => return Err(anyhow!("invalid bool KV override value {rest:?}")), + }; + return Ok(llama_quant_ffi::LlamaModelKvOverride { + tag: llama_quant_ffi::LlamaModelKvOverrideType::Bool, + key, + value: llama_quant_ffi::LlamaModelKvOverrideValue { val_bool }, + }); + } + if let Some(rest) = value.strip_prefix("str:") { + return Ok(llama_quant_ffi::LlamaModelKvOverride { + tag: llama_quant_ffi::LlamaModelKvOverrideType::Str, + key, + value: llama_quant_ffi::LlamaModelKvOverrideValue { + val_str: fixed_c_char_array(rest, "KV override string value")?, + }, + }); + } + Err(anyhow!("invalid KV override type in {raw:?}")) +} + +fn fixed_c_char_array(raw: &str, label: &str) -> Result<[c_char; 128]> { + ensure!(raw.len() < 128, "{label} cannot exceed 127 bytes"); + let cstring = CString::new(raw).with_context(|| format!("{label} contains NUL byte"))?; + let mut out = [0 as c_char; 128]; + for (dst, src) in out.iter_mut().zip(cstring.as_bytes_with_nul()) { + *dst = *src as c_char; + } + Ok(out) +} + +fn optional_ggml_type(raw: Option<&str>, flag: &str) -> Result> { + raw.map(|value| { + let tensor_type = TensorType::parse(value) + .with_context(|| format!("{flag} has unsupported ggml type {value:?}"))?; + tensor_type + .as_ggml_type() + .with_context(|| format!("{flag} requires a raw ggml_type, got {value:?}")) + }) + .transpose() +} + +fn path_to_cstring(path: &Path) -> Result { + let text = path + .to_str() + .with_context(|| format!("path is not valid UTF-8: {}", path.display()))?; + CString::new(text).with_context(|| format!("path contains NUL byte: {}", path.display())) +} + +fn llama_split_output_prefix(path: &Path) -> PathBuf { + if path.extension().and_then(|value| value.to_str()) != Some("gguf") { + return path.to_path_buf(); + } + let Some(stem) = path.file_stem() else { + return path.to_path_buf(); + }; + path.with_file_name(stem) +} + +fn push_optional(command: &mut Vec, flag: &str, value: &Option) { + if let Some(value) = value { + command.push(flag.to_string()); + command.push(value.clone()); + } +} + +#[cfg(test)] +mod tests { + use crate::MANIFEST_VERSION; + use crate::backend::BackendKind; + use crate::types::JobKind; + + use super::*; + + #[test] + fn builds_native_quantize_command() { + let mut args = native_args(); + args.tensor_type = vec!["mtp_head.weight=NVFP4".to_string()]; + args.prune_layers = Some("2,1,2".to_string()); + args.override_kv = vec!["general.name=str:test".to_string()]; + let manifest = manifest(None); + let command = build_native_quantize_command( + &args, + &manifest, + Path::new("/tmp/in/model-00001-of-00002.gguf"), + Path::new("/tmp/out/model-q2"), + SplitWindow { + first_split: 1, + last_split: 2, + }, + ) + .unwrap(); + + assert_eq!(command[0], "llama-api-quantize"); + assert!(!command.contains(&"--native-runtime-library".to_string())); + assert!(command.contains(&"--keep-split".to_string())); + assert!(command.contains(&"--tensor-type".to_string())); + assert!(command.contains(&"--prune-layers".to_string())); + assert!(command.contains(&"--override-kv".to_string())); + assert!(!command.contains(&"--max-memory".to_string())); + assert!(command.contains(&"Q2_K".to_string())); + } + + #[test] + fn passes_explicit_tensor_type_file_to_command() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + let recipe = root.join("glm-5.2-q2-k-mtp-q8.tensor-types.txt"); + fs::write(&recipe, "(^|\\.)nextn\\.=Q8_0").unwrap(); + let args = native_args(); + let manifest = manifest(Some(recipe.clone())); + + let command = build_native_quantize_command( + &args, + &manifest, + Path::new("/tmp/in/model-00001-of-00002.gguf"), + Path::new("/tmp/out/model-q2-mtp-q8"), + SplitWindow { + first_split: 1, + last_split: 2, + }, + ) + .unwrap(); + + assert!(command.contains(&"--tensor-type-file".to_string())); + assert!(command.contains(&recipe.display().to_string())); + assert!(!command.contains(&"(^|\\.)nextn\\.=Q8_0".to_string())); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn builds_skippy_abi_quantize_command_label() { + let mut args = native_args(); + args.backend = BackendKind::SkippyAbi; + let manifest = manifest(None); + + let command = build_native_quantize_command( + &args, + &manifest, + Path::new("/tmp/in/model-00001-of-00002.gguf"), + Path::new("/tmp/out/model-q2"), + SplitWindow { + first_split: 1, + last_split: 2, + }, + ) + .unwrap(); + + assert_eq!(command[0], "skippy-abi-quantize"); + assert!(!command.contains(&"--native-runtime-library".to_string())); + } + + #[test] + fn strips_gguf_extension_from_llama_split_output_prefix() { + let args = native_args(); + let manifest = manifest(None); + + let command = build_native_quantize_command( + &args, + &manifest, + Path::new("/tmp/in/model.gguf"), + Path::new("/tmp/out/model-q4.gguf"), + SplitWindow { + first_split: 1, + last_split: 2, + }, + ) + .unwrap(); + + assert!(command.contains(&"/tmp/out/model-q4".to_string())); + assert!(!command.contains(&"/tmp/out/model-q4.gguf".to_string())); + assert_eq!( + llama_split_output_prefix(Path::new("/tmp/out/model-q4.gguf")), + PathBuf::from("/tmp/out/model-q4") + ); + } + + #[test] + fn rejects_partial_split_window_after_dropping_patched_llama_quantize() { + let args = native_args(); + let manifest = manifest(None); + + let error = build_native_quantize_command( + &args, + &manifest, + Path::new("/tmp/in/model-00001-of-00002.gguf"), + Path::new("/tmp/out/model-q2"), + SplitWindow { + first_split: 1, + last_split: 1, + }, + ) + .unwrap_err(); + + assert!( + error + .to_string() + .contains("no longer supports partial split windows") + ); + } + + #[test] + fn builds_native_inputs_for_tensor_prune_and_kv_overrides() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + let recipe = root.join("tensors.txt"); + fs::write(&recipe, "blk.0.weight=Q8_0").unwrap(); + let mut args = native_args(); + args.tensor_type = vec!["mtp_head.weight=NVFP4".to_string()]; + args.prune_layers = Some("2,1,2".to_string()); + args.override_kv = vec![ + "general.name=str:test".to_string(), + "custom.count=int:7".to_string(), + "custom.enabled=bool:true".to_string(), + ]; + let manifest = manifest(Some(recipe)); + + let inputs = NativeQuantizeInputs::build(&args, &manifest).unwrap(); + + assert_eq!(inputs.tensor_overrides.len(), 3); + assert_eq!(inputs.prune_layers, vec![1, 2, -1]); + assert_eq!(inputs.kv_overrides.len(), 4); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn appends_tensor_type_file_entries_after_explicit_overrides() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + let recipe = root.join("glm-5.2-q2-k-mtp-q8.tensor-types.txt"); + fs::write( + &recipe, + "^token_embd\\.weight$=Q8_0\n(^|\\.)nextn\\.=Q8_0\n", + ) + .unwrap(); + let mut args = native_args(); + args.tensor_type = vec!["nextn\\.pre_projection\\.weight=F16".to_string()]; + let manifest = manifest(Some(recipe)); + + let inputs = NativeQuantizeInputs::build(&args, &manifest).unwrap(); + + assert_eq!(inputs.tensor_overrides.len(), 4); + assert_eq!( + inputs._tensor_patterns[0].to_str().unwrap(), + "nextn\\.pre_projection\\.weight" + ); + assert_eq!( + inputs._tensor_patterns[1].to_str().unwrap(), + "^token_embd\\.weight$" + ); + assert_eq!( + inputs._tensor_patterns[2].to_str().unwrap(), + "(^|\\.)nextn\\." + ); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn loads_legacy_imatrix_with_include_filter() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + let imatrix_path = root.join("imatrix.dat"); + write_legacy_imatrix( + &imatrix_path, + &[ + ("blk.0.attn_q.weight", 2, &[2.0, 4.0]), + ("blk.0.ffn_down.weight", 1, &[9.0, 12.0]), + ], + ); + let mut args = native_args(); + args.imatrix = Some(imatrix_path); + args.include_weights = vec!["attn_q".to_string()]; + let manifest = manifest(None); + + let inputs = NativeQuantizeInputs::build(&args, &manifest).unwrap(); + + assert_eq!(inputs.imatrix.as_ref().unwrap().entry_count(), 1); + assert!(inputs.kv_overrides.len() >= 3); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn loads_gguf_imatrix_for_native_inputs() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + let imatrix_path = root.join("imatrix.gguf"); + write_gguf_imatrix(&imatrix_path); + let mut args = native_args(); + args.imatrix = Some(imatrix_path); + args.include_weights = vec!["attn_q".to_string()]; + let manifest = manifest(None); + + let inputs = NativeQuantizeInputs::build(&args, &manifest).unwrap(); + let imatrix = inputs.imatrix.as_ref().unwrap(); + + assert_eq!(imatrix.entry_count(), 1); + assert_eq!(imatrix.dataset(), Some("calibration.txt")); + assert_eq!(imatrix.chunk_count(), 7); + assert!(inputs.kv_overrides.len() >= 3); + fs::remove_dir_all(root).unwrap(); + } + + fn native_args() -> QuantRunnerArgs { + QuantRunnerArgs { + backend: BackendKind::LlamaApi, + native_runtime_libraries: Vec::new(), + work_dir: PathBuf::from("/tmp/work"), + print_only: false, + dry_run: false, + allow_requantize: true, + pure: false, + imatrix: None, + include_weights: Vec::new(), + exclude_weights: Vec::new(), + output_tensor_type: None, + token_embedding_type: None, + tensor_type: Vec::new(), + prune_layers: None, + override_kv: Vec::new(), + nthreads: Some(8), + leave_output_tensor: true, + no_stage_source: false, + keep_staged_source: false, + spool_dir: None, + keep_spool: false, + watchdog_seconds: None, + max_memory: None, + memory_policy: crate::memory_budget::MemoryPolicy::Hard, + record_dir: None, + json_event_file: None, + json_event_interval_seconds: 120, + json_event_window: 8, + } + } + + fn manifest(tensor_type_file: Option) -> Manifest { + Manifest { + schema_version: MANIFEST_VERSION, + kind: JobKind::QuantizeGguf, + source: PathBuf::from("/tmp/source"), + source_prefix: Some("BF16".to_string()), + target: PathBuf::from("/tmp/target"), + target_prefix: "Q2_K".to_string(), + output_basename: "model-q2".to_string(), + expected_splits: 2, + window_size: 1, + quant: Some("Q2_K".to_string()), + output_type: None, + tensor_type_file, + tensor_type_recipe: None, + } + } + + fn unique_temp_dir() -> PathBuf { + static NEXT_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let id = NEXT_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + std::env::temp_dir().join(format!("skippy-native-quantize-{nanos}-{id}")) + } + + fn write_legacy_imatrix(path: &Path, entries: &[(&str, i32, &[f32])]) { + let mut bytes = Vec::new(); + bytes.extend_from_slice(&(entries.len() as i32).to_le_bytes()); + for (name, ncall, values) in entries { + bytes.extend_from_slice(&(name.len() as i32).to_le_bytes()); + bytes.extend_from_slice(name.as_bytes()); + bytes.extend_from_slice(&ncall.to_le_bytes()); + bytes.extend_from_slice(&(values.len() as i32).to_le_bytes()); + for value in *values { + bytes.extend_from_slice(&value.to_le_bytes()); + } + } + bytes.extend_from_slice(&3_i32.to_le_bytes()); + let dataset = "dataset.txt"; + bytes.extend_from_slice(&(dataset.len() as i32).to_le_bytes()); + bytes.extend_from_slice(dataset.as_bytes()); + fs::write(path, bytes).unwrap(); + } + + fn write_gguf_imatrix(path: &Path) { + const GGUF_MAGIC: &[u8; 4] = b"GGUF"; + const GGML_TYPE_F32: u32 = 0; + const GGUF_TYPE_UINT32: u32 = 4; + const GGUF_TYPE_STRING: u32 = 8; + const GGUF_TYPE_ARRAY: u32 = 9; + const KV_GENERAL_ALIGNMENT: &str = "general.alignment"; + const KV_IMATRIX_DATASETS: &str = "imatrix.datasets"; + const KV_IMATRIX_CHUNK_COUNT: &str = "imatrix.chunk_count"; + + let sums_name = "blk.0.attn_q.weight.in_sum2"; + let counts_name = "blk.0.attn_q.weight.counts"; + let mut bytes = Vec::new(); + bytes.extend_from_slice(GGUF_MAGIC); + bytes.extend_from_slice(&3_u32.to_le_bytes()); + bytes.extend_from_slice(&2_u64.to_le_bytes()); + bytes.extend_from_slice(&4_u64.to_le_bytes()); + write_gguf_kv_string(&mut bytes, "general.type", "imatrix", GGUF_TYPE_STRING); + write_gguf_kv_u32(&mut bytes, KV_GENERAL_ALIGNMENT, 32, GGUF_TYPE_UINT32); + write_gguf_kv_u32(&mut bytes, KV_IMATRIX_CHUNK_COUNT, 7, GGUF_TYPE_UINT32); + write_gguf_kv_string_array( + &mut bytes, + KV_IMATRIX_DATASETS, + &["calibration.txt"], + GGUF_TYPE_ARRAY, + GGUF_TYPE_STRING, + ); + write_gguf_tensor_info(&mut bytes, sums_name, &[2, 2], GGML_TYPE_F32, 0); + write_gguf_tensor_info(&mut bytes, counts_name, &[1, 2], GGML_TYPE_F32, 16); + while bytes.len() % 32 != 0 { + bytes.push(0); + } + for value in [2.0_f32, 4.0, 9.0, 11.0] { + bytes.extend_from_slice(&value.to_le_bytes()); + } + for value in [2.0_f32, 0.0] { + bytes.extend_from_slice(&value.to_le_bytes()); + } + fs::write(path, bytes).unwrap(); + } + + fn write_gguf_kv_string(bytes: &mut Vec, key: &str, value: &str, string_type: u32) { + write_gguf_string(bytes, key); + bytes.extend_from_slice(&string_type.to_le_bytes()); + write_gguf_string(bytes, value); + } + + fn write_gguf_kv_u32(bytes: &mut Vec, key: &str, value: u32, uint32_type: u32) { + write_gguf_string(bytes, key); + bytes.extend_from_slice(&uint32_type.to_le_bytes()); + bytes.extend_from_slice(&value.to_le_bytes()); + } + + fn write_gguf_kv_string_array( + bytes: &mut Vec, + key: &str, + values: &[&str], + array_type: u32, + string_type: u32, + ) { + write_gguf_string(bytes, key); + bytes.extend_from_slice(&array_type.to_le_bytes()); + bytes.extend_from_slice(&string_type.to_le_bytes()); + bytes.extend_from_slice(&(values.len() as u64).to_le_bytes()); + for value in values { + write_gguf_string(bytes, value); + } + } + + fn write_gguf_tensor_info( + bytes: &mut Vec, + name: &str, + dims: &[u64], + tensor_type: u32, + offset: u64, + ) { + write_gguf_string(bytes, name); + bytes.extend_from_slice(&(dims.len() as u32).to_le_bytes()); + for dim in dims { + bytes.extend_from_slice(&dim.to_le_bytes()); + } + bytes.extend_from_slice(&tensor_type.to_le_bytes()); + bytes.extend_from_slice(&offset.to_le_bytes()); + } + + fn write_gguf_string(bytes: &mut Vec, value: &str) { + bytes.extend_from_slice(&(value.len() as u64).to_le_bytes()); + bytes.extend_from_slice(value.as_bytes()); + } +} diff --git a/crates/skippy-quantize/src/output.rs b/crates/skippy-quantize/src/output.rs new file mode 100644 index 0000000000..eadc013717 --- /dev/null +++ b/crates/skippy-quantize/src/output.rs @@ -0,0 +1,395 @@ +use std::collections::VecDeque; +use std::fs; +use std::path::Path; +use std::path::PathBuf; +use std::sync::{ + Arc, Mutex, + mpsc::{self, Sender}, +}; +use std::thread; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use anyhow::{Context, Result}; +use serde::{Deserialize, Serialize}; + +use crate::splits::{ShardRange, SplitWindow}; + +const BAR_WIDTH: usize = 24; +const JSON_EVENT_SCHEMA_VERSION: u32 = 1; + +pub(crate) fn print_json_pretty(value: &impl Serialize) -> Result<()> { + println!("{}", serde_json::to_string_pretty(value)?); + Ok(()) +} + +pub(crate) fn progress_bar(completed: usize, expected: u32) -> String { + let expected = expected as usize; + let filled = completed + .saturating_mul(BAR_WIDTH) + .checked_div(expected) + .unwrap_or_default() + .min(BAR_WIDTH); + format!( + "[{}{}]", + "█".repeat(filled), + "░".repeat(BAR_WIDTH.saturating_sub(filled)) + ) +} + +pub(crate) fn percent(completed: usize, expected: u32) -> f64 { + if expected == 0 { + 0.0 + } else { + (completed as f64 / f64::from(expected)) * 100.0 + } +} + +pub(crate) fn print_progress_line(label: &str, completed: usize, expected: u32) { + println!( + "📊 {label}: {} {completed}/{expected} shards ({:.2}%)", + progress_bar(completed, expected), + percent(completed, expected) + ); +} + +pub(crate) fn print_success(message: impl AsRef) { + println!("✅ {}", message.as_ref()); +} + +pub(crate) fn print_info(message: impl AsRef) { + println!("ℹ️ {}", message.as_ref()); +} + +pub(crate) fn print_warn(message: impl AsRef) { + println!("⚠️ {}", message.as_ref()); +} + +pub(crate) fn print_copy(source: &Path, target: &Path, size_bytes: Option) { + match size_bytes { + Some(size_bytes) => println!( + "📤 Copying {} -> {} ({})", + source.display(), + target.display(), + format_bytes(size_bytes) + ), + None => println!("📤 Copying {} -> {}", source.display(), target.display()), + } +} + +pub(crate) fn print_path_event(emoji: &str, label: &str, path: &Path) { + println!("{emoji} {label}: {}", path.display()); +} + +pub(crate) fn print_window(label: &str, window: SplitWindow) { + println!("🪟 {label}: {}", format_window(window)); +} + +pub(crate) fn format_window(window: SplitWindow) -> String { + if window.first_split == window.last_split { + window.first_split.to_string() + } else { + format!("{}..{}", window.first_split, window.last_split) + } +} + +pub(crate) fn format_shard_ranges(ranges: &[ShardRange]) -> String { + ranges + .iter() + .map(|range| { + if range.first_split == range.last_split { + range.first_split.to_string() + } else { + format!("{}..{}", range.first_split, range.last_split) + } + }) + .collect::>() + .join(", ") +} + +pub(crate) fn format_bytes(bytes: u64) -> String { + const KIB: f64 = 1024.0; + const MIB: f64 = KIB * 1024.0; + const GIB: f64 = MIB * 1024.0; + let bytes_f = bytes as f64; + if bytes_f >= GIB { + format!("{:.2} GiB", bytes_f / GIB) + } else if bytes_f >= MIB { + format!("{:.2} MiB", bytes_f / MIB) + } else if bytes_f >= KIB { + format!("{:.2} KiB", bytes_f / KIB) + } else { + format!("{bytes} B") + } +} + +#[derive(Debug, Clone)] +pub(crate) struct JsonEventConfig { + pub(crate) file: Option, + pub(crate) interval_seconds: u64, + pub(crate) window_size: usize, +} + +#[derive(Debug, Serialize, Deserialize)] +struct JsonEventSnapshot { + schema_version: u32, + event: String, + kind: String, + phase: String, + started_unix_ms: u128, + updated_unix_ms: u128, + interval_seconds: u64, + window_size: usize, + current_window: Option, + recent_events: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct JsonRecentEvent { + unix_ms: u128, + message: String, +} + +#[derive(Debug)] +struct JsonEventState { + kind: String, + phase: String, + started_unix_ms: u128, + updated_unix_ms: u128, + interval_seconds: u64, + window_size: usize, + current_window: Option, + recent_events: VecDeque, +} + +impl JsonEventState { + fn snapshot(&self) -> JsonEventSnapshot { + JsonEventSnapshot { + schema_version: JSON_EVENT_SCHEMA_VERSION, + event: "skippy_quantize_periodic_status".to_string(), + kind: self.kind.clone(), + phase: self.phase.clone(), + started_unix_ms: self.started_unix_ms, + updated_unix_ms: self.updated_unix_ms, + interval_seconds: self.interval_seconds, + window_size: self.window_size, + current_window: self.current_window, + recent_events: self.recent_events.iter().cloned().collect(), + } + } +} + +pub(crate) struct JsonEventReporter { + path: Option, + state: Option>>, + stop: Option>, + thread: Option>, +} + +impl JsonEventReporter { + pub(crate) fn start( + config: JsonEventConfig, + kind: impl Into, + window: Option, + ) -> Result { + let Some(path) = config.file else { + return Ok(Self::disabled()); + }; + let interval_seconds = config.interval_seconds.max(1); + let window_size = config.window_size.max(1); + let now = unix_timestamp_ms(); + let state = Arc::new(Mutex::new(JsonEventState { + kind: kind.into(), + phase: "starting".to_string(), + started_unix_ms: now, + updated_unix_ms: now, + interval_seconds, + window_size, + current_window: window, + recent_events: VecDeque::new(), + })); + write_snapshot(&path, &state)?; + let (stop, thread) = spawn_periodic_writer(path.clone(), Arc::clone(&state)); + Ok(Self { + path: Some(path), + state: Some(state), + stop: Some(stop), + thread: Some(thread), + }) + } + + fn disabled() -> Self { + Self { + path: None, + state: None, + stop: None, + thread: None, + } + } + + pub(crate) fn record(&self, message: impl Into) -> Result<()> { + self.with_state(|state| { + state.updated_unix_ms = unix_timestamp_ms(); + state.recent_events.push_back(JsonRecentEvent { + unix_ms: state.updated_unix_ms, + message: message.into(), + }); + while state.recent_events.len() > state.window_size { + state.recent_events.pop_front(); + } + }) + } + + pub(crate) fn set_phase(&self, phase: impl Into) -> Result<()> { + self.with_state(|state| { + state.updated_unix_ms = unix_timestamp_ms(); + state.phase = phase.into(); + }) + } + + pub(crate) fn finish(mut self, phase: impl Into) -> Result<()> { + self.set_phase(phase)?; + self.write_now()?; + self.stop_thread(); + Ok(()) + } + + pub(crate) fn write_now(&self) -> Result<()> { + if let (Some(path), Some(state)) = (&self.path, &self.state) { + write_snapshot(path, state)?; + } + Ok(()) + } + + fn with_state(&self, update: impl FnOnce(&mut JsonEventState)) -> Result<()> { + if let Some(state) = &self.state { + let mut state = state + .lock() + .map_err(|_| anyhow::anyhow!("json event state lock poisoned"))?; + update(&mut state); + } + Ok(()) + } + + fn stop_thread(&mut self) { + if let Some(stop) = &self.stop { + let _ = stop.send(()); + } + if let Some(thread) = self.thread.take() { + let _ = thread.join(); + } + } +} + +impl Drop for JsonEventReporter { + fn drop(&mut self) { + self.stop_thread(); + } +} + +fn spawn_periodic_writer( + path: PathBuf, + state: Arc>, +) -> (Sender<()>, thread::JoinHandle<()>) { + let (stop_tx, stop_rx) = mpsc::channel(); + let interval = state + .lock() + .map(|state| Duration::from_secs(state.interval_seconds)) + .unwrap_or_else(|_| Duration::from_secs(120)); + let thread = thread::spawn(move || { + loop { + match stop_rx.recv_timeout(interval) { + Ok(()) | Err(mpsc::RecvTimeoutError::Disconnected) => break, + Err(mpsc::RecvTimeoutError::Timeout) => {} + } + if let Ok(mut state) = state.lock() { + state.updated_unix_ms = unix_timestamp_ms(); + } + let _ = write_snapshot(&path, &state); + } + }); + (stop_tx, thread) +} + +fn write_snapshot(path: &Path, state: &Arc>) -> Result<()> { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?; + } + let snapshot = { + let state = state + .lock() + .map_err(|_| anyhow::anyhow!("json event state lock poisoned"))?; + state.snapshot() + }; + let temp = path.with_extension(format!( + "{}.tmp", + path.extension() + .and_then(|extension| extension.to_str()) + .unwrap_or("json") + )); + fs::write(&temp, serde_json::to_vec_pretty(&snapshot)?) + .with_context(|| format!("write {}", temp.display()))?; + fs::rename(&temp, path).with_context(|| { + format!( + "replace json event snapshot {} with {}", + path.display(), + temp.display() + ) + })?; + Ok(()) +} + +fn unix_timestamp_ms() -> u128 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |duration| duration.as_millis()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn json_event_snapshot_keeps_bounded_recent_window() { + let root = unique_temp_dir(); + let path = root.join("events").join("status.json"); + let reporter = JsonEventReporter::start( + JsonEventConfig { + file: Some(path.clone()), + interval_seconds: 60, + window_size: 2, + }, + "quant", + Some(SplitWindow { + first_split: 3, + last_split: 4, + }), + ) + .unwrap(); + + reporter.set_phase("running").unwrap(); + reporter.record("first").unwrap(); + reporter.record("second").unwrap(); + reporter.record("third").unwrap(); + reporter.write_now().unwrap(); + reporter.finish("complete").unwrap(); + + let snapshot: JsonEventSnapshot = + serde_json::from_slice(&fs::read(&path).unwrap()).unwrap(); + assert_eq!(snapshot.phase, "complete"); + assert_eq!(snapshot.recent_events.len(), 2); + assert_eq!(snapshot.recent_events[0].message, "second"); + assert_eq!(snapshot.recent_events[1].message, "third"); + assert_eq!(snapshot.current_window.unwrap().first_split, 3); + fs::remove_dir_all(root).ok(); + } + + fn unique_temp_dir() -> PathBuf { + static NEXT_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos(); + let id = NEXT_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + std::env::temp_dir().join(format!("skippy-quantize-json-events-{nanos}-{id}")) + } +} diff --git a/crates/skippy-quantize/src/plan_convert.rs b/crates/skippy-quantize/src/plan_convert.rs new file mode 100644 index 0000000000..09551d6219 --- /dev/null +++ b/crates/skippy-quantize/src/plan_convert.rs @@ -0,0 +1,151 @@ +use std::path::PathBuf; + +use anyhow::Result; +use clap::Parser; + +use crate::gguf_template::metadata_from_hf_config; +use crate::gguf_writer::{ + RawGgufValidation, RawGgufWriteOptions, TensorSelection, validate_raw_safetensors_gguf, + write_raw_safetensors_gguf, +}; +use crate::hf_checkpoint::{ + HfCheckpointPlan, inspect_hf_checkpoint, verify_hf_checkpoint_tensor_streams, +}; +use crate::memory_budget::MemorySize; +use crate::output::{format_bytes, print_info, print_json_pretty, print_success}; +use crate::tensor_map::TensorNameMap; +use crate::tokenizer_metadata::ensure_native_tokenizer_metadata_supported; +use crate::types::ConvertOutputType; + +#[derive(Debug, Parser)] +pub(crate) struct PlanConvertArgs { + source: PathBuf, + #[arg(long)] + max_memory: Option, + #[arg(long, default_value_t = 0.60)] + staging_fraction: f64, + #[arg(long)] + verify_streaming: bool, + #[arg(long, default_value_t = 8 * 1024 * 1024)] + stream_buffer_bytes: usize, + #[arg(long)] + write_raw_gguf: Option, + #[arg(long)] + hf_config_metadata: bool, + #[arg(long)] + validate_native: bool, + #[arg(long, value_enum, default_value_t = ConvertOutputType::Bf16)] + output_type: ConvertOutputType, + #[arg(long)] + json: bool, +} + +pub(crate) fn run_plan_convert(args: PlanConvertArgs) -> Result<()> { + let mut plan = inspect_hf_checkpoint(&args.source, args.max_memory, args.staging_fraction)?; + let mut native_validation = None; + if args.verify_streaming { + plan.stream_verification = Some(verify_hf_checkpoint_tensor_streams( + &args.source, + args.stream_buffer_bytes, + )?); + } + if args.validate_native { + ensure_native_tokenizer_metadata_supported(&args.source)?; + let metadata = metadata_from_hf_config(&args.source, plan.tensor_count)?; + native_validation = Some(validate_raw_safetensors_gguf( + &args.source, + RawGgufWriteOptions { + buffer_size: args.stream_buffer_bytes, + metadata: Some(metadata), + tensor_name_map: TensorNameMap::HfToGguf, + split: None, + output_type: Some(args.output_type), + tensor_selection: TensorSelection::All, + }, + )?); + } + if let Some(output) = args.write_raw_gguf.as_deref() { + let metadata = if args.hf_config_metadata { + ensure_native_tokenizer_metadata_supported(&args.source)?; + Some(metadata_from_hf_config(&args.source, plan.tensor_count)?) + } else { + None + }; + write_raw_safetensors_gguf( + &args.source, + output, + RawGgufWriteOptions { + buffer_size: args.stream_buffer_bytes, + metadata, + tensor_name_map: if args.hf_config_metadata { + TensorNameMap::HfToGguf + } else { + TensorNameMap::Raw + }, + split: None, + output_type: args.hf_config_metadata.then_some(args.output_type), + tensor_selection: TensorSelection::All, + }, + )?; + } + if args.json { + if let Some(native_validation) = native_validation { + print_json_pretty(&serde_json::json!({ + "plan": plan, + "native_validation": native_validation, + }))?; + } else { + print_json_pretty(&plan)?; + } + } else { + print_plan(&plan); + if let Some(native_validation) = native_validation { + print_native_validation(&native_validation); + } + } + Ok(()) +} + +fn print_plan(plan: &HfCheckpointPlan) { + print_success(format!("Checkpoint: {}", plan.source.display())); + print_info(format!("SafeTensors files: {}", plan.safetensor_count)); + print_info(format!("Tensors: {}", plan.tensor_count)); + print_info(format!( + "Total tensor bytes: {}", + format_bytes(plan.total_tensor_bytes) + )); + print_info(format!( + "Largest tensor: {}", + format_bytes(plan.largest_tensor_bytes) + )); + print_info(format!("Source windows: {}", plan.source_windows.len())); + for window in &plan.source_windows { + print_info(format!( + "Window {}: {} file(s), {}", + window.index, + window.files.len(), + format_bytes(window.total_tensor_bytes) + )); + } + if let Some(verification) = &plan.stream_verification { + print_success(format!( + "Stream verified: {} tensors, {}, buffer {}", + verification.tensor_count, + format_bytes(verification.streamed_bytes), + format_bytes(verification.buffer_size as u64) + )); + } +} + +fn print_native_validation(validation: &RawGgufValidation) { + print_success("Native writer validation passed"); + print_info(format!( + "Selected tensors: {} ({})", + validation.selected_tensor_count, + format_bytes(validation.selected_tensor_bytes) + )); + print_info(format!("Metadata entries: {}", validation.metadata_count)); + if let Some(output_type) = validation.output_type.as_deref() { + print_info(format!("Output type: {output_type}")); + } +} diff --git a/crates/skippy-quantize/src/preflight.rs b/crates/skippy-quantize/src/preflight.rs new file mode 100644 index 0000000000..6576e5919a --- /dev/null +++ b/crates/skippy-quantize/src/preflight.rs @@ -0,0 +1,325 @@ +use std::path::{Path, PathBuf}; + +use anyhow::{Result, ensure}; +use serde::Serialize; + +use crate::backend::{BackendKind, ensure_convert_backend, ensure_quant_backend}; +use crate::manifest::{Manifest, manifest_progress, read_manifest}; +use crate::output::{ + print_info, print_json_pretty, print_progress_line, print_success, print_warn, +}; +use crate::splits::{Progress, SplitWindow, next_missing_window_in_range, split_status}; +use crate::types::JobKind; + +#[derive(Debug, Serialize)] +struct JobPreflight { + kind: JobKind, + manifest_path: PathBuf, + manifest_exists: bool, + manifest_matches: bool, + source_complete: Option, + source_shards: Option, + expected_source_shards: Option, + source_missing_ranges: Option>, + target_shards: usize, + expected_target_shards: u32, + first_missing_target: Option, + target_missing_ranges: Vec, + next_window: Option, + requested_window: Option, + next_requested_window: Option, + backend_kind: String, + backend_path: Option, + backend_ready: bool, + backend_error: Option, +} + +#[derive(Debug, Serialize)] +struct ProgressWindow { + first_split: u32, + last_split: u32, +} + +pub fn run_job_preflight( + manifest_path: &Path, + manifest: &Manifest, + source_split: Option<(&Path, &str)>, + requested_window: Option, + backend_kind: BackendKind, + backend_path: Option<&Path>, + json: bool, +) -> Result<()> { + ensure_backend_supported(manifest.kind, backend_kind)?; + let manifest_exists = manifest_path.exists(); + let manifest_matches = if manifest_exists { + read_manifest(manifest_path)? == *manifest + } else { + true + }; + let target_progress = manifest_progress(manifest)?; + let source_progress = source_split + .map(|(source, prefix)| split_status(source, prefix, None)) + .transpose()?; + let next_requested_window = requested_window.and_then(|requested| { + next_missing_window_in_range(&target_progress.missing_ranges, requested) + }); + let backend_check = check_backend_ready(backend_kind, backend_path); + let report = JobPreflight { + kind: manifest.kind, + manifest_path: manifest_path.to_path_buf(), + manifest_exists, + manifest_matches, + source_complete: source_progress.as_ref().map(is_complete), + source_shards: source_progress + .as_ref() + .map(|progress| progress.completed_count), + expected_source_shards: source_progress + .as_ref() + .map(|progress| progress.expected_splits), + source_missing_ranges: source_progress + .as_ref() + .map(|progress| progress_windows(&progress.missing_ranges)), + target_shards: target_progress.completed_count, + expected_target_shards: target_progress.expected_splits, + first_missing_target: target_progress.first_missing, + target_missing_ranges: progress_windows(&target_progress.missing_ranges), + next_window: target_progress.next_window.map(|window| ProgressWindow { + first_split: window.first_split, + last_split: window.last_split, + }), + requested_window: requested_window.map(ProgressWindow::from), + next_requested_window: next_requested_window.map(ProgressWindow::from), + backend_kind: backend_kind.as_str().to_string(), + backend_path: backend_path.map(Path::to_path_buf), + backend_ready: backend_check.ready, + backend_error: backend_check.error, + }; + print_preflight(&report, json)?; + ensure!( + report.manifest_matches, + "existing manifest does not match requested job" + ); + ensure!( + report.backend_ready, + "backend is not ready for {}: {} ({})", + backend_kind.as_str(), + backend_path + .map(|path| path.display().to_string()) + .unwrap_or_else(|| "".to_string()), + report + .backend_error + .as_deref() + .unwrap_or("no additional error") + ); + if let Some(false) = report.source_complete { + ensure!(false, "source split artifact is incomplete"); + } + Ok(()) +} + +struct BackendReady { + ready: bool, + error: Option, +} + +fn check_backend_ready(backend_kind: BackendKind, backend_path: Option<&Path>) -> BackendReady { + match backend_kind { + BackendKind::LlamaApi => check_llama_quant_runtime_ready(backend_path), + BackendKind::SkippyAbi => check_skippy_runtime_ready(backend_path), + BackendKind::NativeRust => BackendReady { + ready: true, + error: None, + }, + } +} + +fn check_llama_quant_runtime_ready(backend_path: Option<&Path>) -> BackendReady { + check_runtime_ready( + backend_path, + llama_quant_ffi::native_runtime_loaded, + |libraries| unsafe { llama_quant_ffi::load_native_runtime_libraries(libraries) }, + ) +} + +fn check_skippy_runtime_ready(backend_path: Option<&Path>) -> BackendReady { + check_runtime_ready( + backend_path, + skippy_ffi::native_runtime_loaded, + |libraries| unsafe { skippy_ffi::load_native_runtime_libraries(libraries) }, + ) +} + +fn check_runtime_ready( + backend_path: Option<&Path>, + loaded: impl Fn() -> bool, + load: impl FnOnce(&[PathBuf; 1]) -> Result<(), LoadError>, +) -> BackendReady +where + LoadError: std::fmt::Display, +{ + if loaded() { + return BackendReady { + ready: true, + error: None, + }; + } + let Some(path) = backend_path else { + return BackendReady { + ready: false, + error: Some("native runtime library path is missing".to_string()), + }; + }; + if !path.is_file() { + return BackendReady { + ready: false, + error: Some(format!( + "native runtime library is not a file: {}", + path.display() + )), + }; + } + let libraries = [path.to_path_buf()]; + match load(&libraries) { + Ok(()) => BackendReady { + ready: true, + error: None, + }, + Err(error) => BackendReady { + ready: false, + error: Some(error.to_string()), + }, + } +} + +fn ensure_backend_supported(kind: JobKind, backend_kind: BackendKind) -> Result<()> { + match kind { + JobKind::ConvertHf => ensure_convert_backend(backend_kind), + JobKind::QuantizeGguf => ensure_quant_backend(backend_kind), + } +} + +fn is_complete(progress: &Progress) -> bool { + progress.complete +} + +fn progress_windows(ranges: &[crate::splits::ShardRange]) -> Vec { + ranges + .iter() + .map(|range| ProgressWindow { + first_split: range.first_split, + last_split: range.last_split, + }) + .collect() +} + +fn print_preflight(report: &JobPreflight, json: bool) -> Result<()> { + if json { + print_json_pretty(report)?; + } else { + print_info(format!( + "Preflight {:?} with backend {}", + report.kind, report.backend_kind + )); + print_progress_line( + "target", + report.target_shards, + report.expected_target_shards, + ); + if report.manifest_matches { + print_success("Manifest is compatible"); + } else { + print_warn("Existing manifest does not match requested job"); + } + if report.backend_ready { + print_success("Backend is ready"); + } else { + print_warn("Backend is not ready"); + } + if let Some(error) = report.backend_error.as_deref() { + print_warn(format!("Backend error: {error}")); + } + if let (Some(source_shards), Some(expected), Some(complete)) = ( + report.source_shards, + report.expected_source_shards, + report.source_complete, + ) { + print_progress_line("source", source_shards, expected); + if complete { + print_success("Source artifact is complete"); + } else { + print_warn("Source artifact is incomplete"); + } + if let Some(ranges) = report.source_missing_ranges.as_deref() + && !ranges.is_empty() + { + print_info(format!("Source missing ranges: {}", format_ranges(ranges))); + } + } + if !report.target_missing_ranges.is_empty() { + print_info(format!( + "Target missing ranges: {}", + format_ranges(&report.target_missing_ranges) + )); + } + if let Some(requested) = report.requested_window.as_ref() { + print_info(format!( + "Requested window: {}; next requested window: {}", + format_ranges(std::slice::from_ref(requested)), + report + .next_requested_window + .as_ref() + .map(|window| format_ranges(std::slice::from_ref(window))) + .unwrap_or_else(|| "complete".to_string()) + )); + } + } + Ok(()) +} + +impl From for ProgressWindow { + fn from(window: SplitWindow) -> Self { + Self { + first_split: window.first_split, + last_split: window.last_split, + } + } +} + +fn format_ranges(ranges: &[ProgressWindow]) -> String { + ranges + .iter() + .map(|range| { + if range.first_split == range.last_split { + range.first_split.to_string() + } else { + format!("{}..{}", range.first_split, range.last_split) + } + }) + .collect::>() + .join(",") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn native_rust_backend_is_ready_without_path() { + let ready = check_backend_ready(BackendKind::NativeRust, None); + + assert!(ready.ready); + assert!(ready.error.is_none()); + } + + #[test] + fn native_runtime_backend_rejects_missing_or_invalid_library() { + let missing = check_backend_ready(BackendKind::SkippyAbi, None); + let executable = std::env::current_exe().unwrap(); + let invalid_library = check_backend_ready(BackendKind::SkippyAbi, Some(&executable)); + + assert!(!missing.ready); + assert!(missing.error.unwrap().contains("missing")); + assert!(!invalid_library.ready); + assert!(invalid_library.error.is_some()); + } +} diff --git a/crates/skippy-quantize/src/quantize.rs b/crates/skippy-quantize/src/quantize.rs new file mode 100644 index 0000000000..303fa77482 --- /dev/null +++ b/crates/skippy-quantize/src/quantize.rs @@ -0,0 +1,44 @@ +use anyhow::{Result, anyhow, ensure}; + +use crate::types::TensorType; + +pub fn ensure_tensor_type_entry(token: &str) -> Result<()> { + normalize_tensor_type_entry(token).map(|_| ()) +} + +pub fn normalize_tensor_type_entry(token: &str) -> Result { + let (name, raw_type) = token + .split_once('=') + .ok_or_else(|| anyhow!("malformed tensor type entry {token:?}"))?; + ensure!(!name.is_empty(), "tensor type entry has empty tensor name"); + ensure_raw_tensor_type(raw_type).map_err(|error| { + anyhow!("unsupported raw ggml tensor type {raw_type:?} in entry {token:?}: {error}") + })?; + Ok(format!("{}={raw_type}", name.to_ascii_lowercase())) +} + +fn ensure_raw_tensor_type(raw_type: &str) -> Result<()> { + ensure!( + TensorType::parse(raw_type).is_some(), + "unsupported raw ggml tensor type {raw_type:?}" + ); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn normalizes_tensor_type_entry() { + assert_eq!( + normalize_tensor_type_entry("MTP_Head.Weight=NVFP4").unwrap(), + "mtp_head.weight=NVFP4" + ); + } + + #[test] + fn rejects_unknown_tensor_type() { + assert!(normalize_tensor_type_entry("foo=NOT_A_TYPE").is_err()); + } +} diff --git a/crates/skippy-quantize/src/records.rs b/crates/skippy-quantize/src/records.rs new file mode 100644 index 0000000000..e625fb62e8 --- /dev/null +++ b/crates/skippy-quantize/src/records.rs @@ -0,0 +1,115 @@ +use std::fs; +use std::path::{Path, PathBuf}; +use std::time::{SystemTime, UNIX_EPOCH}; + +use anyhow::{Context, Result}; +use serde::Serialize; + +use crate::backend::BackendRunStatus; +use crate::output::print_path_event; +use crate::splits::SplitWindow; +use crate::types::JobKind; + +#[derive(Debug, Serialize)] +struct WindowRunRecord { + schema_version: u32, + kind: JobKind, + started_unix_ms: u128, + duration_ms: u128, + first_split: u32, + last_split: u32, + output_prefix: PathBuf, + command: Vec, + status_code: Option, + success: bool, +} + +pub struct WindowRunRecordInput<'a> { + pub schema_version: u32, + pub kind: JobKind, + pub command: &'a [String], + pub output_prefix: &'a Path, + pub window: SplitWindow, + pub status: BackendRunStatus, + pub duration_ms: u128, + pub started_unix_ms: u128, +} + +pub fn write_window_record( + record_dir: Option<&Path>, + input: WindowRunRecordInput<'_>, +) -> Result<()> { + let Some(record_dir) = record_dir else { + return Ok(()); + }; + fs::create_dir_all(record_dir).with_context(|| format!("create {}", record_dir.display()))?; + let record = WindowRunRecord { + schema_version: input.schema_version, + kind: input.kind, + started_unix_ms: input.started_unix_ms, + duration_ms: input.duration_ms, + first_split: input.window.first_split, + last_split: input.window.last_split, + output_prefix: input.output_prefix.to_path_buf(), + command: input.command.to_vec(), + status_code: input.status.status_code, + success: input.status.success, + }; + let name = format!( + "{:?}-{:05}-{:05}-{}.json", + input.kind, input.window.first_split, input.window.last_split, input.started_unix_ms + ) + .to_lowercase(); + let path = record_dir.join(name); + fs::write(&path, serde_json::to_vec_pretty(&record)?) + .with_context(|| format!("write {}", path.display()))?; + print_path_event("🧾", "Wrote window record", &path); + Ok(()) +} + +pub fn unix_timestamp_ms() -> u128 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |duration| duration.as_millis()) +} + +#[cfg(test)] +mod tests { + use crate::manifest::MANIFEST_VERSION; + use crate::splits::SplitWindow; + use crate::types::JobKind; + + use super::*; + + #[test] + fn writes_window_run_records() { + let root = std::env::temp_dir().join(format!( + "skippy-quantize-record-test-{}", + unix_timestamp_ms() + )); + write_window_record( + Some(&root), + WindowRunRecordInput { + schema_version: MANIFEST_VERSION, + kind: JobKind::QuantizeGguf, + command: &["llama-quantize".to_string()], + output_prefix: Path::new("/target/Q2_K/out.gguf"), + window: SplitWindow { + first_split: 3, + last_split: 4, + }, + status: BackendRunStatus { + status_code: Some(0), + success: true, + }, + duration_ms: 10, + started_unix_ms: 1234, + }, + ) + .unwrap(); + + let entries = fs::read_dir(&root).unwrap().count(); + assert_eq!(entries, 1); + fs::remove_dir_all(root).unwrap(); + } +} diff --git a/crates/skippy-quantize/src/residency.rs b/crates/skippy-quantize/src/residency.rs new file mode 100644 index 0000000000..a82d55d049 --- /dev/null +++ b/crates/skippy-quantize/src/residency.rs @@ -0,0 +1,87 @@ +use std::fs; +use std::io::{Read, Write}; +use std::path::Path; + +use anyhow::{Context, Result}; + +use crate::output::{format_bytes, print_success}; + +const COPY_CHUNK_BYTES: usize = 16 * 1024 * 1024; + +pub fn remove_dir_if_exists(path: &Path) -> Result<()> { + if path.exists() { + fs::remove_dir_all(path).with_context(|| format!("remove {}", path.display()))?; + } + Ok(()) +} + +pub fn copy_file_bounded_with_label(label: &str, source: &Path, target: &Path) -> Result<()> { + let mut input = fs::File::open(source).with_context(|| format!("open {}", source.display()))?; + let mut output = + fs::File::create(target).with_context(|| format!("create {}", target.display()))?; + let mut buffer = vec![0_u8; COPY_CHUNK_BYTES]; + let mut copied = 0_u64; + + loop { + let count = input + .read(&mut buffer) + .with_context(|| format!("read {}", source.display()))?; + if count == 0 { + break; + } + output + .write_all(&buffer[..count]) + .with_context(|| format!("write {}", target.display()))?; + drop_file_cache_range(&input, copied, count as u64); + copied += count as u64; + } + + output + .sync_all() + .with_context(|| format!("sync {}", target.display()))?; + let target_dropped = drop_file_cache_range(&output, 0, copied); + print_success(format!( + "{label} copied {} -> {} ({}, cache_dropped={target_dropped})", + source.display(), + target.display(), + format_bytes(copied) + )); + Ok(()) +} + +#[cfg(unix)] +pub fn symlink_file(source: &Path, target: &Path) -> Result<()> { + std::os::unix::fs::symlink(source, target) + .with_context(|| format!("symlink {} -> {}", target.display(), source.display())) +} + +#[cfg(not(unix))] +pub fn symlink_file(source: &Path, target: &Path) -> Result<()> { + fs::copy(source, target) + .with_context(|| format!("copy {} -> {}", source.display(), target.display()))?; + Ok(()) +} + +#[cfg(target_os = "linux")] +fn drop_file_cache_range(file: &fs::File, offset: u64, len: u64) -> bool { + use std::os::fd::AsRawFd; + + if len == 0 { + return false; + } + + let rc = unsafe { + libc::posix_fadvise( + file.as_raw_fd(), + offset as libc::off_t, + len as libc::off_t, + libc::POSIX_FADV_DONTNEED, + ) + }; + rc == 0 +} + +#[cfg(not(target_os = "linux"))] +fn drop_file_cache_range(_file: &fs::File, _offset: u64, _len: u64) -> bool { + false +} diff --git a/crates/skippy-quantize/src/splits.rs b/crates/skippy-quantize/src/splits.rs new file mode 100644 index 0000000000..1e1b7e7ce8 --- /dev/null +++ b/crates/skippy-quantize/src/splits.rs @@ -0,0 +1,557 @@ +use std::collections::BTreeSet; +use std::fs; +use std::path::{Path, PathBuf}; + +use anyhow::{Context, Result, ensure}; +use serde::{Deserialize, Serialize}; + +use crate::output::{print_copy, print_info}; +use crate::residency::{copy_file_bounded_with_label, remove_dir_if_exists, symlink_file}; + +#[derive(Debug, Serialize)] +pub struct Progress { + pub expected_splits: u32, + pub completed_count: usize, + pub missing_count: usize, + pub missing_ranges: Vec, + pub completed_percent: f64, + pub complete: bool, + pub first_missing: Option, + pub last_present: Option, + pub next_window: Option, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +pub struct SplitWindow { + pub first_split: u32, + pub last_split: u32, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +pub struct ShardRange { + pub first_split: u32, + pub last_split: u32, +} + +pub fn split_status(root: &Path, prefix: &str, expected_splits: Option) -> Result { + let scan_root = root.join(prefix); + let mut seen = BTreeSet::new(); + if scan_root.exists() { + for entry in fs::read_dir(&scan_root) + .with_context(|| format!("read directory {}", scan_root.display()))? + { + let entry = entry?; + let file_name = entry.file_name(); + let file_name = file_name.to_string_lossy(); + if let Some((index, total)) = parse_split_file_name(&file_name) + && expected_splits.is_none_or(|expected| expected == total) + { + seen.insert(index); + } + } + if seen.is_empty() && expected_splits.is_none() && has_single_unsplit_gguf(&scan_root)? { + seen.insert(1); + } + } + + let expected = expected_splits.unwrap_or_else(|| { + seen.iter() + .next_back() + .copied() + .unwrap_or_default() + .max(parse_expected_total_from_dir(root, prefix).unwrap_or_default()) + }); + let first_missing = (1..=expected).find(|index| !seen.contains(index)); + Ok(progress_from_seen(expected, seen, first_missing)) +} + +pub fn split_status_for_basename( + root: &Path, + prefix: &str, + basename: &str, + expected_splits: u32, +) -> Result { + let scan_root = root.join(prefix); + let mut seen = BTreeSet::new(); + if scan_root.exists() { + for index in 1..=expected_splits { + let name = format!("{basename}-{index:05}-of-{expected_splits:05}.gguf"); + if scan_root.join(name).is_file() { + seen.insert(index); + } + } + if expected_splits == 1 && scan_root.join(format!("{basename}.gguf")).is_file() { + seen.insert(1); + } + } + let first_missing = (1..=expected_splits).find(|index| !seen.contains(index)); + Ok(progress_from_seen(expected_splits, seen, first_missing)) +} + +fn progress_from_seen( + expected_splits: u32, + seen: BTreeSet, + first_missing: Option, +) -> Progress { + let completed_count = seen.len(); + let expected_count = expected_splits as usize; + let missing_count = expected_count.saturating_sub(completed_count); + let complete = expected_count > 0 && missing_count == 0 && first_missing.is_none(); + let completed_percent = if expected_splits == 0 { + 0.0 + } else { + (completed_count as f64 / f64::from(expected_splits)) * 100.0 + }; + let missing_ranges = missing_ranges(expected_splits, &seen); + Progress { + expected_splits, + completed_count, + missing_count, + missing_ranges, + completed_percent, + complete, + first_missing, + last_present: seen.iter().next_back().copied(), + next_window: None, + } +} + +fn missing_ranges(expected_splits: u32, seen: &BTreeSet) -> Vec { + let mut ranges = Vec::new(); + let mut current_start = None; + let mut previous_missing = None; + + for index in 1..=expected_splits { + if seen.contains(&index) { + if let Some(start) = current_start.take() { + ranges.push(ShardRange { + first_split: start, + last_split: previous_missing.expect("missing range has previous index"), + }); + } + previous_missing = None; + continue; + } + + current_start.get_or_insert(index); + previous_missing = Some(index); + } + + if let Some(start) = current_start { + ranges.push(ShardRange { + first_split: start, + last_split: previous_missing.expect("missing range has previous index"), + }); + } + + ranges +} + +pub fn parse_split_file_name(file_name: &str) -> Option<(u32, u32)> { + let stem = file_name.strip_suffix(".gguf")?; + let (before_total, total) = stem.rsplit_once("-of-")?; + let (_, index) = before_total.rsplit_once('-')?; + Some((index.parse().ok()?, total.parse().ok()?)) +} + +pub fn shard_name_for(first_shard: &Path, index: u32, total: u32) -> Result { + let first_name = first_shard + .file_name() + .and_then(|name| name.to_str()) + .with_context(|| format!("invalid first shard path {}", first_shard.display()))?; + let stem = first_name + .strip_suffix(".gguf") + .with_context(|| format!("first shard is not a GGUF file: {}", first_shard.display()))?; + if total == 1 && index == 1 && parse_split_file_name(first_name).is_none() { + return Ok(first_name.to_string()); + } + let (before_total, _) = stem + .rsplit_once("-of-") + .with_context(|| format!("invalid GGUF split shard name: {}", first_shard.display()))?; + let (base, _) = before_total + .rsplit_once('-') + .with_context(|| format!("invalid GGUF split shard name: {}", first_shard.display()))?; + Ok(format!("{base}-{index:05}-of-{total:05}.gguf")) +} + +pub fn find_first_shard(source: &Path, source_prefix: &str) -> Result { + let source_root = prefixed_path(source, source_prefix); + let mut candidates = Vec::new(); + let mut unsplit_candidates = Vec::new(); + for entry in fs::read_dir(&source_root) + .with_context(|| format!("read directory {}", source_root.display()))? + { + let entry = entry?; + let file_name = entry.file_name(); + let file_name = file_name.to_string_lossy(); + if let Some((index, _)) = parse_split_file_name(&file_name) + && index == 1 + { + candidates.push(entry.path()); + } else if is_unsplit_gguf_name(&file_name) { + unsplit_candidates.push(entry.path()); + } + } + if candidates.is_empty() && unsplit_candidates.len() == 1 { + return Ok(unsplit_candidates.remove(0)); + } + ensure!( + !candidates.is_empty(), + "no first GGUF split shard found under {}", + source_root.display() + ); + ensure!( + candidates.len() == 1, + "multiple first GGUF split shards found under {}", + source_root.display() + ); + Ok(candidates.remove(0)) +} + +pub fn stage_source_window( + source: &Path, + source_prefix: &str, + first_source_shard: &Path, + stage_path: &Path, + window: SplitWindow, + total: u32, +) -> Result { + remove_dir_if_exists(stage_path)?; + let stage_root = prefixed_path(stage_path, source_prefix); + fs::create_dir_all(&stage_root).with_context(|| format!("create {}", stage_root.display()))?; + let source_root = prefixed_path(source, source_prefix); + + for index in 1..=total { + let name = shard_name_for(first_source_shard, index, total)?; + let source_shard = source_root.join(&name); + ensure!( + source_shard.is_file(), + "source shard does not exist: {}", + source_shard.display() + ); + let staged_shard = stage_root.join(name); + if window.first_split <= index && index <= window.last_split { + print_copy( + &source_shard, + &staged_shard, + source_shard.metadata().ok().map(|m| m.len()), + ); + copy_file_bounded_with_label("stage_source_copy", &source_shard, &staged_shard)?; + } else { + symlink_file(&source_shard, &staged_shard)?; + } + } + + let staged_first = stage_root.join(shard_name_for(first_source_shard, 1, total)?); + print_info(format!( + "Staged source window {}..{} at {}", + window.first_split, + window.last_split, + stage_path.display() + )); + Ok(staged_first) +} + +pub fn next_missing_window(missing_ranges: &[ShardRange], window_size: u32) -> Option { + let first_range = missing_ranges.first()?; + let capped_last = first_range + .first_split + .saturating_add(window_size.max(1).saturating_sub(1)) + .min(first_range.last_split); + Some(SplitWindow { + first_split: first_range.first_split, + last_split: capped_last, + }) +} + +pub fn next_missing_window_in_range( + missing_ranges: &[ShardRange], + requested: SplitWindow, +) -> Option { + missing_ranges.iter().find_map(|range| { + let first_split = range.first_split.max(requested.first_split); + let last_split = range.last_split.min(requested.last_split); + (first_split <= last_split).then_some(SplitWindow { + first_split, + last_split, + }) + }) +} + +pub fn validate_split_window(window: SplitWindow, expected_splits: u32) -> Result<()> { + ensure!( + window.first_split > 0, + "first split must be greater than zero" + ); + ensure!( + window.first_split <= window.last_split, + "first split {} must be <= last split {}", + window.first_split, + window.last_split + ); + ensure!( + window.last_split <= expected_splits, + "last split {} exceeds expected split count {}", + window.last_split, + expected_splits + ); + Ok(()) +} + +fn parse_expected_total_from_dir(root: &Path, prefix: &str) -> Option { + let scan_root = root.join(prefix); + let entries = fs::read_dir(scan_root).ok()?; + entries.filter_map(Result::ok).find_map(|entry| { + parse_split_file_name(&entry.file_name().to_string_lossy()).map(|(_, total)| total) + }) +} + +fn has_single_unsplit_gguf(root: &Path) -> Result { + let mut count = 0_u32; + for entry in fs::read_dir(root).with_context(|| format!("read directory {}", root.display()))? { + let entry = entry?; + if is_unsplit_gguf_name(&entry.file_name().to_string_lossy()) { + count += 1; + } + } + Ok(count == 1) +} + +fn is_unsplit_gguf_name(file_name: &str) -> bool { + file_name.ends_with(".gguf") && parse_split_file_name(file_name).is_none() +} + +fn prefixed_path(root: &Path, prefix: &str) -> PathBuf { + if prefix.is_empty() { + root.to_path_buf() + } else { + root.join(prefix) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_split_file_names() { + assert_eq!( + parse_split_file_name("GLM-5.2-Q2_K-MTP-Q8-00174-of-00306.gguf"), + Some((174, 306)) + ); + assert_eq!(parse_split_file_name("not-a-shard.gguf"), None); + } + + #[test] + fn plans_next_window_from_first_missing_range() { + let ranges = vec![ + ShardRange { + first_split: 2, + last_split: 2, + }, + ShardRange { + first_split: 5, + last_split: 9, + }, + ]; + assert_eq!( + next_missing_window(&ranges, 4).map(|w| (w.first_split, w.last_split)), + Some((2, 2)) + ); + assert_eq!( + next_missing_window(&ranges[1..], 3).map(|w| (w.first_split, w.last_split)), + Some((5, 7)) + ); + assert_eq!( + next_missing_window(&ranges[1..], 0).map(|w| (w.first_split, w.last_split)), + Some((5, 5)) + ); + assert!(next_missing_window(&[], 4).is_none()); + } + + #[test] + fn plans_next_missing_window_inside_requested_range() { + let ranges = vec![ + ShardRange { + first_split: 2, + last_split: 3, + }, + ShardRange { + first_split: 7, + last_split: 9, + }, + ]; + + assert_eq!( + next_missing_window_in_range( + &ranges, + SplitWindow { + first_split: 3, + last_split: 8 + } + ) + .map(|w| (w.first_split, w.last_split)), + Some((3, 3)) + ); + assert_eq!( + next_missing_window_in_range( + &ranges[1..], + SplitWindow { + first_split: 3, + last_split: 8 + } + ) + .map(|w| (w.first_split, w.last_split)), + Some((7, 8)) + ); + assert!( + next_missing_window_in_range( + &ranges, + SplitWindow { + first_split: 4, + last_split: 6 + } + ) + .is_none() + ); + } + + #[test] + fn validates_requested_split_window_bounds() { + assert!( + validate_split_window( + SplitWindow { + first_split: 1, + last_split: 3 + }, + 3 + ) + .is_ok() + ); + assert!( + validate_split_window( + SplitWindow { + first_split: 0, + last_split: 1 + }, + 3 + ) + .is_err() + ); + assert!( + validate_split_window( + SplitWindow { + first_split: 3, + last_split: 2 + }, + 3 + ) + .is_err() + ); + assert!( + validate_split_window( + SplitWindow { + first_split: 1, + last_split: 4 + }, + 3 + ) + .is_err() + ); + } + + #[test] + fn derives_matching_shard_names_from_first_shard() { + let first = Path::new("/repo/BF16/GLM-5.2-BF16-00001-of-00306.gguf"); + assert_eq!( + shard_name_for(first, 174, 306).unwrap(), + "GLM-5.2-BF16-00174-of-00306.gguf" + ); + } + + #[test] + fn derives_single_shard_name_from_unsplit_gguf() { + let first = Path::new("/repo/BF16/model.gguf"); + assert_eq!(shard_name_for(first, 1, 1).unwrap(), "model.gguf"); + } + + #[test] + fn validates_exact_basename_splits() { + let root = std::env::temp_dir().join(format!( + "skippy-quantize-test-{}", + crate::unix_timestamp_ms() + )); + let prefix_root = root.join("Q2_K"); + fs::create_dir_all(&prefix_root).unwrap(); + fs::write(prefix_root.join("out-00001-of-00003.gguf"), b"1").unwrap(); + fs::write(prefix_root.join("out-00003-of-00003.gguf"), b"3").unwrap(); + + let progress = split_status_for_basename(&root, "Q2_K", "out", 3).unwrap(); + assert_eq!(progress.expected_splits, 3); + assert_eq!(progress.completed_count, 2); + assert_eq!(progress.missing_count, 1); + assert_eq!( + progress.missing_ranges, + vec![ShardRange { + first_split: 2, + last_split: 2, + }] + ); + assert!((progress.completed_percent - (200.0 / 3.0)).abs() < 1e-9); + assert!(!progress.complete); + assert_eq!(progress.first_missing, Some(2)); + assert_eq!(progress.last_present, Some(3)); + + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn validates_exact_unsplit_single_file() { + let root = std::env::temp_dir().join(format!( + "skippy-quantize-unsplit-status-test-{}", + crate::unix_timestamp_ms() + )); + let prefix_root = root.join("Q2_K"); + fs::create_dir_all(&prefix_root).unwrap(); + fs::write(prefix_root.join("out.gguf"), b"one").unwrap(); + + let progress = split_status_for_basename(&root, "Q2_K", "out", 1).unwrap(); + assert_eq!(progress.completed_count, 1); + assert_eq!(progress.missing_count, 0); + assert!(progress.missing_ranges.is_empty()); + assert_eq!(progress.completed_percent, 100.0); + assert!(progress.complete); + assert_eq!(progress.first_missing, None); + + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn reports_compact_missing_ranges() { + let mut seen = BTreeSet::new(); + seen.insert(1); + seen.insert(4); + seen.insert(8); + + let progress = progress_from_seen(9, seen, Some(2)); + + assert_eq!(progress.missing_count, 6); + assert_eq!( + progress.missing_ranges, + vec![ + ShardRange { + first_split: 2, + last_split: 3, + }, + ShardRange { + first_split: 5, + last_split: 7, + }, + ShardRange { + first_split: 9, + last_split: 9, + }, + ] + ); + } +} diff --git a/crates/skippy-quantize/src/tensor_map.rs b/crates/skippy-quantize/src/tensor_map.rs new file mode 100644 index 0000000000..3fa4b3920b --- /dev/null +++ b/crates/skippy-quantize/src/tensor_map.rs @@ -0,0 +1,412 @@ +use anyhow::{Result, anyhow, bail}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum TensorNameMap { + Raw, + HfToGguf, + HfToGgufWithMtp { layer_start: u32 }, +} + +impl TensorNameMap { + pub(crate) fn map_tensor_name(self, name: &str) -> Result { + match self { + Self::Raw => Ok(name.to_string()), + Self::HfToGguf => map_hf_to_gguf(name, None), + Self::HfToGgufWithMtp { layer_start } => map_hf_to_gguf(name, Some(layer_start)), + } + } +} + +fn map_hf_to_gguf(name: &str, mtp_layer_start: Option) -> Result { + if let Some(layer_start) = mtp_layer_start + && let Some(normalized) = normalize_qwen_mtp_source_name(name, layer_start)? + { + return map_hf_to_gguf(&normalized, mtp_layer_start); + } + if name == "model.embed_tokens.weight" { + return Ok("token_embd.weight".to_string()); + } + if name == "embed_tokens.weight" { + return Ok("token_embd.weight".to_string()); + } + if name == "lm_head.weight" { + return Ok("output.weight".to_string()); + } + if name == "model.norm.weight" { + return Ok("output_norm.weight".to_string()); + } + if name == "norm.weight" { + return Ok("output_norm.weight".to_string()); + } + if let Some(name) = map_mtp_source_tensor(name)? { + return Ok(name); + } + let Some(layer) = HfLayerTensor::parse(name)? else { + bail!("no HF->GGUF tensor mapping for {name}"); + }; + layer.map() +} + +pub(crate) fn is_mtp_source_tensor(name: &str) -> bool { + is_qwen_mtp_source_tensor(name) + || map_mtp_source_tensor(name).is_ok_and(|mapped| mapped.is_some()) +} + +pub(crate) fn hf_layer_id(name: &str) -> Result> { + let Some(rest) = name.strip_prefix("model.layers.") else { + return Ok(None); + }; + let Some((layer, _suffix)) = rest.split_once('.') else { + bail!("malformed layer tensor name {name}"); + }; + layer + .parse::() + .map(Some) + .map_err(|err| anyhow!("malformed layer id in {name}: {err}")) +} + +pub(crate) fn is_shared_mtp_context_tensor(name: &str) -> bool { + matches!( + name, + "model.embed_tokens.weight" + | "embed_tokens.weight" + | "model.norm.weight" + | "norm.weight" + | "lm_head.weight" + ) +} + +fn map_mtp_source_tensor(name: &str) -> Result> { + let mapped = match name { + "pre_projection" | "pre_projection.weight" => "nextn.pre_projection.weight".to_string(), + "post_projection" | "post_projection.weight" => "nextn.post_projection.weight".to_string(), + "d2t" | "d2t.weight" => "d2t.weight".to_string(), + _ => { + let Some(layer) = HfLayerTensor::parse(name)? else { + return Ok(None); + }; + let bid = layer.layer; + match layer.suffix { + "eh_proj" | "eh_proj.weight" => format!("blk.{bid}.nextn.eh_proj.weight"), + "embed_tokens" | "embed_tokens.weight" => { + format!("blk.{bid}.nextn.embed_tokens.weight") + } + "enorm" | "enorm.weight" => format!("blk.{bid}.nextn.enorm.weight"), + "hnorm" | "hnorm.weight" => format!("blk.{bid}.nextn.hnorm.weight"), + "shared_head.head" | "shared_head.head.weight" | "shared_head.output.weight" => { + format!("blk.{bid}.nextn.shared_head_head.weight") + } + "shared_head.norm" | "shared_head.norm.weight" => { + format!("blk.{bid}.nextn.shared_head_norm.weight") + } + _ => return Ok(None), + } + } + }; + Ok(Some(mapped)) +} + +fn is_qwen_mtp_source_tensor(name: &str) -> bool { + let name = name.strip_prefix("model.").unwrap_or(name); + name.starts_with("mtp.") +} + +fn normalize_qwen_mtp_source_name(name: &str, layer_start: u32) -> Result> { + let name = name.strip_prefix("model.").unwrap_or(name); + if !name.starts_with("mtp.") { + return Ok(None); + } + let parts = name.splitn(4, '.').collect::>(); + if parts.len() == 4 && parts[1] == "layers" { + let mtp_idx = parts[2] + .parse::() + .map_err(|err| anyhow!("malformed MTP layer id in {name}: {err}"))?; + return Ok(Some(format!( + "model.layers.{}.{}", + layer_start + mtp_idx, + parts[3] + ))); + } + if parts.len() == 3 { + let suffix = match parts[1] { + "fc" => "eh_proj", + "pre_fc_norm_embedding" => "enorm", + "pre_fc_norm_hidden" => "hnorm", + "norm" => "shared_head.norm", + _ => return Ok(None), + }; + return Ok(Some(format!( + "model.layers.{layer_start}.{suffix}.{}", + parts[2] + ))); + } + Ok(None) +} + +struct HfLayerTensor<'a> { + layer: u32, + suffix: &'a str, +} + +impl<'a> HfLayerTensor<'a> { + fn parse(name: &'a str) -> Result> { + let Some(rest) = name.strip_prefix("model.layers.") else { + return Ok(None); + }; + let Some((layer, suffix)) = rest.split_once('.') else { + bail!("malformed layer tensor name {name}"); + }; + let layer = layer + .parse::() + .map_err(|err| anyhow!("malformed layer id in {name}: {err}"))?; + Ok(Some(Self { layer, suffix })) + } + + fn map(&self) -> Result { + let bid = self.layer; + match self.suffix { + "input_layernorm.weight" => Ok(format!("blk.{bid}.attn_norm.weight")), + "post_attention_layernorm.weight" => Ok(format!("blk.{bid}.ffn_norm.weight")), + "self_attn.q_proj.weight" => Ok(format!("blk.{bid}.attn_q.weight")), + "self_attn.k_proj.weight" => Ok(format!("blk.{bid}.attn_k.weight")), + "self_attn.v_proj.weight" => Ok(format!("blk.{bid}.attn_v.weight")), + "self_attn.q_proj.bias" => Ok(format!("blk.{bid}.attn_q.bias")), + "self_attn.k_proj.bias" => Ok(format!("blk.{bid}.attn_k.bias")), + "self_attn.v_proj.bias" => Ok(format!("blk.{bid}.attn_v.bias")), + "self_attn.o_proj.weight" => Ok(format!("blk.{bid}.attn_output.weight")), + "self_attn.q_norm.weight" => Ok(format!("blk.{bid}.attn_q_norm.weight")), + "self_attn.k_norm.weight" => Ok(format!("blk.{bid}.attn_k_norm.weight")), + "self_attn.q_a_proj.weight" => Ok(format!("blk.{bid}.attn_q_a.weight")), + "self_attn.q_b_proj.weight" => Ok(format!("blk.{bid}.attn_q_b.weight")), + "self_attn.q_a_layernorm.weight" => Ok(format!("blk.{bid}.attn_q_a_norm.weight")), + "self_attn.kv_a_proj_with_mqa.weight" => Ok(format!("blk.{bid}.attn_kv_a_mqa.weight")), + "self_attn.kv_b_proj.weight" => Ok(format!("blk.{bid}.attn_kv_b.weight")), + "self_attn.kv_a_layernorm.weight" => Ok(format!("blk.{bid}.attn_kv_a_norm.weight")), + "self_attn.indexer.k_norm.weight" => Ok(format!("blk.{bid}.indexer.k_norm.weight")), + "self_attn.indexer.k_norm.bias" => Ok(format!("blk.{bid}.indexer.k_norm.bias")), + "self_attn.indexer.weights_proj.weight" => Ok(format!("blk.{bid}.indexer.proj.weight")), + "self_attn.indexer.wk.weight" => Ok(format!("blk.{bid}.indexer.attn_k.weight")), + "self_attn.indexer.wq_b.weight" => Ok(format!("blk.{bid}.indexer.attn_q_b.weight")), + "mlp.down_proj.weight" => Ok(format!("blk.{bid}.ffn_down.weight")), + "mlp.gate_proj.weight" => Ok(format!("blk.{bid}.ffn_gate.weight")), + "mlp.up_proj.weight" => Ok(format!("blk.{bid}.ffn_up.weight")), + "mlp.gate.weight" => Ok(format!("blk.{bid}.ffn_gate_inp.weight")), + "mlp.shared_expert_gate" => Ok(format!("blk.{bid}.ffn_gate_inp_shexp.weight")), + "mlp.gate.e_score_correction_bias" => Ok(format!("blk.{bid}.exp_probs_b.bias")), + "mlp.shared_expert.down_proj.weight" => Ok(format!("blk.{bid}.ffn_down_shexp.weight")), + "mlp.shared_expert.gate_proj.weight" => Ok(format!("blk.{bid}.ffn_gate_shexp.weight")), + "mlp.shared_expert.up_proj.weight" => Ok(format!("blk.{bid}.ffn_up_shexp.weight")), + "mlp.shared_experts.down_proj.weight" => Ok(format!("blk.{bid}.ffn_down_shexp.weight")), + "mlp.shared_experts.gate_proj.weight" => Ok(format!("blk.{bid}.ffn_gate_shexp.weight")), + "mlp.shared_experts.up_proj.weight" => Ok(format!("blk.{bid}.ffn_up_shexp.weight")), + suffix if suffix.starts_with("mlp.experts.") => { + bail!( + "expert source tensor {suffix} requires streaming expert merge before GGUF write" + ) + } + suffix => bail!("no HF->GGUF mapping for layer tensor suffix {suffix:?}"), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn maps_glm_moe_lite_direct_tensor_names() { + assert_eq!( + TensorNameMap::HfToGguf + .map_tensor_name("model.embed_tokens.weight") + .unwrap(), + "token_embd.weight" + ); + assert_eq!( + TensorNameMap::HfToGguf + .map_tensor_name("model.layers.7.self_attn.kv_a_proj_with_mqa.weight") + .unwrap(), + "blk.7.attn_kv_a_mqa.weight" + ); + assert_eq!( + TensorNameMap::HfToGguf + .map_tensor_name("model.layers.7.mlp.shared_experts.gate_proj.weight") + .unwrap(), + "blk.7.ffn_gate_shexp.weight" + ); + } + + #[test] + fn maps_glm_dsa_indexer_tensor_names() { + assert_eq!( + TensorNameMap::HfToGguf + .map_tensor_name("model.layers.7.self_attn.indexer.k_norm.weight") + .unwrap(), + "blk.7.indexer.k_norm.weight" + ); + assert_eq!( + TensorNameMap::HfToGguf + .map_tensor_name("model.layers.7.self_attn.indexer.k_norm.bias") + .unwrap(), + "blk.7.indexer.k_norm.bias" + ); + assert_eq!( + TensorNameMap::HfToGguf + .map_tensor_name("model.layers.7.self_attn.indexer.weights_proj.weight") + .unwrap(), + "blk.7.indexer.proj.weight" + ); + assert_eq!( + TensorNameMap::HfToGguf + .map_tensor_name("model.layers.7.self_attn.indexer.wk.weight") + .unwrap(), + "blk.7.indexer.attn_k.weight" + ); + assert_eq!( + TensorNameMap::HfToGguf + .map_tensor_name("model.layers.7.self_attn.indexer.wq_b.weight") + .unwrap(), + "blk.7.indexer.attn_q_b.weight" + ); + } + + #[test] + fn rejects_expert_source_tensors_until_merge_exists() { + let err = TensorNameMap::HfToGguf + .map_tensor_name("model.layers.1.mlp.experts.0.down_proj.weight") + .unwrap_err() + .to_string(); + + assert!(err.contains("requires streaming expert merge")); + } + + #[test] + fn maps_qwen_dense_tensor_names() { + assert_eq!( + TensorNameMap::HfToGguf + .map_tensor_name("model.layers.3.self_attn.q_proj.weight") + .unwrap(), + "blk.3.attn_q.weight" + ); + assert_eq!( + TensorNameMap::HfToGguf + .map_tensor_name("model.layers.3.self_attn.k_norm.weight") + .unwrap(), + "blk.3.attn_k_norm.weight" + ); + assert_eq!( + TensorNameMap::HfToGguf + .map_tensor_name("model.layers.3.self_attn.v_proj.bias") + .unwrap(), + "blk.3.attn_v.bias" + ); + assert_eq!( + TensorNameMap::HfToGguf + .map_tensor_name("model.layers.3.mlp.up_proj.weight") + .unwrap(), + "blk.3.ffn_up.weight" + ); + } + + #[test] + fn maps_qwen2_moe_shared_expert_tensor_names() { + assert_eq!( + TensorNameMap::HfToGguf + .map_tensor_name("model.layers.3.mlp.shared_expert_gate") + .unwrap(), + "blk.3.ffn_gate_inp_shexp.weight" + ); + assert_eq!( + TensorNameMap::HfToGguf + .map_tensor_name("model.layers.3.mlp.shared_expert.gate_proj.weight") + .unwrap(), + "blk.3.ffn_gate_shexp.weight" + ); + assert_eq!( + TensorNameMap::HfToGguf + .map_tensor_name("model.layers.3.mlp.shared_expert.down_proj.weight") + .unwrap(), + "blk.3.ffn_down_shexp.weight" + ); + assert_eq!( + TensorNameMap::HfToGguf + .map_tensor_name("model.layers.3.mlp.shared_expert.up_proj.weight") + .unwrap(), + "blk.3.ffn_up_shexp.weight" + ); + } + + #[test] + fn recognizes_known_nextn_mtp_source_tensors() { + for name in [ + "pre_projection", + "pre_projection.weight", + "post_projection", + "d2t", + "model.layers.47.eh_proj.weight", + "model.layers.47.embed_tokens.weight", + "model.layers.47.enorm.weight", + "model.layers.47.hnorm.weight", + "model.layers.47.shared_head.head.weight", + "model.layers.47.shared_head.norm.weight", + ] { + assert!(is_mtp_source_tensor(name), "{name}"); + } + assert!(!is_mtp_source_tensor( + "model.layers.0.self_attn.q_proj.weight" + )); + assert!(!is_mtp_source_tensor("model.embed_tokens.weight")); + } + + #[test] + fn maps_known_nextn_mtp_source_tensors() { + assert_eq!( + TensorNameMap::HfToGguf + .map_tensor_name("model.layers.47.eh_proj.weight") + .unwrap(), + "blk.47.nextn.eh_proj.weight" + ); + assert_eq!( + TensorNameMap::HfToGguf + .map_tensor_name("model.layers.47.shared_head.norm.weight") + .unwrap(), + "blk.47.nextn.shared_head_norm.weight" + ); + assert_eq!( + TensorNameMap::HfToGguf + .map_tensor_name("embed_tokens.weight") + .unwrap(), + "token_embd.weight" + ); + } + + #[test] + fn recognizes_and_maps_qwen_style_mtp_source_tensors() { + for name in [ + "mtp.fc.weight", + "model.mtp.pre_fc_norm_embedding.weight", + "mtp.pre_fc_norm_hidden.weight", + "mtp.norm.weight", + "mtp.layers.0.self_attn.q_proj.weight", + ] { + assert!(is_mtp_source_tensor(name), "{name}"); + } + assert_eq!( + TensorNameMap::HfToGgufWithMtp { layer_start: 32 } + .map_tensor_name("mtp.fc.weight") + .unwrap(), + "blk.32.nextn.eh_proj.weight" + ); + assert_eq!( + TensorNameMap::HfToGgufWithMtp { layer_start: 32 } + .map_tensor_name("model.mtp.norm.weight") + .unwrap(), + "blk.32.nextn.shared_head_norm.weight" + ); + assert_eq!( + TensorNameMap::HfToGgufWithMtp { layer_start: 32 } + .map_tensor_name("mtp.layers.1.self_attn.q_proj.weight") + .unwrap(), + "blk.33.attn_q.weight" + ); + } +} diff --git a/crates/skippy-quantize/src/tokenizer_metadata.rs b/crates/skippy-quantize/src/tokenizer_metadata.rs new file mode 100644 index 0000000000..a6377f9cc0 --- /dev/null +++ b/crates/skippy-quantize/src/tokenizer_metadata.rs @@ -0,0 +1,553 @@ +use std::collections::{BTreeMap, BTreeSet}; +use std::fs; +use std::path::Path; + +use anyhow::{Context, Result, ensure}; +use serde_json::Value; + +use crate::gguf_writer::GgufKv; + +const TOKEN_TYPE_NORMAL: i32 = 1; +const TOKEN_TYPE_CONTROL: i32 = 3; + +pub(crate) fn push_tokenizer_metadata( + metadata: &mut Vec, + source: &Path, + config: &Value, +) -> Result<()> { + let tokenizer_path = source.join("tokenizer.json"); + if !tokenizer_path.exists() { + return Ok(()); + } + let tokenizer: Value = serde_json::from_slice( + &fs::read(&tokenizer_path).with_context(|| format!("read {}", tokenizer_path.display()))?, + ) + .with_context(|| format!("parse {}", tokenizer_path.display()))?; + let tokenizer_config = read_optional_json(&source.join("tokenizer_config.json"))?; + let vocab = read_byte_level_bpe(&tokenizer, config)?; + + metadata.push(GgufKv::string("tokenizer.ggml.model", "gpt2")); + metadata.push(GgufKv::string("tokenizer.ggml.pre", tokenizer_pre(config)?)); + metadata.push(GgufKv::array_string("tokenizer.ggml.tokens", vocab.tokens)); + metadata.push(GgufKv::array_i32( + "tokenizer.ggml.token_type", + vocab.token_types, + )); + metadata.push(GgufKv::array_f32("tokenizer.ggml.scores", vocab.scores)); + if !vocab.merges.is_empty() { + metadata.push(GgufKv::array_string("tokenizer.ggml.merges", vocab.merges)); + } + push_special_token_ids(metadata, config, &tokenizer_config, &vocab.added_tokens); + push_chat_template(metadata, source, &tokenizer_config)?; + Ok(()) +} + +pub(crate) fn ensure_native_tokenizer_metadata_supported(source: &Path) -> Result<()> { + if source.join("tokenizer.json").exists() { + return Ok(()); + } + if source.join("tokenizer.model").exists() { + anyhow::bail!( + "native tokenizer metadata does not yet support SentencePiece tokenizer.model; use external convert_hf_to_gguf.py for this checkpoint" + ); + } + anyhow::bail!( + "native tokenizer metadata requires tokenizer.json; use external convert_hf_to_gguf.py for this checkpoint" + ) +} + +struct BpeVocabMetadata { + tokens: Vec, + token_types: Vec, + scores: Vec, + merges: Vec, + added_tokens: BTreeMap, +} + +fn read_byte_level_bpe(tokenizer: &Value, _config: &Value) -> Result { + let model = tokenizer + .get("model") + .and_then(Value::as_object) + .context("tokenizer.json missing object field model")?; + ensure!( + model.get("type").and_then(Value::as_str) == Some("BPE"), + "native tokenizer metadata currently supports tokenizer.json model.type=BPE only" + ); + ensure!( + tokenizer + .get("decoder") + .and_then(|decoder| decoder.get("type")) + .and_then(Value::as_str) + == Some("ByteLevel"), + "native tokenizer metadata currently supports ByteLevel BPE decoders only" + ); + let raw_vocab = model + .get("vocab") + .and_then(Value::as_object) + .context("tokenizer.json model missing object field vocab")?; + let added_tokens = collect_added_tokens(tokenizer); + let vocab_size = tokenizer_vocab_size(raw_vocab, &added_tokens)?; + let mut tokens = vec![String::new(); vocab_size]; + let mut token_types = vec![TOKEN_TYPE_NORMAL; vocab_size]; + let mut scores = vec![0.0_f32; vocab_size]; + for (token, id) in raw_vocab { + let id = u32_value(id).with_context(|| format!("invalid vocab id for token {token:?}"))?; + let index = usize::try_from(id).context("vocab id does not fit usize")?; + ensure!( + index < tokens.len(), + "token {token:?} id {id} is outside configured vocab_size {vocab_size}" + ); + tokens[index] = token.clone(); + } + + for added in &added_tokens { + let index = usize::try_from(added.id).context("added token id does not fit usize")?; + ensure!( + index < tokens.len(), + "added token {:?} id {} is outside configured vocab_size {vocab_size}", + added.content, + added.id + ); + tokens[index] = added.content.clone(); + scores[index] = -1000.0; + if added.special { + token_types[index] = TOKEN_TYPE_CONTROL; + } + } + + let missing = tokens.iter().position(String::is_empty); + ensure!( + missing.is_none(), + "tokenizer vocab has a gap at token id {}", + missing.unwrap_or_default() + ); + Ok(BpeVocabMetadata { + tokens, + token_types, + scores, + merges: normalize_merges(model)?, + added_tokens: added_tokens + .into_iter() + .map(|token| (token.content, token.id)) + .collect(), + }) +} + +#[derive(Debug)] +struct AddedToken { + id: u32, + content: String, + special: bool, +} + +fn collect_added_tokens(tokenizer: &Value) -> Vec { + let mut seen = BTreeSet::new(); + let mut tokens = tokenizer + .get("added_tokens") + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter_map(|value| { + let id = value.get("id").and_then(u32_value)?; + let content = value.get("content")?.as_str()?.to_string(); + let special = value + .get("special") + .and_then(Value::as_bool) + .unwrap_or(false); + Some(AddedToken { + id, + content, + special, + }) + }) + .filter(|token| seen.insert(token.id)) + .collect::>(); + tokens.sort_by_key(|token| token.id); + tokens +} + +fn tokenizer_vocab_size( + raw_vocab: &serde_json::Map, + added_tokens: &[AddedToken], +) -> Result { + let mut max_id = raw_vocab + .values() + .map(u32_value) + .chain(added_tokens.iter().map(|token| Some(token.id))) + .collect::>>() + .context("tokenizer vocab contains a non-u32 id")? + .into_iter() + .max() + .unwrap_or(0); + max_id = max_id + .checked_add(1) + .context("tokenizer vocab size overflow")?; + usize::try_from(max_id).context("tokenizer vocab size does not fit usize") +} + +fn normalize_merges(model: &serde_json::Map) -> Result> { + let Some(merges) = model.get("merges").and_then(Value::as_array) else { + return Ok(Vec::new()); + }; + merges + .iter() + .map(|merge| { + if let Some(value) = merge.as_str() { + return Ok(value.to_string()); + } + let pair = merge + .as_array() + .context("tokenizer merge entries must be strings or 2-item arrays")?; + ensure!(pair.len() == 2, "tokenizer merge pair must contain 2 items"); + let left = pair[0] + .as_str() + .context("tokenizer merge left item must be a string")?; + let right = pair[1] + .as_str() + .context("tokenizer merge right item must be a string")?; + Ok(format!( + "{} {}", + encode_merge_spaces(left), + encode_merge_spaces(right) + )) + }) + .collect() +} + +fn encode_merge_spaces(value: &str) -> String { + value + .chars() + .map(|ch| { + if ch == ' ' { + char::from_u32(u32::from(ch) + 256).unwrap_or(ch) + } else { + ch + } + }) + .collect() +} + +fn tokenizer_pre(config: &Value) -> Result<&'static str> { + let model_type = config + .get("model_type") + .and_then(Value::as_str) + .unwrap_or_default(); + if model_type.starts_with("glm") { + return Ok("glm4"); + } + if model_type.starts_with("qwen2") || model_type.starts_with("qwen3") { + return Ok("qwen2"); + } + if matches!(model_type, "llama" | "mistral") { + return Ok("llama-bpe"); + } + anyhow::bail!( + "native tokenizer metadata needs an explicit tokenizer.ggml.pre mapping for model_type={model_type:?}" + ) +} + +fn push_special_token_ids( + metadata: &mut Vec, + config: &Value, + tokenizer_config: &Value, + added_tokens: &BTreeMap, +) { + let model_type = config + .get("model_type") + .and_then(Value::as_str) + .unwrap_or_default(); + if model_type.starts_with("glm") { + push_added_token_id( + metadata, + "tokenizer.ggml.bos_token_id", + added_tokens, + "[gMASK]", + ); + push_added_token_id( + metadata, + "tokenizer.ggml.eot_token_id", + added_tokens, + "<|user|>", + ); + push_added_token_id( + metadata, + "tokenizer.ggml.eom_token_id", + added_tokens, + "<|observation|>", + ); + push_added_token_id( + metadata, + "tokenizer.ggml.unknown_token_id", + added_tokens, + "<|endoftext|>", + ); + } + for (key, config_key) in [ + ("tokenizer.ggml.eos_token_id", "eos_token"), + ("tokenizer.ggml.padding_token_id", "pad_token"), + ("tokenizer.ggml.mask_token_id", "mask_token"), + ] { + if let Some(content) = tokenizer_config.get(config_key).and_then(token_content) + && let Some(id) = added_tokens.get(content) + { + metadata.push(GgufKv::u32(key, *id)); + } + } + push_tokenizer_bool(metadata, tokenizer_config, "add_bos_token"); + push_tokenizer_bool(metadata, tokenizer_config, "add_eos_token"); +} + +fn push_added_token_id( + metadata: &mut Vec, + key: &str, + added_tokens: &BTreeMap, + content: &str, +) { + if let Some(id) = added_tokens.get(content) { + metadata.push(GgufKv::u32(key, *id)); + } +} + +fn push_chat_template( + metadata: &mut Vec, + source: &Path, + tokenizer_config: &Value, +) -> Result<()> { + if let Some(template) = tokenizer_config + .get("chat_template") + .and_then(Value::as_str) + { + metadata.push(GgufKv::string("tokenizer.chat_template", template)); + return Ok(()); + } + let template_path = source.join("chat_template.jinja"); + if template_path.exists() { + metadata.push(GgufKv::string( + "tokenizer.chat_template", + &fs::read_to_string(&template_path) + .with_context(|| format!("read {}", template_path.display()))?, + )); + } + Ok(()) +} + +fn token_content(value: &Value) -> Option<&str> { + value + .as_str() + .or_else(|| value.get("content").and_then(Value::as_str)) +} + +fn push_tokenizer_bool(metadata: &mut Vec, tokenizer_config: &Value, config_key: &str) { + if let Some(value) = tokenizer_config.get(config_key).and_then(Value::as_bool) { + metadata.push(GgufKv::bool(&format!("tokenizer.ggml.{config_key}"), value)); + } +} + +fn read_optional_json(path: &Path) -> Result { + if !path.exists() { + return Ok(Value::Null); + } + serde_json::from_slice(&fs::read(path).with_context(|| format!("read {}", path.display()))?) + .with_context(|| format!("parse {}", path.display())) +} + +fn u32_value(value: &Value) -> Option { + value.as_u64().and_then(|value| u32::try_from(value).ok()) +} + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + + use super::*; + + #[test] + fn builds_glm_byte_level_bpe_tokenizer_metadata() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + fs::write( + root.join("tokenizer.json"), + r#"{ + "model": { + "type": "BPE", + "vocab": {"a": 0, "b": 1, "[gMASK]": 2, "<|endoftext|>": 3, "<|user|>": 4, "<|observation|>": 5}, + "merges": [["a", "b"]] + }, + "decoder": {"type": "ByteLevel"}, + "added_tokens": [ + {"id": 2, "content": "[gMASK]", "special": true}, + {"id": 3, "content": "<|endoftext|>", "special": true}, + {"id": 4, "content": "<|user|>", "special": true}, + {"id": 5, "content": "<|observation|>", "special": true} + ] + }"#, + ) + .unwrap(); + fs::write( + root.join("tokenizer_config.json"), + r#"{"eos_token": "<|endoftext|>", "pad_token": "<|endoftext|>"}"#, + ) + .unwrap(); + let config: Value = + serde_json::from_str(r#"{"model_type":"glm4_moe_lite","vocab_size":6}"#).unwrap(); + let mut metadata = Vec::new(); + + push_tokenizer_metadata(&mut metadata, &root, &config).unwrap(); + let text = format!("{metadata:?}"); + + assert!(text.contains("tokenizer.ggml.tokens")); + assert!(text.contains("tokenizer.ggml.pre")); + assert!(text.contains("glm4")); + assert!(text.contains("tokenizer.ggml.bos_token_id")); + assert!(text.contains("tokenizer.ggml.eot_token_id")); + assert!(text.contains("tokenizer.ggml.eom_token_id")); + assert!(text.contains("tokenizer.ggml.merges")); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn builds_qwen_byte_level_bpe_tokenizer_metadata() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + fs::write( + root.join("tokenizer.json"), + r#"{ + "model": { + "type": "BPE", + "vocab": {"a": 0, "b": 1, "<|endoftext|>": 2, "<|im_end|>": 3}, + "merges": ["a b"] + }, + "decoder": {"type": "ByteLevel"}, + "added_tokens": [ + {"id": 2, "content": "<|endoftext|>", "special": true}, + {"id": 3, "content": "<|im_end|>", "special": true} + ] + }"#, + ) + .unwrap(); + fs::write( + root.join("tokenizer_config.json"), + r#"{"eos_token": "<|im_end|>", "pad_token": "<|endoftext|>", "add_bos_token": false}"#, + ) + .unwrap(); + let config: Value = + serde_json::from_str(r#"{"model_type":"qwen3","vocab_size":4}"#).unwrap(); + let mut metadata = Vec::new(); + + push_tokenizer_metadata(&mut metadata, &root, &config).unwrap(); + let text = format!("{metadata:?}"); + + assert!(text.contains("tokenizer.ggml.pre")); + assert!(text.contains("qwen2")); + assert!(text.contains("tokenizer.ggml.eos_token_id")); + assert!(text.contains("tokenizer.ggml.padding_token_id")); + assert!(text.contains("tokenizer.ggml.add_bos_token")); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn ignores_trailing_embedding_vocab_padding() { + let tokenizer: Value = serde_json::from_str( + r#"{ + "model": { + "type": "BPE", + "vocab": {"a": 0, "b": 1}, + "merges": ["a b"] + }, + "decoder": {"type": "ByteLevel"}, + "added_tokens": [ + {"id": 2, "content": "<|endoftext|>", "special": true} + ] + }"#, + ) + .unwrap(); + let config: Value = + serde_json::from_str(r#"{"model_type":"qwen2","vocab_size":8}"#).unwrap(); + + let metadata = read_byte_level_bpe(&tokenizer, &config).unwrap(); + + assert_eq!(metadata.tokens.len(), 3); + assert_eq!(metadata.tokens[2], "<|endoftext|>"); + } + + #[test] + fn builds_llama_byte_level_bpe_tokenizer_metadata() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + fs::write( + root.join("tokenizer.json"), + r#"{ + "model": { + "type": "BPE", + "vocab": {"a": 0, "b": 1, "<|end_of_text|>": 2}, + "merges": ["a b"] + }, + "decoder": {"type": "ByteLevel"}, + "added_tokens": [ + {"id": 2, "content": "<|end_of_text|>", "special": true} + ] + }"#, + ) + .unwrap(); + fs::write( + root.join("tokenizer_config.json"), + r#"{"eos_token": "<|end_of_text|>", "add_bos_token": true}"#, + ) + .unwrap(); + let config: Value = + serde_json::from_str(r#"{"model_type":"llama","vocab_size":3}"#).unwrap(); + let mut metadata = Vec::new(); + + push_tokenizer_metadata(&mut metadata, &root, &config).unwrap(); + let text = format!("{metadata:?}"); + + assert!(text.contains("tokenizer.ggml.pre")); + assert!(text.contains("llama-bpe")); + assert!(text.contains("tokenizer.ggml.eos_token_id")); + assert!(text.contains("tokenizer.ggml.add_bos_token")); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn rejects_missing_tokenizer_json() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + let config: Value = serde_json::from_str(r#"{"model_type":"qwen3"}"#).unwrap(); + let mut metadata = Vec::new(); + + push_tokenizer_metadata(&mut metadata, &root, &config).unwrap(); + let error = ensure_native_tokenizer_metadata_supported(&root) + .unwrap_err() + .to_string(); + + assert!(error.contains("requires tokenizer.json")); + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn rejects_sentencepiece_tokenizer_model_without_tokenizer_json() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + fs::write(root.join("tokenizer.model"), b"not-a-real-spm").unwrap(); + let config: Value = serde_json::from_str(r#"{"model_type":"llama"}"#).unwrap(); + let mut metadata = Vec::new(); + + push_tokenizer_metadata(&mut metadata, &root, &config).unwrap(); + let error = ensure_native_tokenizer_metadata_supported(&root) + .unwrap_err() + .to_string(); + + assert!(error.contains("SentencePiece tokenizer.model")); + fs::remove_dir_all(root).unwrap(); + } + + fn unique_temp_dir() -> PathBuf { + static NEXT_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let id = NEXT_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + std::env::temp_dir().join(format!("skippy-tokenizer-metadata-{nanos}-{id}")) + } +} diff --git a/crates/skippy-quantize/src/tool_paths.rs b/crates/skippy-quantize/src/tool_paths.rs new file mode 100644 index 0000000000..5f04c69fbf --- /dev/null +++ b/crates/skippy-quantize/src/tool_paths.rs @@ -0,0 +1,65 @@ +use std::path::{Path, PathBuf}; + +const LLAMA_CLI_ENV: &str = "SKIPPY_QUANTIZE_LLAMA_CLI"; +const LLAMA_CLI_CANDIDATES: &[&str] = &[ + ".deps/llama.cpp/build-cli/bin/llama-cli", + ".deps/llama.cpp/build-cli/bin/llama", + ".deps/llama.cpp/build-cli/bin/llama-simple", + ".deps/llama.cpp/build/bin/llama-cli", + ".deps/llama.cpp/build/bin/llama", + ".deps/llama.cpp/build/bin/llama-simple", + ".deps/llama.cpp/build/bin/Release/llama-cli", + ".deps/llama.cpp/build/bin/Release/llama", + ".deps/llama.cpp/build/bin/Release/llama-simple", + "../../.deps/llama.cpp/build-cli/bin/llama-cli", + "../../.deps/llama.cpp/build-cli/bin/llama", + "../../.deps/llama.cpp/build-cli/bin/llama-simple", + "../../.deps/llama.cpp/build/bin/llama-cli", + "../../.deps/llama.cpp/build/bin/llama", + "../../.deps/llama.cpp/build/bin/llama-simple", + "../../.deps/llama.cpp/build/bin/Release/llama-cli", + "../../.deps/llama.cpp/build/bin/Release/llama", + "../../.deps/llama.cpp/build/bin/Release/llama-simple", +]; + +pub(crate) fn resolve_llama_cli(explicit: Option<&Path>) -> Option { + resolve_tool(explicit, LLAMA_CLI_ENV, LLAMA_CLI_CANDIDATES) +} + +fn resolve_tool( + explicit: Option<&Path>, + env_var: &str, + relative_candidates: &[&str], +) -> Option { + if let Some(explicit) = explicit { + return Some(explicit.to_path_buf()); + } + if let Some(from_env) = std::env::var_os(env_var).map(PathBuf::from) + && from_env.is_file() + { + return Some(from_env); + } + for root in candidate_roots() { + for relative in relative_candidates { + let candidate = root.join(relative); + if candidate.is_file() { + return Some(candidate); + } + } + } + None +} + +fn candidate_roots() -> Vec { + let mut roots = Vec::new(); + if let Ok(current) = std::env::current_dir() { + roots.push(current); + } + let manifest_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")); + roots.push(manifest_dir.clone()); + if let Some(repo_root) = manifest_dir.parent().and_then(Path::parent) { + roots.push(repo_root.to_path_buf()); + } + roots.dedup(); + roots +} diff --git a/crates/skippy-quantize/src/type_catalog.rs b/crates/skippy-quantize/src/type_catalog.rs new file mode 100644 index 0000000000..d06633ae48 --- /dev/null +++ b/crates/skippy-quantize/src/type_catalog.rs @@ -0,0 +1,70 @@ +use anyhow::Result; +use clap::Parser; +use serde::Serialize; + +use crate::output::{print_info, print_json_pretty, print_success}; +use crate::types::{QuantType, TensorType}; + +#[derive(Debug, Parser)] +pub(crate) struct TypeCatalogArgs { + #[arg(long)] + json: bool, +} + +#[derive(Debug, Serialize)] +struct QuantCatalog { + whole_model_quant_modes: Vec<&'static str>, +} + +#[derive(Debug, Serialize)] +struct TensorTypeCatalog { + raw_tensor_types: Vec<&'static str>, +} + +pub(crate) fn list_quants(args: TypeCatalogArgs) -> Result<()> { + let names = QuantType::ALL + .iter() + .map(|quant| quant.as_llama_name()) + .collect::>(); + if args.json { + print_json_pretty(&QuantCatalog { + whole_model_quant_modes: names, + })?; + } else { + print_success("Whole-model quant modes"); + for name in names { + println!(" • {name}"); + } + print_info("Use --tensor-type-file for custom tensor recipes"); + } + Ok(()) +} + +pub(crate) fn list_tensor_types(args: TypeCatalogArgs) -> Result<()> { + let names = TensorType::ALL + .iter() + .map(|tensor_type| tensor_type.as_ggml_name()) + .collect::>(); + if args.json { + print_json_pretty(&TensorTypeCatalog { + raw_tensor_types: names, + })?; + } else { + print_success("Raw tensor override types"); + for name in names { + println!(" • {name}"); + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn catalogs_are_not_empty() { + assert!(!QuantType::ALL.is_empty()); + assert!(!TensorType::ALL.is_empty()); + } +} diff --git a/crates/skippy-quantize/src/types.rs b/crates/skippy-quantize/src/types.rs new file mode 100644 index 0000000000..80f7f42c81 --- /dev/null +++ b/crates/skippy-quantize/src/types.rs @@ -0,0 +1,757 @@ +use std::str::FromStr; + +use clap::ValueEnum; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, ValueEnum)] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub enum ConvertOutputType { + F32, + F16, + Bf16, + Q8_0, + TQ1_0, + TQ2_0, + Auto, +} + +impl ConvertOutputType { + pub fn as_arg(self) -> &'static str { + match self { + Self::F32 => "f32", + Self::F16 => "f16", + Self::Bf16 => "bf16", + Self::Q8_0 => "q8_0", + Self::TQ1_0 => "tq1_0", + Self::TQ2_0 => "tq2_0", + Self::Auto => "auto", + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub enum JobKind { + ConvertHf, + QuantizeGguf, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub enum QuantType { + Q1_0, + Q2K, + Q2KS, + Q3K, + Q3KS, + Q3KM, + Q3KL, + Q4_0, + Q4_1, + Q4K, + Q4KS, + Q4KM, + Q5_0, + Q5_1, + Q5K, + Q5KS, + Q5KM, + Q6K, + Q8_0, + IQ1S, + IQ1M, + IQ2XXS, + IQ2XS, + IQ2S, + IQ2M, + IQ3XXS, + IQ3XS, + IQ3S, + IQ3M, + IQ4NL, + IQ4XS, + TQ1_0, + TQ2_0, + Mxfp4Moe, + F16, + Bf16, + F32, + Copy, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct QuantSpec { + base_quant: QuantType, +} + +impl FromStr for QuantSpec { + type Err = String; + + fn from_str(raw: &str) -> std::result::Result { + let base_quant = raw.parse::()?; + Ok(Self { base_quant }) + } +} + +impl QuantSpec { + pub fn base_quant(&self) -> QuantType { + self.base_quant + } + + pub fn output_name(&self) -> &'static str { + self.base_quant.as_llama_name() + } + + pub fn validate_recipe_requirements( + &self, + _has_tensor_type_file: bool, + ) -> std::result::Result<(), String> { + Ok(()) + } +} + +impl From for QuantSpec { + fn from(base_quant: QuantType) -> Self { + Self { base_quant } + } +} + +impl FromStr for QuantType { + type Err = String; + + fn from_str(raw: &str) -> std::result::Result { + if let Ok(ftype) = raw.parse::() { + return Self::from_llama_ftype_id(ftype) + .ok_or_else(|| format!("unsupported quant ftype id {ftype}")); + } + let normalized = normalize_type_name(raw); + let quant = match normalized.as_str() { + "Q10" => Self::Q1_0, + "Q2K" => Self::Q2K, + "Q2KS" => Self::Q2KS, + "Q3K" => Self::Q3K, + "Q3KS" => Self::Q3KS, + "Q3KM" => Self::Q3KM, + "Q3KL" => Self::Q3KL, + "Q40" => Self::Q4_0, + "Q41" => Self::Q4_1, + "Q4K" => Self::Q4K, + "Q4KS" => Self::Q4KS, + "Q4KM" => Self::Q4KM, + "Q50" => Self::Q5_0, + "Q51" => Self::Q5_1, + "Q5K" => Self::Q5K, + "Q5KS" => Self::Q5KS, + "Q5KM" => Self::Q5KM, + "Q6K" => Self::Q6K, + "Q80" => Self::Q8_0, + "IQ1S" => Self::IQ1S, + "IQ1M" => Self::IQ1M, + "IQ2XXS" => Self::IQ2XXS, + "IQ2XS" => Self::IQ2XS, + "IQ2S" => Self::IQ2S, + "IQ2M" => Self::IQ2M, + "IQ3XXS" => Self::IQ3XXS, + "IQ3XS" => Self::IQ3XS, + "IQ3S" => Self::IQ3S, + "IQ3M" => Self::IQ3M, + "IQ4NL" => Self::IQ4NL, + "IQ4XS" => Self::IQ4XS, + "TQ10" => Self::TQ1_0, + "TQ20" => Self::TQ2_0, + "MXFP4MOE" => Self::Mxfp4Moe, + "F16" => Self::F16, + "BF16" => Self::Bf16, + "F32" => Self::F32, + "COPY" => Self::Copy, + _ => return Err(unsupported_quant_type_error(raw, &normalized)), + }; + Ok(quant) + } +} + +impl QuantType { + pub const ALL: &'static [Self] = &[ + Self::Q1_0, + Self::Q2K, + Self::Q2KS, + Self::Q3K, + Self::Q3KS, + Self::Q3KM, + Self::Q3KL, + Self::Q4_0, + Self::Q4_1, + Self::Q4K, + Self::Q4KS, + Self::Q4KM, + Self::Q5_0, + Self::Q5_1, + Self::Q5K, + Self::Q5KS, + Self::Q5KM, + Self::Q6K, + Self::Q8_0, + Self::IQ1S, + Self::IQ1M, + Self::IQ2XXS, + Self::IQ2XS, + Self::IQ2S, + Self::IQ2M, + Self::IQ3XXS, + Self::IQ3XS, + Self::IQ3S, + Self::IQ3M, + Self::IQ4NL, + Self::IQ4XS, + Self::TQ1_0, + Self::TQ2_0, + Self::Mxfp4Moe, + Self::F16, + Self::Bf16, + Self::F32, + Self::Copy, + ]; + + pub fn as_llama_name(self) -> &'static str { + match self { + Self::Q1_0 => "Q1_0", + Self::Q2K => "Q2_K", + Self::Q2KS => "Q2_K_S", + Self::Q3K => "Q3_K", + Self::Q3KS => "Q3_K_S", + Self::Q3KM => "Q3_K_M", + Self::Q3KL => "Q3_K_L", + Self::Q4_0 => "Q4_0", + Self::Q4_1 => "Q4_1", + Self::Q4K => "Q4_K", + Self::Q4KS => "Q4_K_S", + Self::Q4KM => "Q4_K_M", + Self::Q5_0 => "Q5_0", + Self::Q5_1 => "Q5_1", + Self::Q5K => "Q5_K", + Self::Q5KS => "Q5_K_S", + Self::Q5KM => "Q5_K_M", + Self::Q6K => "Q6_K", + Self::Q8_0 => "Q8_0", + Self::IQ1S => "IQ1_S", + Self::IQ1M => "IQ1_M", + Self::IQ2XXS => "IQ2_XXS", + Self::IQ2XS => "IQ2_XS", + Self::IQ2S => "IQ2_S", + Self::IQ2M => "IQ2_M", + Self::IQ3XXS => "IQ3_XXS", + Self::IQ3XS => "IQ3_XS", + Self::IQ3S => "IQ3_S", + Self::IQ3M => "IQ3_M", + Self::IQ4NL => "IQ4_NL", + Self::IQ4XS => "IQ4_XS", + Self::TQ1_0 => "TQ1_0", + Self::TQ2_0 => "TQ2_0", + Self::Mxfp4Moe => "MXFP4_MOE", + Self::F16 => "F16", + Self::Bf16 => "BF16", + Self::F32 => "F32", + Self::Copy => "COPY", + } + } + + pub fn from_llama_ftype_id(ftype: i32) -> Option { + match ftype { + 0 => Some(Self::F32), + 1 => Some(Self::F16), + 2 => Some(Self::Q4_0), + 3 => Some(Self::Q4_1), + 7 => Some(Self::Q8_0), + 8 => Some(Self::Q5_0), + 9 => Some(Self::Q5_1), + 10 => Some(Self::Q2K), + 11 => Some(Self::Q3KS), + 12 => Some(Self::Q3K), + 13 => Some(Self::Q3KL), + 14 => Some(Self::Q4KS), + 15 => Some(Self::Q4K), + 16 => Some(Self::Q5KS), + 17 => Some(Self::Q5K), + 18 => Some(Self::Q6K), + 19 => Some(Self::IQ2XXS), + 20 => Some(Self::IQ2XS), + 21 => Some(Self::Q2KS), + 22 => Some(Self::IQ3XS), + 23 => Some(Self::IQ3XXS), + 24 => Some(Self::IQ1S), + 25 => Some(Self::IQ4NL), + 26 => Some(Self::IQ3S), + 27 => Some(Self::IQ3M), + 28 => Some(Self::IQ2S), + 29 => Some(Self::IQ2M), + 30 => Some(Self::IQ4XS), + 31 => Some(Self::IQ1M), + 32 => Some(Self::Bf16), + 36 => Some(Self::TQ1_0), + 37 => Some(Self::TQ2_0), + 38 => Some(Self::Mxfp4Moe), + 40 => Some(Self::Q1_0), + _ => None, + } + } + + pub fn as_llama_file_type(self) -> llama_quant_ffi::LlamaFileType { + match self { + Self::Q1_0 => llama_quant_ffi::LlamaFileType::MostlyQ1_0, + Self::Q2K => llama_quant_ffi::LlamaFileType::MostlyQ2K, + Self::Q2KS => llama_quant_ffi::LlamaFileType::MostlyQ2KS, + Self::Q3K | Self::Q3KM => llama_quant_ffi::LlamaFileType::MostlyQ3KM, + Self::Q3KS => llama_quant_ffi::LlamaFileType::MostlyQ3KS, + Self::Q3KL => llama_quant_ffi::LlamaFileType::MostlyQ3KL, + Self::Q4_0 => llama_quant_ffi::LlamaFileType::MostlyQ4_0, + Self::Q4_1 => llama_quant_ffi::LlamaFileType::MostlyQ4_1, + Self::Q4K | Self::Q4KM => llama_quant_ffi::LlamaFileType::MostlyQ4KM, + Self::Q4KS => llama_quant_ffi::LlamaFileType::MostlyQ4KS, + Self::Q5_0 => llama_quant_ffi::LlamaFileType::MostlyQ5_0, + Self::Q5_1 => llama_quant_ffi::LlamaFileType::MostlyQ5_1, + Self::Q5K | Self::Q5KM => llama_quant_ffi::LlamaFileType::MostlyQ5KM, + Self::Q5KS => llama_quant_ffi::LlamaFileType::MostlyQ5KS, + Self::Q6K => llama_quant_ffi::LlamaFileType::MostlyQ6K, + Self::Q8_0 => llama_quant_ffi::LlamaFileType::MostlyQ8_0, + Self::IQ1S => llama_quant_ffi::LlamaFileType::MostlyIQ1S, + Self::IQ1M => llama_quant_ffi::LlamaFileType::MostlyIQ1M, + Self::IQ2XXS => llama_quant_ffi::LlamaFileType::MostlyIQ2XXS, + Self::IQ2XS => llama_quant_ffi::LlamaFileType::MostlyIQ2XS, + Self::IQ2S => llama_quant_ffi::LlamaFileType::MostlyIQ2S, + Self::IQ2M => llama_quant_ffi::LlamaFileType::MostlyIQ2M, + Self::IQ3XXS => llama_quant_ffi::LlamaFileType::MostlyIQ3XXS, + Self::IQ3XS => llama_quant_ffi::LlamaFileType::MostlyIQ3XS, + Self::IQ3S => llama_quant_ffi::LlamaFileType::MostlyIQ3S, + Self::IQ3M => llama_quant_ffi::LlamaFileType::MostlyIQ3M, + Self::IQ4NL => llama_quant_ffi::LlamaFileType::MostlyIQ4NL, + Self::IQ4XS => llama_quant_ffi::LlamaFileType::MostlyIQ4XS, + Self::TQ1_0 => llama_quant_ffi::LlamaFileType::MostlyTQ1_0, + Self::TQ2_0 => llama_quant_ffi::LlamaFileType::MostlyTQ2_0, + Self::Mxfp4Moe => llama_quant_ffi::LlamaFileType::MostlyMxfp4Moe, + Self::F16 => llama_quant_ffi::LlamaFileType::MostlyF16, + Self::Bf16 => llama_quant_ffi::LlamaFileType::MostlyBf16, + Self::F32 | Self::Copy => llama_quant_ffi::LlamaFileType::AllF32, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub enum TensorType { + F32, + F16, + Q1_0, + Q4_0, + Q4_1, + Q5_0, + Q5_1, + Q8_0, + Q8_1, + Q2K, + Q3K, + Q4K, + Q5K, + Q6K, + Q8K, + IQ1S, + IQ1M, + IQ2XXS, + IQ2XS, + IQ2S, + IQ3XXS, + IQ3S, + IQ4NL, + IQ4XS, + I8, + I16, + I32, + I64, + F64, + TQ1_0, + TQ2_0, + Mxfp4, + Nvfp4, + Bf16, +} + +impl TensorType { + pub fn as_ggml_type(self) -> Option { + match self { + Self::F32 => Some(llama_quant_ffi::GgmlType::F32), + Self::F16 => Some(llama_quant_ffi::GgmlType::F16), + Self::Q1_0 => Some(llama_quant_ffi::GgmlType::Q1_0), + Self::Q4_0 => Some(llama_quant_ffi::GgmlType::Q4_0), + Self::Q4_1 => Some(llama_quant_ffi::GgmlType::Q4_1), + Self::Q5_0 => Some(llama_quant_ffi::GgmlType::Q5_0), + Self::Q5_1 => Some(llama_quant_ffi::GgmlType::Q5_1), + Self::Q8_0 => Some(llama_quant_ffi::GgmlType::Q8_0), + Self::Q8_1 => Some(llama_quant_ffi::GgmlType::Q8_1), + Self::Q2K => Some(llama_quant_ffi::GgmlType::Q2K), + Self::Q3K => Some(llama_quant_ffi::GgmlType::Q3K), + Self::Q4K => Some(llama_quant_ffi::GgmlType::Q4K), + Self::Q5K => Some(llama_quant_ffi::GgmlType::Q5K), + Self::Q6K => Some(llama_quant_ffi::GgmlType::Q6K), + Self::Q8K => Some(llama_quant_ffi::GgmlType::Q8K), + Self::IQ1S => Some(llama_quant_ffi::GgmlType::IQ1S), + Self::IQ1M => Some(llama_quant_ffi::GgmlType::IQ1M), + Self::IQ2XXS => Some(llama_quant_ffi::GgmlType::IQ2XXS), + Self::IQ2XS => Some(llama_quant_ffi::GgmlType::IQ2XS), + Self::IQ2S => Some(llama_quant_ffi::GgmlType::IQ2S), + Self::IQ3XXS => Some(llama_quant_ffi::GgmlType::IQ3XXS), + Self::IQ3S => Some(llama_quant_ffi::GgmlType::IQ3S), + Self::IQ4NL => Some(llama_quant_ffi::GgmlType::IQ4NL), + Self::IQ4XS => Some(llama_quant_ffi::GgmlType::IQ4XS), + Self::I8 => Some(llama_quant_ffi::GgmlType::I8), + Self::I16 => Some(llama_quant_ffi::GgmlType::I16), + Self::I32 => Some(llama_quant_ffi::GgmlType::I32), + Self::I64 => Some(llama_quant_ffi::GgmlType::I64), + Self::F64 => Some(llama_quant_ffi::GgmlType::F64), + Self::TQ1_0 => Some(llama_quant_ffi::GgmlType::TQ1_0), + Self::TQ2_0 => Some(llama_quant_ffi::GgmlType::TQ2_0), + Self::Mxfp4 => Some(llama_quant_ffi::GgmlType::Mxfp4), + Self::Nvfp4 => Some(llama_quant_ffi::GgmlType::Nvfp4), + Self::Bf16 => Some(llama_quant_ffi::GgmlType::Bf16), + } + } + + pub const ALL: &'static [Self] = &[ + Self::F32, + Self::F16, + Self::Q1_0, + Self::Q4_0, + Self::Q4_1, + Self::Q5_0, + Self::Q5_1, + Self::Q8_0, + Self::Q8_1, + Self::Q2K, + Self::Q3K, + Self::Q4K, + Self::Q5K, + Self::Q6K, + Self::Q8K, + Self::IQ1S, + Self::IQ1M, + Self::IQ2XXS, + Self::IQ2XS, + Self::IQ2S, + Self::IQ3XXS, + Self::IQ3S, + Self::IQ4NL, + Self::IQ4XS, + Self::I8, + Self::I16, + Self::I32, + Self::I64, + Self::F64, + Self::TQ1_0, + Self::TQ2_0, + Self::Mxfp4, + Self::Nvfp4, + Self::Bf16, + ]; + + pub fn as_ggml_name(self) -> &'static str { + match self { + Self::F32 => "F32", + Self::F16 => "F16", + Self::Q1_0 => "Q1_0", + Self::Q4_0 => "Q4_0", + Self::Q4_1 => "Q4_1", + Self::Q5_0 => "Q5_0", + Self::Q5_1 => "Q5_1", + Self::Q8_0 => "Q8_0", + Self::Q8_1 => "Q8_1", + Self::Q2K => "Q2_K", + Self::Q3K => "Q3_K", + Self::Q4K => "Q4_K", + Self::Q5K => "Q5_K", + Self::Q6K => "Q6_K", + Self::Q8K => "Q8_K", + Self::IQ1S => "IQ1_S", + Self::IQ1M => "IQ1_M", + Self::IQ2XXS => "IQ2_XXS", + Self::IQ2XS => "IQ2_XS", + Self::IQ2S => "IQ2_S", + Self::IQ3XXS => "IQ3_XXS", + Self::IQ3S => "IQ3_S", + Self::IQ4NL => "IQ4_NL", + Self::IQ4XS => "IQ4_XS", + Self::I8 => "I8", + Self::I16 => "I16", + Self::I32 => "I32", + Self::I64 => "I64", + Self::F64 => "F64", + Self::TQ1_0 => "TQ1_0", + Self::TQ2_0 => "TQ2_0", + Self::Mxfp4 => "MXFP4", + Self::Nvfp4 => "NVFP4", + Self::Bf16 => "BF16", + } + } + + pub fn parse(raw: &str) -> Option { + let normalized = normalize_type_name(raw); + match normalized.as_str() { + "F32" => Some(Self::F32), + "F16" => Some(Self::F16), + "Q10" => Some(Self::Q1_0), + "Q40" => Some(Self::Q4_0), + "Q41" => Some(Self::Q4_1), + "Q50" => Some(Self::Q5_0), + "Q51" => Some(Self::Q5_1), + "Q80" => Some(Self::Q8_0), + "Q81" => Some(Self::Q8_1), + "Q2K" => Some(Self::Q2K), + "Q3K" => Some(Self::Q3K), + "Q4K" => Some(Self::Q4K), + "Q5K" => Some(Self::Q5K), + "Q6K" => Some(Self::Q6K), + "Q8K" => Some(Self::Q8K), + "IQ1S" => Some(Self::IQ1S), + "IQ1M" => Some(Self::IQ1M), + "IQ2XXS" => Some(Self::IQ2XXS), + "IQ2XS" => Some(Self::IQ2XS), + "IQ2S" => Some(Self::IQ2S), + "IQ3XXS" => Some(Self::IQ3XXS), + "IQ3S" => Some(Self::IQ3S), + "IQ4NL" => Some(Self::IQ4NL), + "IQ4XS" => Some(Self::IQ4XS), + "I8" => Some(Self::I8), + "I16" => Some(Self::I16), + "I32" => Some(Self::I32), + "I64" => Some(Self::I64), + "F64" => Some(Self::F64), + "TQ10" => Some(Self::TQ1_0), + "TQ20" => Some(Self::TQ2_0), + "MXFP4" => Some(Self::Mxfp4), + "NVFP4" => Some(Self::Nvfp4), + "BF16" => Some(Self::Bf16), + _ => None, + } + } +} + +fn normalize_type_name(raw: &str) -> String { + raw.chars() + .filter(|ch| !matches!(*ch, '_' | '-')) + .flat_map(char::to_uppercase) + .collect() +} + +fn unsupported_quant_type_error(raw: &str, normalized: &str) -> String { + if let Some(base) = normalized.strip_prefix("UD") + && let Some(base_quant) = base_quant_from_profile_suffix(base) + { + return format!( + "unsupported quant type {raw:?}: UD-* labels are custom tensor-type recipes, \ + not upstream llama-quantize whole-model modes; use base quant {base_quant:?} \ + with --tensor-type-file for the dynamic recipe" + ); + } + if normalized == "Q4KXL" { + return "unsupported quant type \"Q4_K_XL\": Q4_K_XL is a custom high-quality recipe, \ + not an upstream llama-quantize whole-model mode; use base quant \"Q4_K_M\" \ + with --tensor-type-file for the XL recipe" + .to_string(); + } + if normalized.ends_with("MTPQ8") { + return format!( + "unsupported quant type {raw:?}: MTP-Q8 is a custom artifact profile, \ + not a whole-model quant mode; pass the base quant with --quant and \ + the tensor policy with --tensor-type-file" + ); + } + format!("unsupported quant type {raw:?}") +} + +fn base_quant_from_profile_suffix(normalized_suffix: &str) -> Option<&'static str> { + match normalized_suffix { + "Q2K" => Some("Q2_K"), + "Q2KS" => Some("Q2_K_S"), + "Q3K" => Some("Q3_K"), + "Q3KS" => Some("Q3_K_S"), + "Q3KM" => Some("Q3_K_M"), + "Q3KL" => Some("Q3_K_L"), + "Q4K" => Some("Q4_K"), + "Q4KS" => Some("Q4_K_S"), + "Q4KM" => Some("Q4_K_M"), + "Q5K" => Some("Q5_K"), + "Q5KS" => Some("Q5_K_S"), + "Q5KM" => Some("Q5_K_M"), + "Q6K" => Some("Q6_K"), + other => QuantType::from_str(other) + .ok() + .map(QuantType::as_llama_name), + } +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeSet; + use std::fs; + use std::path::PathBuf; + + use super::*; + + #[test] + fn accepts_raw_tensor_types_but_not_ftype_mixtures() { + assert!(TensorType::parse("Q1_0").is_some()); + assert!(TensorType::parse("Q3_K").is_some()); + assert!(TensorType::parse("q4_K").is_some()); + assert!(TensorType::parse("Q6_K").is_some()); + assert!(TensorType::parse("Q8_0").is_some()); + assert!(TensorType::parse("MXFP4").is_some()); + assert!(TensorType::parse("NVFP4").is_some()); + assert!(TensorType::parse("I8").is_some()); + assert!(TensorType::parse("F64").is_some()); + assert!(TensorType::parse("IQ2_M").is_none()); + assert!(TensorType::parse("IQ3_XS").is_none()); + assert!(TensorType::parse("IQ3_M").is_none()); + assert!(TensorType::parse("Q3_K_S").is_none()); + assert!(TensorType::parse("Q4_K_M").is_none()); + } + + #[test] + fn quant_names_match_llama_cli() { + assert_eq!(QuantType::Q2K.as_llama_name(), "Q2_K"); + assert_eq!(QuantType::Q3KS.as_llama_name(), "Q3_K_S"); + assert_eq!(QuantType::Mxfp4Moe.as_llama_name(), "MXFP4_MOE"); + } + + #[test] + fn parses_llama_quant_names() { + assert_eq!("Q2_K".parse::().unwrap(), QuantType::Q2K); + assert_eq!("q2-k".parse::().unwrap(), QuantType::Q2K); + assert_eq!("q2k".parse::().unwrap(), QuantType::Q2K); + assert_eq!( + "MXFP4_MOE".parse::().unwrap(), + QuantType::Mxfp4Moe + ); + assert!("NVFP4".parse::().is_err()); + } + + #[test] + fn parses_quant_specs_from_llama_quant_names() { + assert!("UD-Q3_K_S".parse::().is_err()); + assert!("Q4_K_XL".parse::().is_err()); + assert!("Q2_K-MTP-Q8".parse::().is_err()); + let regular = "Q4_K_M".parse::().unwrap(); + assert_eq!(regular.base_quant(), QuantType::Q4KM); + assert_eq!(regular.output_name(), "Q4_K_M"); + assert!(regular.validate_recipe_requirements(false).is_ok()); + } + + #[test] + fn parses_llama_numeric_ftype_ids() { + assert_eq!("0".parse::().unwrap(), QuantType::F32); + assert_eq!("12".parse::().unwrap(), QuantType::Q3K); + assert_eq!("15".parse::().unwrap(), QuantType::Q4K); + assert_eq!("17".parse::().unwrap(), QuantType::Q5K); + assert_eq!("40".parse::().unwrap(), QuantType::Q1_0); + assert!("999".parse::().is_err()); + } + + #[test] + fn parses_every_current_llama_quant_cli_name() { + for name in pinned_llama_quant_option_names() { + assert!( + name.parse::().is_ok(), + "{name} should parse as a llama-quantize mode" + ); + } + } + + #[test] + fn rejects_quant_modes_not_in_current_llama_cli() { + let xl_error = "Q4_K_XL".parse::().unwrap_err(); + assert!(xl_error.contains("custom high-quality recipe")); + assert!(xl_error.contains("Q4_K_M")); + + let ud_error = "UD-Q3_K_S".parse::().unwrap_err(); + assert!(ud_error.contains("custom tensor-type recipes")); + assert!(ud_error.contains("Q3_K_S")); + } + + #[test] + fn parses_current_llama_ftype_names() { + let names = ["MXFP4_MOE", "Q1_0"]; + for name in names { + assert!( + name.parse::().is_ok(), + "{name} should parse as a llama ftype-backed mode" + ); + } + } + + #[test] + fn public_quant_catalog_is_parseable() { + assert!(QuantType::ALL.contains(&QuantType::Copy)); + assert!(QuantType::ALL.contains(&QuantType::Mxfp4Moe)); + for quant in QuantType::ALL { + assert_eq!(quant.as_llama_name().parse::().unwrap(), *quant); + } + } + + #[test] + fn public_quant_catalog_covers_pinned_llama_quantize_table() { + let pinned = pinned_llama_quant_option_names() + .into_iter() + .collect::>(); + let local = QuantType::ALL + .iter() + .map(|quant| quant.as_llama_name().to_string()) + .collect::>(); + + let missing = pinned.difference(&local).collect::>(); + assert!( + missing.is_empty(), + "local quant catalog is missing pinned llama-quantize modes: {missing:?}" + ); + + let extra = local.difference(&pinned).collect::>(); + assert!( + extra.is_empty(), + "local quant catalog has modes missing from pinned llama-quantize: {extra:?}" + ); + } + + #[test] + fn public_tensor_catalog_is_parseable() { + assert!(TensorType::ALL.contains(&TensorType::Nvfp4)); + assert!(TensorType::ALL.contains(&TensorType::Mxfp4)); + for tensor_type in TensorType::ALL { + assert_eq!( + TensorType::parse(tensor_type.as_ggml_name()).unwrap(), + *tensor_type + ); + } + } + + fn pinned_llama_quant_option_names() -> Vec { + let quantize_cpp = repo_root().join(".deps/llama.cpp/tools/quantize/quantize.cpp"); + let source = fs::read_to_string(&quantize_cpp) + .unwrap_or_else(|err| panic!("read {}: {err}", quantize_cpp.display())); + let table = source + .split("static const std::vector QUANT_OPTIONS = {") + .nth(1) + .and_then(|rest| rest.split_once("};").map(|(table, _)| table)) + .expect("find QUANT_OPTIONS table"); + table + .lines() + .filter_map(|line| { + let line = line.trim(); + line.strip_prefix("{ \"") + .and_then(|rest| rest.split_once('"').map(|(name, _)| name.to_string())) + }) + .collect() + } + + fn repo_root() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .parent() + .and_then(|path| path.parent()) + .expect("crate lives under crates/skippy-quantize") + .to_path_buf() + } +} diff --git a/crates/skippy-quantize/src/validation_commands.rs b/crates/skippy-quantize/src/validation_commands.rs new file mode 100644 index 0000000000..8d4a05fa3e --- /dev/null +++ b/crates/skippy-quantize/src/validation_commands.rs @@ -0,0 +1,136 @@ +use std::fs; +use std::path::Path; + +use anyhow::{Context, Result, ensure}; + +use crate::command_reports::{SplitValidation, TensorTypeValidation}; +use crate::manifest::{manifest_progress, read_manifest}; +use crate::output::{ + format_shard_ranges, format_window, print_info, print_json_pretty, print_progress_line, + print_success, print_warn, +}; +use crate::quantize::ensure_tensor_type_entry; +use crate::splits::{Progress, split_status, split_status_for_basename}; + +pub(crate) fn run_status(manifest_path: &Path, json: bool) -> Result<()> { + let manifest = read_manifest(manifest_path)?; + let progress = manifest_progress(&manifest)?; + if json { + print_json_pretty(&progress)?; + } else { + print_progress(&progress); + } + Ok(()) +} + +pub(crate) fn run_next_window(manifest_path: &Path, json: bool) -> Result<()> { + let manifest = read_manifest(manifest_path)?; + let progress = manifest_progress(&manifest)?; + if json { + print_json_pretty(&progress.next_window)?; + } else if let Some(window) = progress.next_window { + print_info(format!("Next window: {}", format_window(window))); + } else { + print_success("No next window; job is complete"); + } + Ok(()) +} + +pub(crate) fn validate_tensor_types_command(path: &Path, json: bool) -> Result<()> { + let validation = validate_tensor_types(path)?; + if json { + print_json_pretty(&validation)?; + } else { + print_success(format!( + "Valid tensor type file: {} entries", + validation.entry_count + )); + } + Ok(()) +} + +pub(crate) fn validate_splits_command( + root: &Path, + prefix: &str, + expected_splits: Option, + basename: Option<&str>, + json: bool, +) -> Result<()> { + let progress = if let Some(basename) = basename { + split_status_for_basename( + root, + prefix, + basename, + expected_splits.context("--expected-splits is required with --basename")?, + )? + } else { + split_status(root, prefix, expected_splits)? + }; + let validation = SplitValidation { + root: root.to_path_buf(), + prefix: prefix.to_string(), + expected_splits: progress.expected_splits, + completed_count: progress.completed_count, + first_missing: progress.first_missing, + last_present: progress.last_present, + complete: progress.complete, + }; + if json { + print_json_pretty(&validation)?; + } else if validation.complete { + print_progress_line( + "split artifact", + validation.completed_count, + validation.expected_splits, + ); + print_success("Split artifact is complete"); + } else { + print_progress_line( + "split artifact", + validation.completed_count, + validation.expected_splits, + ); + print_warn(format!( + "Split artifact is incomplete; first missing shard: {:?}", + validation.first_missing + )); + } + ensure!(validation.complete, "split artifact is incomplete"); + Ok(()) +} + +pub(crate) fn validate_tensor_types(path: &Path) -> Result { + let data = fs::read_to_string(path).with_context(|| format!("read {}", path.display()))?; + let mut entry_count = 0; + for token in data.split_whitespace() { + ensure_tensor_type_entry(token)?; + entry_count += 1; + } + Ok(TensorTypeValidation { + valid: true, + entry_count, + }) +} + +fn print_progress(progress: &Progress) { + print_progress_line( + "job status", + progress.completed_count, + progress.expected_splits, + ); + if progress.complete { + print_success("All shards complete"); + } else { + print_warn(format!("Missing shards: {}", progress.missing_count)); + } + if !progress.missing_ranges.is_empty() { + print_info(format!( + "Missing ranges: {}", + format_shard_ranges(&progress.missing_ranges) + )); + } + match progress.next_window { + Some(window) => print_info(format!("Next window: {}", format_window(window))), + None => print_success("Next window: complete"), + } +} diff --git a/crates/skippy-quantize/src/verify.rs b/crates/skippy-quantize/src/verify.rs new file mode 100644 index 0000000000..28fe4348ca --- /dev/null +++ b/crates/skippy-quantize/src/verify.rs @@ -0,0 +1,171 @@ +use std::path::Path; + +use anyhow::{Result, ensure}; +use serde::Serialize; + +use crate::llama_load::{LlamaLoadOptions, validate_llama_load}; +use crate::manifest::{Manifest, manifest_progress, read_manifest}; +use crate::output::{print_info, print_success}; +use crate::splits::{Progress, split_status_for_basename}; + +#[derive(Debug, Serialize)] +pub struct VerificationReport { + pub root: String, + pub prefix: String, + pub basename: String, + pub expected_splits: u32, + pub completed_count: usize, + pub first_missing: Option, + pub last_present: Option, + pub first_shard: String, + pub last_shard: String, + pub complete: bool, +} + +pub fn verify_manifest(manifest: &Manifest) -> Result { + let progress = split_status_for_basename( + &manifest.target, + &manifest.target_prefix, + &manifest.output_basename, + manifest.expected_splits, + )?; + let report = report_from_progress(manifest, progress); + ensure!( + report.complete, + "manifest artifact is incomplete: {}/{} shards first_missing={:?}", + report.completed_count, + report.expected_splits, + report.first_missing + ); + Ok(report) +} + +pub fn verify_manifest_path_if_complete(path: &Path) -> Result> { + let manifest = read_manifest(path)?; + let progress = manifest_progress(&manifest)?; + if !progress.complete { + return Ok(None); + } + verify_manifest(&manifest).map(Some) +} + +#[derive(Debug, Clone, Copy)] +pub(crate) struct VerifyOnCompleteOptions<'a> { + pub(crate) enabled: bool, + pub(crate) llama_load: bool, + pub(crate) llama_cli: Option<&'a Path>, + pub(crate) check_tensors: bool, +} + +pub fn print_verify_on_complete( + manifest_path: &Path, + options: VerifyOnCompleteOptions<'_>, +) -> Result<()> { + if !options.enabled { + return Ok(()); + } + let manifest = read_manifest(manifest_path)?; + let Some(report) = verify_manifest_path_if_complete(manifest_path)? else { + return Ok(()); + }; + print_success(format!( + "Verified complete artifact: {}/{} shards prefix={} basename={}", + report.completed_count, report.expected_splits, report.prefix, report.basename + )); + if options.llama_load || options.llama_cli.is_some() { + let llama_load = validate_llama_load( + &first_artifact_path(&manifest), + options.llama_cli, + LlamaLoadOptions { + check_tensors: options.check_tensors, + }, + )?; + print_info(format!( + "llama load validation passed: model={} llama_cli={}", + llama_load.model.display(), + llama_load.llama_cli.display() + )); + } + Ok(()) +} + +pub(crate) fn first_artifact_path(manifest: &Manifest) -> std::path::PathBuf { + let root = manifest.target.join(&manifest.target_prefix); + let unsplit = root.join(format!("{}.gguf", manifest.output_basename)); + if manifest.expected_splits == 1 && unsplit.is_file() { + return unsplit; + } + root.join(format!( + "{}-00001-of-{:05}.gguf", + manifest.output_basename, manifest.expected_splits + )) +} + +fn report_from_progress(manifest: &Manifest, progress: Progress) -> VerificationReport { + VerificationReport { + root: manifest.target.display().to_string(), + prefix: manifest.target_prefix.clone(), + basename: manifest.output_basename.clone(), + expected_splits: manifest.expected_splits, + completed_count: progress.completed_count, + first_missing: progress.first_missing, + last_present: progress.last_present, + first_shard: shard_name(&manifest.output_basename, 1, manifest.expected_splits), + last_shard: shard_name( + &manifest.output_basename, + manifest.expected_splits, + manifest.expected_splits, + ), + complete: progress.complete, + } +} + +fn shard_name(basename: &str, index: u32, total: u32) -> String { + format!("{basename}-{index:05}-of-{total:05}.gguf") +} + +#[cfg(test)] +mod tests { + use std::fs; + use std::path::PathBuf; + + use crate::manifest::MANIFEST_VERSION; + use crate::records::unix_timestamp_ms; + use crate::types::JobKind; + + use super::*; + + #[test] + fn verifies_exact_manifest_artifact() { + let root = std::env::temp_dir().join(format!( + "skippy-quantize-verify-test-{}", + unix_timestamp_ms() + )); + let prefix_root = root.join("Q2_K"); + fs::create_dir_all(&prefix_root).unwrap(); + fs::write(prefix_root.join("out-00001-of-00002.gguf"), b"1").unwrap(); + fs::write(prefix_root.join("out-00002-of-00002.gguf"), b"2").unwrap(); + + let manifest = Manifest { + schema_version: MANIFEST_VERSION, + kind: JobKind::QuantizeGguf, + source: PathBuf::from("/source"), + source_prefix: Some("BF16".to_string()), + target: root.clone(), + target_prefix: "Q2_K".to_string(), + output_basename: "out".to_string(), + expected_splits: 2, + window_size: 1, + quant: Some("Q2_K".to_string()), + output_type: None, + tensor_type_file: None, + tensor_type_recipe: None, + }; + + let report = verify_manifest(&manifest).unwrap(); + assert!(report.complete); + assert_eq!(report.first_shard, "out-00001-of-00002.gguf"); + assert_eq!(report.last_shard, "out-00002-of-00002.gguf"); + fs::remove_dir_all(root).unwrap(); + } +} diff --git a/crates/skippy-quantize/src/verify_command.rs b/crates/skippy-quantize/src/verify_command.rs new file mode 100644 index 0000000000..0bb934bf78 --- /dev/null +++ b/crates/skippy-quantize/src/verify_command.rs @@ -0,0 +1,51 @@ +use std::path::Path; + +use anyhow::Result; + +use crate::llama_load::{LlamaLoadOptions, validate_llama_load}; +use crate::manifest::read_manifest; +use crate::output::{print_info, print_json_pretty, print_success}; +use crate::verify::{first_artifact_path, verify_manifest}; + +pub(crate) fn verify_job( + manifest_path: &Path, + llama_load: bool, + llama_cli: Option<&Path>, + check_tensors: bool, + json: bool, +) -> Result<()> { + let manifest = read_manifest(manifest_path)?; + let report = verify_manifest(&manifest)?; + let llama_load = if llama_load || llama_cli.is_some() { + Some(validate_llama_load( + &first_artifact_path(&manifest), + llama_cli, + LlamaLoadOptions { check_tensors }, + )?) + } else { + None + }; + if json { + if let Some(llama_load) = llama_load { + print_json_pretty(&serde_json::json!({ + "artifact": report, + "llama_load": llama_load, + }))?; + } else { + print_json_pretty(&report)?; + } + } else { + print_success(format!( + "Verified artifact: {}/{} shards prefix={} basename={}", + report.completed_count, report.expected_splits, report.prefix, report.basename + )); + if let Some(llama_load) = llama_load { + print_info(format!( + "llama load valid: model={} llama_cli={}", + llama_load.model.display(), + llama_load.llama_cli.display() + )); + } + } + Ok(()) +} diff --git a/crates/skippy-quantize/src/window_loop.rs b/crates/skippy-quantize/src/window_loop.rs new file mode 100644 index 0000000000..b817021707 --- /dev/null +++ b/crates/skippy-quantize/src/window_loop.rs @@ -0,0 +1,45 @@ +use anyhow::Result; + +use crate::output::{print_info, print_success}; + +pub(crate) fn run_window_loop( + label: &str, + max_windows: Option, + mut run_once: F, +) -> Result<()> +where + F: FnMut() -> Result, +{ + let mut completed = 0_u32; + loop { + if max_windows.is_some_and(|max| completed >= max) { + print_info(format!( + "{label} loop stopped after {completed} completed window(s)" + )); + return Ok(()); + } + if !run_once()? { + print_success(format!( + "{label} loop complete after {completed} completed window(s)" + )); + return Ok(()); + } + completed += 1; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn run_loop_honors_max_windows() { + let mut calls = 0_u32; + run_window_loop("test", Some(2), || { + calls += 1; + Ok(true) + }) + .unwrap(); + assert_eq!(calls, 2); + } +} diff --git a/crates/skippy-quantize/tests/direct_convert_cli.rs b/crates/skippy-quantize/tests/direct_convert_cli.rs new file mode 100644 index 0000000000..cf8d7f0140 --- /dev/null +++ b/crates/skippy-quantize/tests/direct_convert_cli.rs @@ -0,0 +1,89 @@ +use std::fs; +use std::path::{Path, PathBuf}; +use std::process::Command; +use std::sync::atomic::{AtomicU64, Ordering}; + +static TEMP_DIR_COUNTER: AtomicU64 = AtomicU64::new(0); + +#[test] +fn direct_convert_native_preflight_resolves_auto_output_type() { + let root = unique_temp_dir(); + let checkpoint = root.join("checkpoint"); + fs::create_dir_all(&checkpoint).unwrap(); + write_safetensor( + &checkpoint.join("model.safetensors"), + &[( + "model.layers.0.self_attn.q_proj.weight", + "BF16", + &[2, 2], + &[1, 2, 3, 4, 5, 6, 7, 8], + )], + ); + + let output = Command::new(env!("CARGO_BIN_EXE_skippy-quantize")) + .args(["convert", "--preflight-only", "--json"]) + .arg(&checkpoint) + .output() + .expect("skippy-quantize command should run"); + + assert!( + output.status.success(), + "preflight should succeed: stderr={}", + String::from_utf8_lossy(&output.stderr) + ); + let stdout = String::from_utf8(output.stdout).expect("stdout should be utf-8"); + let report: serde_json::Value = serde_json::from_str(&stdout) + .unwrap_or_else(|err| panic!("parse preflight JSON: {err}\n{stdout}")); + + assert_eq!(report["backend_kind"], "native-rust"); + assert_eq!(report["backend_ready"], true); + assert_eq!(report["expected_target_shards"], 1); + assert_eq!(report["next_window"]["first_split"], 1); + assert_eq!(report["next_window"]["last_split"], 1); + let manifest_path = report["manifest_path"] + .as_str() + .expect("manifest_path should be a string"); + assert!( + manifest_path.ends_with("/checkpoint/.checkpoint-bf16.bf16.skippy-convert.json"), + "native auto output should resolve to bf16 manifest path, got {manifest_path}" + ); + + fs::remove_dir_all(root).unwrap(); +} + +fn unique_temp_dir() -> PathBuf { + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let counter = TEMP_DIR_COUNTER.fetch_add(1, Ordering::Relaxed); + std::env::temp_dir().join(format!( + "skippy-convert-cli-test-{}-{nanos}-{counter}", + std::process::id() + )) +} + +fn write_safetensor(path: &Path, tensors: &[(&str, &str, &[u64], &[u8])]) { + let mut offset = 0_u64; + let mut entries = serde_json::Map::new(); + for (name, dtype, shape, bytes) in tensors { + let end = offset + bytes.len() as u64; + entries.insert( + (*name).to_string(), + serde_json::json!({ + "dtype": dtype, + "shape": shape, + "data_offsets": [offset, end], + }), + ); + offset = end; + } + let header = serde_json::Value::Object(entries).to_string(); + let mut bytes = Vec::new(); + bytes.extend_from_slice(&(header.len() as u64).to_le_bytes()); + bytes.extend_from_slice(header.as_bytes()); + for (_, _, _, tensor_bytes) in tensors { + bytes.extend_from_slice(tensor_bytes); + } + fs::write(path, bytes).unwrap(); +} diff --git a/crates/skippy-quantize/tests/direct_quantize_cli.rs b/crates/skippy-quantize/tests/direct_quantize_cli.rs new file mode 100644 index 0000000000..0ac335c697 --- /dev/null +++ b/crates/skippy-quantize/tests/direct_quantize_cli.rs @@ -0,0 +1,190 @@ +use std::fs; +use std::path::PathBuf; +use std::process::Command; +use std::sync::atomic::{AtomicU64, Ordering}; + +static TEMP_DIR_COUNTER: AtomicU64 = AtomicU64::new(0); + +#[test] +fn direct_quantize_preflight_reports_requested_window_for_native_backends() { + let root = unique_temp_dir(); + let source = root.join("source"); + let target = root.join("target"); + let source_prefix = source.join("BF16"); + let target_prefix = target.join("Q4"); + fs::create_dir_all(&source_prefix).unwrap(); + fs::create_dir_all(&target_prefix).unwrap(); + for index in 1..=3 { + fs::write( + source_prefix.join(format!("model-0000{index}-of-00003.gguf")), + b"source shard", + ) + .unwrap(); + } + fs::write( + target_prefix.join("model-q4-00002-of-00003.gguf"), + b"already done", + ) + .unwrap(); + + let output = Command::new(env!("CARGO_BIN_EXE_skippy-quantize")) + .args([ + "quantize", + "--backend", + "llama-api", + "--preflight-only", + "--json", + "--keep-split", + "--first-split", + "2", + "--last-split", + "3", + ]) + .arg(source_prefix.join("model-00001-of-00003.gguf")) + .arg(target_prefix.join("model-q4.gguf")) + .arg("Q4_K") + .output() + .expect("skippy-quantize command should run"); + + let stdout = String::from_utf8(output.stdout).expect("stdout should be utf-8"); + assert!( + stdout.contains(r#""requested_window""#), + "preflight should report requested_window, got: {stdout}" + ); + assert!( + stdout.contains(r#""next_requested_window""#), + "preflight should report next_requested_window, got: {stdout}" + ); + assert!( + stdout.contains(r#""first_split": 2"#) && stdout.contains(r#""last_split": 3"#), + "preflight should include requested 2..3 window, got: {stdout}" + ); + assert!( + stdout.contains(r#""first_split": 3"#), + "preflight should skip completed split 2 and report split 3 next, got: {stdout}" + ); + + fs::remove_dir_all(root).unwrap(); +} + +#[test] +fn direct_quantize_preflight_supports_current_directory_no_output_shape() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + fs::write(root.join("model.gguf"), b"source shard").unwrap(); + + let output = Command::new(env!("CARGO_BIN_EXE_skippy-quantize")) + .current_dir(&root) + .args([ + "quantize", + "--backend", + "llama-api", + "--preflight-only", + "--json", + "model.gguf", + "Q4_K", + ]) + .output() + .expect("skippy-quantize command should run"); + + let stdout = String::from_utf8(output.stdout).expect("stdout should be utf-8"); + let report: serde_json::Value = serde_json::from_str(&stdout) + .unwrap_or_else(|err| panic!("parse preflight JSON: {err}\n{stdout}")); + + assert_eq!(report["backend_kind"], "llama-api"); + assert_eq!(report["backend_ready"], true); + assert_eq!(report["source_complete"], true); + assert_eq!(report["expected_source_shards"], 1); + assert_eq!(report["next_window"]["first_split"], 1); + let manifest_path = report["manifest_path"] + .as_str() + .expect("manifest_path should be a string"); + assert!( + manifest_path.ends_with(".ggml-model-Q4_K.Q4_K.skippy-quantize.json"), + "no-output quantize shape should derive upstream-style sidecar, got {manifest_path}" + ); + + fs::remove_dir_all(root).unwrap(); +} + +#[test] +fn direct_quantize_preflight_accepts_base_quant_with_tensor_file() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + fs::write(root.join("model.gguf"), b"source shard").unwrap(); + fs::write(root.join("tensor-types.txt"), b"blk.0.weight=Q8_0\n").unwrap(); + + let output = Command::new(env!("CARGO_BIN_EXE_skippy-quantize")) + .current_dir(&root) + .args([ + "quantize", + "--backend", + "llama-api", + "--tensor-type-file", + "tensor-types.txt", + "--preflight-only", + "--json", + "model.gguf", + "Q3_K_S", + ]) + .output() + .expect("skippy-quantize command should run"); + + let stdout = String::from_utf8(output.stdout).expect("stdout should be utf-8"); + let report: serde_json::Value = serde_json::from_str(&stdout) + .unwrap_or_else(|err| panic!("parse preflight JSON: {err}\n{stdout}")); + let manifest_path = report["manifest_path"] + .as_str() + .expect("manifest_path should be a string"); + assert!( + manifest_path.ends_with(".ggml-model-Q3_K_S.Q3_K_S.skippy-quantize.json"), + "base quant should drive output and sidecar names, got {manifest_path}" + ); + + fs::remove_dir_all(root).unwrap(); +} + +#[test] +fn direct_quantize_preflight_rejects_profile_quant_label() { + let root = unique_temp_dir(); + fs::create_dir_all(&root).unwrap(); + fs::write(root.join("model.gguf"), b"source shard").unwrap(); + + let output = Command::new(env!("CARGO_BIN_EXE_skippy-quantize")) + .current_dir(&root) + .args([ + "quantize", + "--backend", + "llama-api", + "--preflight-only", + "--json", + "model.gguf", + "UD-Q3_K_S", + ]) + .output() + .expect("skippy-quantize command should run"); + + assert!( + !output.status.success(), + "preflight should reject custom profile labels as quant modes" + ); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!( + stderr.contains("custom tensor-type recipes"), + "error should explain profile labels are not quant modes, got: {stderr}" + ); + + fs::remove_dir_all(root).unwrap(); +} + +fn unique_temp_dir() -> PathBuf { + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let counter = TEMP_DIR_COUNTER.fetch_add(1, Ordering::Relaxed); + std::env::temp_dir().join(format!( + "skippy-quantize-cli-test-{}-{nanos}-{counter}", + std::process::id() + )) +} diff --git a/crates/skippy-runtime/src/lib.rs b/crates/skippy-runtime/src/lib.rs index 3d9a90b4d0..d817a39d0a 100644 --- a/crates/skippy-runtime/src/lib.rs +++ b/crates/skippy-runtime/src/lib.rs @@ -18,15 +18,16 @@ use skippy_ffi::{ ChatMessage as RawChatMessage, Error as RawError, GenerationSignalWindow as RawGenerationSignalWindow, KvPageDesc as RawKvPageDesc, LoadMode, LogitBias as RawLogitBias, Model as RawModel, ModelInfo as RawModelInfo, - RuntimeConfig as RawRuntimeConfig, SamplingConfig as RawSamplingConfig, Session as RawSession, - SlicePlan as RawSlicePlan, TensorInfo as RawTensorInfo, TensorRole, - TokenSignal as RawTokenSignal, + NativeMtpDraft as RawNativeMtpDraft, RuntimeConfig as RawRuntimeConfig, + SamplingConfig as RawSamplingConfig, Session as RawSession, SlicePlan as RawSlicePlan, + TensorInfo as RawTensorInfo, TensorRole, TokenSignal as RawTokenSignal, }; use tokio::sync::mpsc; mod devices; pub mod package; mod runtime_events; +pub mod spd; pub const MAX_LOGIT_BIAS: usize = 256; pub const GGML_TYPE_F16: u32 = 1; @@ -1059,6 +1060,21 @@ impl From for ActivationDesc { } } +fn empty_raw_activation_desc() -> RawActivationDesc { + RawActivationDesc { + version: 0, + dtype: ActivationDType::Unknown, + layout: ActivationLayout::Opaque, + producer_stage_index: -1, + layer_start: 0, + layer_end: 0, + token_count: 0, + sequence_count: 0, + payload_bytes: 0, + flags: 0, + } +} + #[derive(Debug, Clone, PartialEq, Eq)] pub struct ActivationFrame { pub desc: ActivationDesc, @@ -1198,6 +1214,41 @@ pub struct StageSession { token_count: u64, } +pub struct DecodeBatchRequest<'a> { + pub session: &'a mut StageSession, + pub token_id: i32, + pub sampling: Option<&'a SamplingConfig>, +} + +pub struct DecodeFrameBatchRequest<'a> { + pub session: &'a mut StageSession, + pub token_id: i32, + pub sampling: Option<&'a SamplingConfig>, + pub input: Option<&'a ActivationFrame>, +} + +pub struct DecodeFrameBatchOutput { + pub predicted_token: i32, + pub output: ActivationFrame, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct NativeMtpDraft { + pub token_id: i32, + pub proposal_compute_us: i64, +} + +const NATIVE_MTP_DRAFT_VERSION: u32 = 1; + +impl NativeMtpDraft { + fn from_raw(raw: RawNativeMtpDraft) -> Option { + (raw.available && raw.version == NATIVE_MTP_DRAFT_VERSION).then_some(Self { + token_id: raw.token_id, + proposal_compute_us: raw.proposal_compute_us, + }) + } +} + #[derive(Debug, Clone, PartialEq, Eq)] pub struct MediaInput { pub bytes: Vec, @@ -2702,6 +2753,55 @@ impl StageSession { Ok(predicted_token) } + pub fn decode_batch_sampled(requests: &mut [DecodeBatchRequest<'_>]) -> Result> { + if requests.is_empty() { + return Ok(Vec::new()); + } + + let sessions = requests + .iter_mut() + .map(|request| request.session.raw) + .collect::>(); + let token_ids = requests + .iter() + .map(|request| request.token_id) + .collect::>(); + let raw_sampling = requests + .iter() + .map(|request| request.sampling.map(SamplingConfig::as_raw)) + .collect::>(); + let sampling = raw_sampling + .iter() + .map(|sampling| { + sampling + .as_ref() + .map_or(ptr::null(), |sampling| sampling as *const RawSamplingConfig) + }) + .collect::>(); + let mut predicted_tokens = vec![0_i32; requests.len()]; + let mut error = ptr::null_mut(); + let status = unsafe { + skippy_ffi::skippy_decode_batch_sampled( + sessions.as_ptr(), + token_ids.as_ptr(), + sampling.as_ptr(), + requests.len(), + predicted_tokens.as_mut_ptr(), + predicted_tokens.len(), + &mut error, + ) + }; + ensure_ok(status, error)?; + for request in requests { + request.session.token_count = request + .session + .token_count + .checked_add(1) + .context("session token count overflow")?; + } + Ok(predicted_tokens) + } + pub fn last_token_signal(&mut self) -> Result { let mut signal = RawTokenSignal::default(); let mut error = ptr::null_mut(); @@ -3028,6 +3128,25 @@ impl StageSession { )) } + pub fn decode_step_frame_sampled_mtp_n1( + &mut self, + token_id: i32, + sampling: Option<&SamplingConfig>, + input: Option<&ActivationFrame>, + output_capacity: usize, + ) -> Result<(i32, Option, ActivationFrame)> { + let (predicted_token, mtp_draft, output_desc, output_payload) = + self.decode_step_frame_mtp_n1_raw(token_id, sampling, input, output_capacity)?; + Ok(( + predicted_token, + mtp_draft, + ActivationFrame { + desc: output_desc.into(), + payload: output_payload, + }, + )) + } + fn decode_step_frame_raw( &mut self, token_id: i32, @@ -3088,17 +3207,236 @@ impl StageSession { Ok((predicted_token, output_desc, output_payload)) } + fn decode_step_frame_mtp_n1_raw( + &mut self, + token_id: i32, + sampling: Option<&SamplingConfig>, + input: Option<&ActivationFrame>, + output_capacity: usize, + ) -> Result<(i32, Option, RawActivationDesc, Vec)> { + let input_desc = input.map(|frame| frame.desc.as_raw()); + let input_desc_ptr = input_desc + .as_ref() + .map_or(ptr::null(), |desc| desc as *const RawActivationDesc); + let input_payload_ptr = input.map_or(ptr::null(), |frame| frame.payload.as_ptr().cast()); + let mut output_desc = RawActivationDesc { + version: 0, + dtype: ActivationDType::Unknown, + layout: ActivationLayout::Opaque, + producer_stage_index: -1, + layer_start: 0, + layer_end: 0, + token_count: 0, + sequence_count: 0, + payload_bytes: 0, + flags: 0, + }; + let mut output_payload = vec![0_u8; output_capacity]; + let mut output_bytes = 0usize; + let mut predicted_token = 0_i32; + let mut mtp_draft = RawNativeMtpDraft::default(); + let mut error = ptr::null_mut(); + let raw_sampling = sampling.map(SamplingConfig::as_raw); + let sampling_ptr = raw_sampling + .as_ref() + .map_or(ptr::null(), |sampling| sampling as *const RawSamplingConfig); + let status = unsafe { + skippy_ffi::skippy_decode_step_frame_sampled_mtp_n1( + self.raw, + token_id, + sampling_ptr, + input_desc_ptr, + input_payload_ptr, + &mut output_desc, + output_payload.as_mut_ptr().cast(), + output_payload.len(), + &mut output_bytes, + &mut predicted_token, + &mut mtp_draft, + &mut error, + ) + }; + if status == Status::BufferTooSmall && output_bytes > output_payload.len() { + free_error(error); + return self.decode_step_frame_mtp_n1_raw(token_id, sampling, input, output_bytes); + } + ensure_ok(status, error)?; + output_payload.truncate(output_bytes); + self.token_count = self + .token_count + .checked_add(1) + .context("session token count overflow")?; + Ok(( + predicted_token, + NativeMtpDraft::from_raw(mtp_draft), + output_desc, + output_payload, + )) + } + + pub fn decode_step_frame_batch_sampled( + requests: &mut [DecodeFrameBatchRequest<'_>], + ) -> Result> { + Self::decode_step_frame_batch_sampled_raw(requests, &vec![0; requests.len()]) + } + + fn decode_step_frame_batch_sampled_raw( + requests: &mut [DecodeFrameBatchRequest<'_>], + output_capacities: &[usize], + ) -> Result> { + if requests.is_empty() { + return Ok(Vec::new()); + } + let sessions = requests + .iter_mut() + .map(|request| request.session.raw) + .collect::>(); + let token_ids = requests + .iter() + .map(|request| request.token_id) + .collect::>(); + let raw_sampling = requests + .iter() + .map(|request| request.sampling.map(SamplingConfig::as_raw)) + .collect::>(); + let sampling = raw_sampling + .iter() + .map(|sampling| { + sampling + .as_ref() + .map_or(ptr::null(), |sampling| sampling as *const RawSamplingConfig) + }) + .collect::>(); + let input_descs = requests + .iter() + .map(|request| request.input.map(|frame| frame.desc.as_raw())) + .collect::>(); + let input_desc_ptrs = input_descs + .iter() + .map(|desc| { + desc.as_ref() + .map_or(ptr::null(), |desc| desc as *const RawActivationDesc) + }) + .collect::>(); + let input_payloads = requests + .iter() + .map(|request| { + request + .input + .map_or(ptr::null(), |frame| frame.payload.as_ptr().cast()) + }) + .collect::>(); + let mut output_descs = vec![empty_raw_activation_desc(); requests.len()]; + let mut output_payloads = output_capacities + .iter() + .map(|capacity| vec![0_u8; *capacity]) + .collect::>(); + let output_payload_ptrs = output_payloads + .iter_mut() + .map(|payload| payload.as_mut_ptr().cast()) + .collect::>(); + let mut output_bytes = vec![0_usize; requests.len()]; + let mut predicted_tokens = vec![0_i32; requests.len()]; + let mut error = ptr::null_mut(); + let status = unsafe { + skippy_ffi::skippy_decode_step_frame_batch_sampled( + sessions.as_ptr(), + token_ids.as_ptr(), + sampling.as_ptr(), + input_desc_ptrs.as_ptr(), + input_payloads.as_ptr(), + output_descs.as_mut_ptr(), + output_payload_ptrs.as_ptr(), + output_capacities.as_ptr(), + output_bytes.as_mut_ptr(), + predicted_tokens.as_mut_ptr(), + predicted_tokens.len(), + requests.len(), + &mut error, + ) + }; + if status == Status::BufferTooSmall { + free_error(error); + error = ptr::null_mut(); + if output_bytes + .iter() + .zip(output_capacities.iter()) + .any(|(required, capacity)| required > capacity) + { + return Self::decode_step_frame_batch_sampled_raw(requests, &output_bytes); + } + } + if status == Status::Unsupported { + free_error(error); + return Self::decode_step_frame_batch_sampled_serial(requests); + } + ensure_ok(status, error)?; + for request in requests.iter_mut() { + request.session.token_count = request + .session + .token_count + .checked_add(1) + .context("session token count overflow")?; + } + Ok(output_payloads + .into_iter() + .zip(output_descs) + .zip(output_bytes) + .zip(predicted_tokens) + .map(|(((mut payload, desc), bytes), predicted_token)| { + payload.truncate(bytes); + DecodeFrameBatchOutput { + predicted_token, + output: ActivationFrame { + desc: desc.into(), + payload, + }, + } + }) + .collect()) + } + + fn decode_step_frame_batch_sampled_serial( + requests: &mut [DecodeFrameBatchRequest<'_>], + ) -> Result> { + requests + .iter_mut() + .map(|request| { + let (predicted_token, output) = request.session.decode_step_frame_sampled( + request.token_id, + request.sampling, + request.input, + 0, + )?; + Ok(DecodeFrameBatchOutput { + predicted_token, + output, + }) + }) + .collect() + } + pub fn verify_tokens_frame( &mut self, token_ids: &[i32], input: Option<&ActivationFrame>, output_capacity: usize, + ) -> Result<(Vec, ActivationFrame)> { + self.verify_tokens_frame_sampled(token_ids, None, input, output_capacity) + } + + pub fn verify_tokens_frame_sampled( + &mut self, + token_ids: &[i32], + sampling: Option<&SamplingConfig>, + input: Option<&ActivationFrame>, + output_capacity: usize, ) -> Result<(Vec, ActivationFrame)> { if token_ids.is_empty() { return Err(anyhow!("verify_tokens_frame requires at least one token")); } let (predicted_tokens, output_desc, output_payload) = - self.verify_tokens_frame_raw(token_ids, input, output_capacity)?; + self.verify_tokens_frame_raw(token_ids, sampling, input, output_capacity)?; Ok(( predicted_tokens, ActivationFrame { @@ -3111,6 +3449,7 @@ impl StageSession { fn verify_tokens_frame_raw( &mut self, token_ids: &[i32], + sampling: Option<&SamplingConfig>, input: Option<&ActivationFrame>, output_capacity: usize, ) -> Result<(Vec, RawActivationDesc, Vec)> { @@ -3133,14 +3472,19 @@ impl StageSession { }; let mut output_payload = vec![0_u8; output_capacity]; let mut output_bytes = 0usize; - let mut predicted = vec![0_i32; token_ids.len()]; + let mut predicted = vec![0_i32; token_ids.len().saturating_add(3)]; let mut output_token_count = 0usize; let mut error = ptr::null_mut(); + let raw_sampling = sampling.map(SamplingConfig::as_raw); + let sampling_ptr = raw_sampling + .as_ref() + .map_or(ptr::null(), |sampling| sampling as *const RawSamplingConfig); let status = unsafe { - skippy_ffi::skippy_verify_tokens_frame( + skippy_ffi::skippy_verify_tokens_frame_sampled( self.raw, token_ids.as_ptr(), token_ids.len(), + sampling_ptr, input_desc_ptr, input_payload_ptr, &mut output_desc, @@ -3155,7 +3499,7 @@ impl StageSession { }; if status == Status::BufferTooSmall && output_bytes > output_payload.len() { free_error(error); - return self.verify_tokens_frame_raw(token_ids, input, output_bytes); + return self.verify_tokens_frame_raw(token_ids, sampling, input, output_bytes); } ensure_ok(status, error)?; predicted.truncate(output_token_count); diff --git a/crates/skippy-runtime/src/spd.rs b/crates/skippy-runtime/src/spd.rs new file mode 100644 index 0000000000..a6e35121b0 --- /dev/null +++ b/crates/skippy-runtime/src/spd.rs @@ -0,0 +1,1065 @@ +use std::{ + collections::BTreeMap, + fs, + io::Read, + path::{Component, Path, PathBuf}, +}; + +use anyhow::{Context, Result, bail}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; + +pub const SPD_HEAD_MANIFEST_SCHEMA: &str = "skippy-spd-head/v1"; +pub const TORCH_SPD_HEAD_FORMAT_V10: &str = "torch-speculation-head-v10"; +pub const GENERIC_LAYER_TAP_HEAD_FORMAT_V1: &str = "generic-layer-tap-sidecar-v1"; +pub const SPD_SERVING_CHECKPOINT_FORMAT_SAFETENSORS_V1: &str = "safetensors-spd-head-v1"; +pub const SPD_HEAD_KIND_FIXED_STAGE_V1: &str = "fixed-stage-v1"; +pub const SPD_HEAD_KIND_GENERIC_LAYER_TAP_V1: &str = "generic-layer-tap-v1"; + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +pub struct SpdHeadManifest { + pub schema: String, + pub checkpoint: SpdHeadCheckpoint, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub serving_checkpoint: Option, + pub source: SpdHeadSource, + pub topology: SpdHeadTopology, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +pub struct SpdHeadCheckpoint { + pub path: String, + pub sha256: String, + pub bytes: u64, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +pub struct SpdHeadServingCheckpoint { + pub path: String, + pub sha256: String, + pub bytes: u64, + pub format: String, + pub tensor_count: u32, + pub dtype: String, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +pub struct SpdHeadSource { + pub format: String, + pub reference_repo: Option, + pub base_model_path: String, + pub model_type: Option, + pub checkpoint_version: u32, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +pub struct SpdHeadTopology { + pub hidden_size: u32, + pub vocab_size: u32, + pub draft_vocab_size: u32, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub head_kind: Option, + pub num_stages: u32, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub stage_layer_boundaries: Option>, + pub num_spec_layers: u32, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub max_taps: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tap_feature_size: Option, + pub trained_with_use_deepest: bool, + pub shallow_hidden_layer_indices: Vec>, + pub spec_init_from_base_layers: Option>, + pub draft_token_ids: Option>, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SpdHeadRuntimeProfile<'a> { + pub base_model_path: Option<&'a str>, + pub hidden_size: u32, + pub vocab_size: u32, + pub num_stages: u32, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SpdSafetensorsIndex { + pub tensors: BTreeMap, + pub metadata: BTreeMap, + pub data_start: u64, + pub data_len: u64, +} + +#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] +pub struct SpdSafetensorsTensor { + pub dtype: String, + pub shape: Vec, + pub data_offsets: [u64; 2], +} + +impl SpdHeadManifest { + pub fn from_path(path: impl AsRef) -> Result { + let path = path.as_ref(); + let manifest: Self = serde_json::from_slice( + &fs::read(path) + .with_context(|| format!("read SPD head manifest {}", path.display()))?, + ) + .with_context(|| format!("parse SPD head manifest {}", path.display()))?; + manifest.validate()?; + Ok(manifest) + } + + pub fn validate(&self) -> Result<()> { + if self.schema != SPD_HEAD_MANIFEST_SCHEMA { + bail!( + "unsupported SPD head manifest schema {}; expected {}", + self.schema, + SPD_HEAD_MANIFEST_SCHEMA + ); + } + validate_source_format(&self.source)?; + if self.source.base_model_path.trim().is_empty() { + bail!("SPD head manifest base_model_path must not be empty"); + } + self.checkpoint.validate()?; + if let Some(serving_checkpoint) = &self.serving_checkpoint { + serving_checkpoint.validate()?; + } + self.topology.validate()?; + Ok(()) + } + + pub fn checkpoint_path(&self, manifest_path: impl AsRef) -> Result { + let manifest_path = manifest_path.as_ref(); + let base = manifest_path + .parent() + .with_context(|| format!("resolve parent for {}", manifest_path.display()))?; + Ok(base.join(safe_relative_manifest_path(&self.checkpoint.path)?)) + } + + pub fn serving_checkpoint_path(&self, manifest_path: impl AsRef) -> Result { + let serving_checkpoint = self + .serving_checkpoint + .as_ref() + .context("SPD manifest does not include a serving checkpoint")?; + let manifest_path = manifest_path.as_ref(); + let base = manifest_path + .parent() + .with_context(|| format!("resolve parent for {}", manifest_path.display()))?; + Ok(base.join(safe_relative_manifest_path(&serving_checkpoint.path)?)) + } + + pub fn verify_checkpoint(&self, manifest_path: impl AsRef) -> Result<()> { + let checkpoint_path = self.checkpoint_path(manifest_path)?; + verify_checkpoint_artifact( + "SPD checkpoint", + &checkpoint_path, + self.checkpoint.bytes, + &self.checkpoint.sha256, + ) + } + + pub fn verify_serving_checkpoint( + &self, + manifest_path: impl AsRef, + ) -> Result { + let serving_checkpoint = self + .serving_checkpoint + .as_ref() + .context("SPD manifest does not include a serving checkpoint")?; + let path = self.serving_checkpoint_path(manifest_path)?; + verify_checkpoint_artifact( + "SPD serving checkpoint", + &path, + serving_checkpoint.bytes, + &serving_checkpoint.sha256, + )?; + let index = SpdSafetensorsIndex::from_path(&path)?; + serving_checkpoint.validate_index(&index)?; + Ok(index) + } + + pub fn serving_checkpoint_index( + &self, + manifest_path: impl AsRef, + ) -> Result { + let serving_checkpoint = self + .serving_checkpoint + .as_ref() + .context("SPD manifest does not include a serving checkpoint")?; + let index = SpdSafetensorsIndex::from_path(self.serving_checkpoint_path(manifest_path)?)?; + serving_checkpoint.validate_index(&index)?; + Ok(index) + } + + pub fn ensure_serving_checkpoint_for_runtime( + &self, + manifest_path: impl AsRef, + ) -> Result { + let index = self.verify_serving_checkpoint(manifest_path)?; + self.ensure_serving_tensor_shapes(&index)?; + Ok(index) + } + + fn ensure_serving_tensor_shapes(&self, index: &SpdSafetensorsIndex) -> Result<()> { + if self.topology.head_kind() == SPD_HEAD_KIND_GENERIC_LAYER_TAP_V1 { + return self.ensure_generic_layer_tap_tensor_shapes(index); + } + self.ensure_fixed_stage_tensor_shapes(index) + } + + fn ensure_fixed_stage_tensor_shapes(&self, index: &SpdSafetensorsIndex) -> Result<()> { + let hidden_size = self.topology.hidden_size as u64; + for (stage, indices) in self + .topology + .shallow_hidden_layer_indices + .iter() + .enumerate() + { + let expected_width = hidden_size + .checked_mul(indices.len() as u64) + .context("SPD stage projection width overflow")?; + index.ensure_tensor_shape( + &format!("stage_projs.{stage}.weight"), + &[hidden_size, expected_width], + )?; + } + index.ensure_tensor_shape("g0_proj.weight", &[hidden_size, hidden_size])?; + index.ensure_tensor_shape( + "lm_head.weight", + &[self.topology.draft_vocab_size as u64, hidden_size], + )?; + for layer in 0..self.topology.num_spec_layers { + index.ensure_tensor_shape( + &format!("spec_layers.{layer}.input_layernorm.weight"), + &[hidden_size], + )?; + index.ensure_tensor_shape( + &format!("spec_layers.{layer}.post_attention_layernorm.weight"), + &[hidden_size], + )?; + } + Ok(()) + } + + fn ensure_generic_layer_tap_tensor_shapes(&self, index: &SpdSafetensorsIndex) -> Result<()> { + let hidden_size = self.topology.hidden_size as u64; + let tap_feature_size = + self.topology + .tap_feature_size + .context("generic SPD head missing tap_feature_size")? as u64; + index.ensure_tensor_shape("tap_proj.weight", &[hidden_size, hidden_size])?; + index.ensure_tensor_shape("tap_proj.bias", &[hidden_size])?; + index.ensure_tensor_shape("depth_proj.weight", &[hidden_size, tap_feature_size])?; + index.ensure_tensor_shape("depth_proj.bias", &[hidden_size])?; + index.ensure_tensor_shape("tap_norm.weight", &[hidden_size])?; + index.ensure_tensor_shape("tap_norm.bias", &[hidden_size])?; + index.ensure_tensor_shape("output_norm.weight", &[hidden_size])?; + index.ensure_tensor_shape("output_norm.bias", &[hidden_size])?; + for layer in 0..self.topology.num_spec_layers { + index.ensure_tensor_shape( + &format!("draft_heads.{layer}.weight"), + &[self.topology.draft_vocab_size as u64, hidden_size], + )?; + index.ensure_tensor_shape( + &format!("draft_heads.{layer}.bias"), + &[self.topology.draft_vocab_size as u64], + )?; + } + Ok(()) + } + + pub fn ensure_runtime_compatible(&self, profile: &SpdHeadRuntimeProfile<'_>) -> Result<()> { + match profile.base_model_path { + Some(base_model_path) if self.source.base_model_path != base_model_path => { + bail!( + "SPD head was trained for base model {}; runtime model is {}", + self.source.base_model_path, + base_model_path + ) + } + _ => {} + } + if self.topology.hidden_size != profile.hidden_size { + bail!( + "SPD head hidden_size {} does not match runtime hidden_size {}", + self.topology.hidden_size, + profile.hidden_size + ); + } + if self.topology.vocab_size != profile.vocab_size { + bail!( + "SPD head vocab_size {} does not match runtime vocab_size {}", + self.topology.vocab_size, + profile.vocab_size + ); + } + if self.topology.head_kind() != SPD_HEAD_KIND_GENERIC_LAYER_TAP_V1 + && self.topology.num_stages != profile.num_stages + { + bail!( + "SPD head num_stages {} does not match runtime num_stages {}", + self.topology.num_stages, + profile.num_stages + ); + } + Ok(()) + } +} + +impl SpdHeadCheckpoint { + fn validate(&self) -> Result<()> { + let _ = safe_relative_manifest_path(&self.path)?; + if self.bytes == 0 { + bail!("SPD checkpoint bytes must be greater than zero"); + } + validate_sha256_digest("SPD checkpoint sha256", &self.sha256) + } +} + +impl SpdHeadServingCheckpoint { + fn validate(&self) -> Result<()> { + let _ = safe_relative_manifest_path(&self.path)?; + if self.bytes == 0 { + bail!("SPD serving checkpoint bytes must be greater than zero"); + } + if self.format != SPD_SERVING_CHECKPOINT_FORMAT_SAFETENSORS_V1 { + bail!( + "unsupported SPD serving checkpoint format {}; expected {}", + self.format, + SPD_SERVING_CHECKPOINT_FORMAT_SAFETENSORS_V1 + ); + } + if self.tensor_count == 0 { + bail!("SPD serving checkpoint tensor_count must be greater than zero"); + } + if self.dtype.trim().is_empty() { + bail!("SPD serving checkpoint dtype must not be empty"); + } + validate_sha256_digest("SPD serving checkpoint sha256", &self.sha256) + } + + fn validate_index(&self, index: &SpdSafetensorsIndex) -> Result<()> { + if index.tensors.len() != self.tensor_count as usize { + bail!( + "SPD serving checkpoint tensor_count mismatch: expected {}, got {}", + self.tensor_count, + index.tensors.len() + ); + } + if self.dtype != "mixed" + && index + .tensors + .values() + .any(|tensor| tensor.dtype != self.dtype) + { + bail!( + "SPD serving checkpoint dtype mismatch: expected all tensors to be {}", + self.dtype + ); + } + match index.metadata.get("format") { + Some(format) if format != SPD_SERVING_CHECKPOINT_FORMAT_SAFETENSORS_V1 => { + bail!( + "SPD serving checkpoint metadata format {}; expected {}", + format, + SPD_SERVING_CHECKPOINT_FORMAT_SAFETENSORS_V1 + ); + } + _ => {} + } + Ok(()) + } +} + +impl SpdSafetensorsIndex { + pub fn from_path(path: impl AsRef) -> Result { + let path = path.as_ref(); + let mut file = fs::File::open(path) + .with_context(|| format!("open SPD safetensors checkpoint {}", path.display()))?; + let file_len = file + .metadata() + .with_context(|| format!("stat SPD safetensors checkpoint {}", path.display()))? + .len(); + let mut header_len_bytes = [0_u8; 8]; + file.read_exact(&mut header_len_bytes) + .with_context(|| format!("read SPD safetensors header length {}", path.display()))?; + let header_len = u64::from_le_bytes(header_len_bytes); + if header_len == 0 { + bail!("SPD safetensors header must not be empty"); + } + if header_len > file_len.saturating_sub(8) { + bail!( + "SPD safetensors header length {} exceeds file length {}", + header_len, + file_len + ); + } + let header_len_usize = + usize::try_from(header_len).context("SPD safetensors header is too large")?; + let mut header = vec![0_u8; header_len_usize]; + file.read_exact(&mut header) + .with_context(|| format!("read SPD safetensors header {}", path.display()))?; + Self::from_header_bytes(&header, 8 + header_len, file_len) + } + + fn from_header_bytes(header: &[u8], data_start: u64, file_len: u64) -> Result { + if data_start > file_len { + bail!("SPD safetensors data section starts past end of file"); + } + let data_len = file_len - data_start; + let mut metadata = BTreeMap::new(); + let mut tensors = BTreeMap::new(); + let value: serde_json::Value = + serde_json::from_slice(header).context("parse SPD safetensors header JSON")?; + let serde_json::Value::Object(entries) = value else { + bail!("SPD safetensors header must be a JSON object"); + }; + + for (name, value) in entries { + if name == "__metadata__" { + metadata = + serde_json::from_value(value).context("parse SPD safetensors metadata map")?; + continue; + } + if name.trim().is_empty() { + bail!("SPD safetensors tensor names must not be empty"); + } + let tensor: SpdSafetensorsTensor = serde_json::from_value(value) + .with_context(|| format!("parse SPD safetensors tensor metadata {name}"))?; + tensor.validate(&name, data_len)?; + tensors.insert(name, tensor); + } + validate_safetensors_ranges(&tensors, data_len)?; + Ok(Self { + tensors, + metadata, + data_start, + data_len, + }) + } + + pub fn ensure_tensor_shape(&self, name: &str, expected_shape: &[u64]) -> Result<()> { + let tensor = self + .tensors + .get(name) + .with_context(|| format!("SPD serving checkpoint is missing tensor {name}"))?; + if tensor.shape != expected_shape { + bail!( + "SPD serving checkpoint tensor {name} shape mismatch: expected {:?}, got {:?}", + expected_shape, + tensor.shape + ); + } + Ok(()) + } +} + +impl SpdSafetensorsTensor { + fn validate(&self, name: &str, data_len: u64) -> Result<()> { + let [start, end] = self.data_offsets; + if start > end || end > data_len { + bail!( + "SPD safetensors tensor {name} has invalid data offsets {:?} for data length {}", + self.data_offsets, + data_len + ); + } + let expected_bytes = tensor_byte_len(&self.dtype, &self.shape) + .with_context(|| format!("validate SPD safetensors tensor {name}"))?; + if end - start != expected_bytes { + bail!( + "SPD safetensors tensor {name} byte length mismatch: offsets describe {}, shape/dtype describe {}", + end - start, + expected_bytes + ); + } + Ok(()) + } +} + +impl SpdHeadTopology { + fn validate(&self) -> Result<()> { + let head_kind = self.head_kind(); + if head_kind != SPD_HEAD_KIND_FIXED_STAGE_V1 + && head_kind != SPD_HEAD_KIND_GENERIC_LAYER_TAP_V1 + { + bail!("unsupported SPD head kind {head_kind}"); + } + if self.hidden_size == 0 { + bail!("SPD head hidden_size must be greater than zero"); + } + if self.vocab_size == 0 { + bail!("SPD head vocab_size must be greater than zero"); + } + if self.draft_vocab_size == 0 || self.draft_vocab_size > self.vocab_size { + bail!( + "SPD head draft_vocab_size {} must be in 1..={}", + self.draft_vocab_size, + self.vocab_size + ); + } + if self.num_stages == 0 { + bail!("SPD head num_stages must be greater than zero"); + } + if self.num_spec_layers == 0 { + bail!("SPD head num_spec_layers must be greater than zero"); + } + if head_kind == SPD_HEAD_KIND_FIXED_STAGE_V1 { + self.validate_fixed_stage_topology()?; + } else { + self.validate_generic_layer_tap_topology()?; + } + match &self.spec_init_from_base_layers { + Some(indices) if indices.len() != self.num_spec_layers as usize => { + bail!( + "SPD head spec_init_from_base_layers length {} must match num_spec_layers {}", + indices.len(), + self.num_spec_layers + ) + } + _ => {} + } + if let Some(ids) = &self.draft_token_ids { + if ids.len() != self.draft_vocab_size as usize { + bail!( + "SPD head draft_token_ids length {} must match draft_vocab_size {}", + ids.len(), + self.draft_vocab_size + ); + } + validate_sorted_unique_indices("draft_token_ids", ids)?; + if ids.iter().any(|id| *id >= self.vocab_size) { + bail!( + "SPD head draft_token_ids must all be less than vocab_size {}", + self.vocab_size + ); + } + } + Ok(()) + } + + fn head_kind(&self) -> &str { + self.head_kind + .as_deref() + .unwrap_or(SPD_HEAD_KIND_FIXED_STAGE_V1) + } + + fn validate_fixed_stage_topology(&self) -> Result<()> { + if let Some(boundaries) = &self.stage_layer_boundaries { + if boundaries.len() != self.num_stages as usize { + bail!( + "SPD head stage_layer_boundaries length {} must match num_stages {}", + boundaries.len(), + self.num_stages + ); + } + validate_sorted_unique_indices("stage_layer_boundaries", boundaries)?; + } + if self.shallow_hidden_layer_indices.len() != self.num_stages as usize { + bail!( + "SPD head shallow_hidden_layer_indices length {} must match num_stages {}", + self.shallow_hidden_layer_indices.len(), + self.num_stages + ); + } + for (stage, indices) in self.shallow_hidden_layer_indices.iter().enumerate() { + validate_sorted_unique_indices( + &format!("shallow_hidden_layer_indices[{stage}]"), + indices, + )?; + } + Ok(()) + } + + fn validate_generic_layer_tap_topology(&self) -> Result<()> { + match self.max_taps { + Some(max_taps) if max_taps > 0 => {} + _ => bail!("generic SPD head max_taps must be greater than zero"), + } + match self.tap_feature_size { + Some(tap_feature_size) if tap_feature_size > 0 => {} + _ => bail!("generic SPD head tap_feature_size must be greater than zero"), + } + if let Some(boundaries) = &self.stage_layer_boundaries { + validate_sorted_unique_indices("stage_layer_boundaries", boundaries)?; + } + if self.shallow_hidden_layer_indices.is_empty() { + bail!("generic SPD head shallow_hidden_layer_indices must include representative taps"); + } + for (row, indices) in self.shallow_hidden_layer_indices.iter().enumerate() { + validate_sorted_unique_indices( + &format!("shallow_hidden_layer_indices[{row}]"), + indices, + )?; + } + Ok(()) + } +} + +fn validate_source_format(source: &SpdHeadSource) -> Result<()> { + match (source.format.as_str(), source.checkpoint_version) { + (TORCH_SPD_HEAD_FORMAT_V10, 10) | (GENERIC_LAYER_TAP_HEAD_FORMAT_V1, 1) => Ok(()), + _ => bail!( + "unsupported SPD head format {} version {}; expected {} version 10 or {} version 1", + source.format, + source.checkpoint_version, + TORCH_SPD_HEAD_FORMAT_V10, + GENERIC_LAYER_TAP_HEAD_FORMAT_V1 + ), + } +} + +fn verify_checkpoint_artifact( + label: &str, + path: &Path, + expected_bytes: u64, + expected_sha256: &str, +) -> Result<()> { + let metadata = + fs::metadata(path).with_context(|| format!("read {label} metadata {}", path.display()))?; + if metadata.len() != expected_bytes { + bail!( + "{label} byte size mismatch for {}: expected {}, got {}", + path.display(), + expected_bytes, + metadata.len() + ); + } + let actual = file_sha256(path)?; + if actual != expected_sha256 { + bail!( + "{label} checksum mismatch for {}: expected {}, got {}", + path.display(), + expected_sha256, + actual + ); + } + Ok(()) +} + +fn validate_safetensors_ranges( + tensors: &BTreeMap, + data_len: u64, +) -> Result<()> { + let mut ranges: Vec<(&str, [u64; 2])> = tensors + .iter() + .map(|(name, tensor)| (name.as_str(), tensor.data_offsets)) + .collect(); + ranges.sort_by_key(|(_, [start, _])| *start); + let mut previous_end = 0; + for (name, [start, end]) in ranges { + if start < previous_end { + bail!("SPD safetensors tensor {name} overlaps a previous tensor range"); + } + if start > previous_end { + bail!("SPD safetensors tensor {name} leaves a gap before its data range"); + } + previous_end = end; + } + if previous_end != data_len { + bail!( + "SPD safetensors tensor data ends at {}, but data section is {} bytes", + previous_end, + data_len + ); + } + Ok(()) +} + +fn tensor_byte_len(dtype: &str, shape: &[u64]) -> Result { + let element_bytes = safetensors_dtype_size(dtype)?; + let elements = shape.iter().try_fold(1_u64, |acc, dimension| { + acc.checked_mul(*dimension) + .context("SPD safetensors tensor shape element count overflow") + })?; + elements + .checked_mul(element_bytes) + .context("SPD safetensors tensor byte length overflow") +} + +fn safetensors_dtype_size(dtype: &str) -> Result { + match dtype { + "BOOL" | "I8" | "U8" => Ok(1), + "F16" | "BF16" | "I16" | "U16" => Ok(2), + "F32" | "I32" | "U32" => Ok(4), + "F64" | "I64" | "U64" => Ok(8), + _ => bail!("unsupported SPD safetensors dtype {dtype}"), + } +} + +fn validate_sorted_unique_indices(label: &str, values: &[u32]) -> Result<()> { + if values.is_empty() { + bail!("SPD head {label} must not be empty"); + } + if values.windows(2).any(|pair| pair[0] >= pair[1]) { + bail!("SPD head {label} must be sorted and unique"); + } + Ok(()) +} + +fn validate_sha256_digest(label: &str, value: &str) -> Result<()> { + if value.len() != 64 || !value.bytes().all(|byte| byte.is_ascii_hexdigit()) { + bail!("{label} must be a 64-character hex digest"); + } + Ok(()) +} + +fn safe_relative_manifest_path(path: &str) -> Result { + let path = Path::new(path); + if path.as_os_str().is_empty() || path.is_absolute() { + bail!("SPD checkpoint path must be a non-empty relative path"); + } + for component in path.components() { + match component { + Component::Normal(_) => {} + _ => bail!( + "SPD checkpoint path must not contain prefix, root, dot, or parent components" + ), + } + } + Ok(path.to_path_buf()) +} + +fn file_sha256(path: &Path) -> Result { + let mut file = fs::File::open(path) + .with_context(|| format!("open SPD checkpoint for hashing {}", path.display()))?; + let mut hasher = Sha256::new(); + std::io::copy(&mut file, &mut hasher) + .with_context(|| format!("hash SPD checkpoint {}", path.display()))?; + Ok(format!("{:x}", hasher.finalize())) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn valid_manifest() -> SpdHeadManifest { + SpdHeadManifest { + schema: SPD_HEAD_MANIFEST_SCHEMA.to_string(), + checkpoint: SpdHeadCheckpoint { + path: "speculation_head_final.pt".to_string(), + sha256: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + .to_string(), + bytes: 4, + }, + serving_checkpoint: None, + source: SpdHeadSource { + format: TORCH_SPD_HEAD_FORMAT_V10.to_string(), + reference_repo: Some("https://example.invalid/spd.git".to_string()), + base_model_path: "Qwen/Qwen3-0.6B".to_string(), + model_type: Some("qwen3".to_string()), + checkpoint_version: 10, + }, + topology: SpdHeadTopology { + hidden_size: 1024, + vocab_size: 10, + draft_vocab_size: 3, + head_kind: None, + num_stages: 2, + stage_layer_boundaries: Some(vec![7, 14]), + num_spec_layers: 1, + max_taps: None, + tap_feature_size: None, + trained_with_use_deepest: true, + shallow_hidden_layer_indices: vec![vec![0, 7, 14], vec![0, 14]], + spec_init_from_base_layers: Some(vec![20]), + draft_token_ids: Some(vec![1, 3, 5]), + }, + } + } + + fn valid_manifest_with_serving_checkpoint() -> SpdHeadManifest { + let mut manifest = valid_manifest(); + manifest.serving_checkpoint = Some(SpdHeadServingCheckpoint { + path: "spd-head.safetensors".to_string(), + sha256: "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb".to_string(), + bytes: 0, + format: SPD_SERVING_CHECKPOINT_FORMAT_SAFETENSORS_V1.to_string(), + tensor_count: 6, + dtype: "F32".to_string(), + }); + manifest + } + + fn valid_generic_layer_tap_manifest() -> SpdHeadManifest { + let mut manifest = valid_manifest_with_serving_checkpoint(); + manifest.source.format = GENERIC_LAYER_TAP_HEAD_FORMAT_V1.to_string(); + manifest.source.checkpoint_version = 1; + manifest.topology.head_kind = Some(SPD_HEAD_KIND_GENERIC_LAYER_TAP_V1.to_string()); + manifest.topology.num_stages = 6; + manifest.topology.stage_layer_boundaries = None; + manifest.topology.num_spec_layers = 2; + manifest.topology.max_taps = Some(8); + manifest.topology.tap_feature_size = Some(2); + manifest.topology.trained_with_use_deepest = false; + manifest.topology.shallow_hidden_layer_indices = + vec![vec![0, 7, 14], vec![0, 4, 10, 14], vec![0, 2, 6, 9, 14]]; + manifest.topology.spec_init_from_base_layers = None; + manifest.serving_checkpoint.as_mut().unwrap().tensor_count = 12; + manifest + } + + fn write_test_safetensors(path: &Path, tensors: &[(&str, &str, &[u64])]) { + let mut header_entries = serde_json::Map::new(); + let mut data = Vec::new(); + for (name, dtype, shape) in tensors { + let start = data.len() as u64; + let bytes = tensor_byte_len(dtype, shape).unwrap(); + data.resize(data.len() + bytes as usize, 0); + let end = data.len() as u64; + header_entries.insert( + (*name).to_string(), + serde_json::json!({ + "dtype": dtype, + "shape": shape, + "data_offsets": [start, end], + }), + ); + } + header_entries.insert( + "__metadata__".to_string(), + serde_json::json!({"format": SPD_SERVING_CHECKPOINT_FORMAT_SAFETENSORS_V1}), + ); + let header = serde_json::to_vec(&serde_json::Value::Object(header_entries)).unwrap(); + let mut bytes = Vec::new(); + bytes.extend_from_slice(&(header.len() as u64).to_le_bytes()); + bytes.extend_from_slice(&header); + bytes.extend_from_slice(&data); + fs::write(path, bytes).unwrap(); + } + + #[test] + fn validates_reference_manifest_shape() { + valid_manifest().validate().unwrap(); + } + + #[test] + fn rejects_draft_vocab_size_mismatch() { + let mut manifest = valid_manifest(); + manifest.topology.draft_token_ids = Some(vec![1, 3]); + let error = manifest.validate().unwrap_err().to_string(); + assert!(error.contains("draft_token_ids length")); + } + + #[test] + fn rejects_stage_layer_boundary_count_mismatch() { + let mut manifest = valid_manifest(); + manifest.topology.stage_layer_boundaries = Some(vec![7]); + let error = manifest.validate().unwrap_err().to_string(); + assert!(error.contains("stage_layer_boundaries length")); + } + + #[test] + fn rejects_unsafe_checkpoint_path() { + let mut manifest = valid_manifest(); + manifest.checkpoint.path = "../speculation_head_final.pt".to_string(); + let error = manifest.validate().unwrap_err().to_string(); + assert!(error.contains("parent components")); + } + + #[test] + fn verifies_checkpoint_checksum_relative_to_manifest() { + let temp = tempfile::tempdir().unwrap(); + let checkpoint = temp.path().join("speculation_head_final.pt"); + fs::write(&checkpoint, b"head").unwrap(); + let sha256 = file_sha256(&checkpoint).unwrap(); + + let mut manifest = valid_manifest(); + manifest.checkpoint.sha256 = sha256; + manifest.checkpoint.bytes = 4; + let manifest_path = temp.path().join("skippy-spd-head.json"); + fs::write(&manifest_path, serde_json::to_vec(&manifest).unwrap()).unwrap(); + + let parsed = SpdHeadManifest::from_path(&manifest_path).unwrap(); + parsed.verify_checkpoint(&manifest_path).unwrap(); + } + + #[test] + fn verifies_serving_checkpoint_checksum_and_shapes() { + let temp = tempfile::tempdir().unwrap(); + let checkpoint = temp.path().join("spd-head.safetensors"); + write_test_safetensors( + &checkpoint, + &[ + ("stage_projs.0.weight", "F32", &[1024, 3072]), + ("stage_projs.1.weight", "F32", &[1024, 2048]), + ("g0_proj.weight", "F32", &[1024, 1024]), + ("lm_head.weight", "F32", &[3, 1024]), + ("spec_layers.0.input_layernorm.weight", "F32", &[1024]), + ( + "spec_layers.0.post_attention_layernorm.weight", + "F32", + &[1024], + ), + ], + ); + let sha256 = file_sha256(&checkpoint).unwrap(); + + let mut manifest = valid_manifest_with_serving_checkpoint(); + let serving_checkpoint = manifest.serving_checkpoint.as_mut().unwrap(); + serving_checkpoint.sha256 = sha256; + serving_checkpoint.bytes = fs::metadata(&checkpoint).unwrap().len(); + let manifest_path = temp.path().join("skippy-spd-head.json"); + fs::write(&manifest_path, serde_json::to_vec(&manifest).unwrap()).unwrap(); + + let parsed = SpdHeadManifest::from_path(&manifest_path).unwrap(); + let index = parsed + .ensure_serving_checkpoint_for_runtime(&manifest_path) + .unwrap(); + assert_eq!(index.tensors.len(), 6); + assert_eq!( + index.metadata.get("format").unwrap(), + SPD_SERVING_CHECKPOINT_FORMAT_SAFETENSORS_V1 + ); + } + + #[test] + fn verifies_generic_layer_tap_serving_checkpoint_shapes() { + let temp = tempfile::tempdir().unwrap(); + let checkpoint = temp.path().join("spd-head.safetensors"); + write_test_safetensors( + &checkpoint, + &[ + ("tap_proj.weight", "F32", &[1024, 1024]), + ("tap_proj.bias", "F32", &[1024]), + ("depth_proj.weight", "F32", &[1024, 2]), + ("depth_proj.bias", "F32", &[1024]), + ("tap_norm.weight", "F32", &[1024]), + ("tap_norm.bias", "F32", &[1024]), + ("output_norm.weight", "F32", &[1024]), + ("output_norm.bias", "F32", &[1024]), + ("draft_heads.0.weight", "F32", &[3, 1024]), + ("draft_heads.0.bias", "F32", &[3]), + ("draft_heads.1.weight", "F32", &[3, 1024]), + ("draft_heads.1.bias", "F32", &[3]), + ], + ); + let sha256 = file_sha256(&checkpoint).unwrap(); + + let mut manifest = valid_generic_layer_tap_manifest(); + let serving_checkpoint = manifest.serving_checkpoint.as_mut().unwrap(); + serving_checkpoint.sha256 = sha256; + serving_checkpoint.bytes = fs::metadata(&checkpoint).unwrap().len(); + let manifest_path = temp.path().join("skippy-spd-head.json"); + fs::write(&manifest_path, serde_json::to_vec(&manifest).unwrap()).unwrap(); + + let parsed = SpdHeadManifest::from_path(&manifest_path).unwrap(); + let index = parsed + .ensure_serving_checkpoint_for_runtime(&manifest_path) + .unwrap(); + assert_eq!(index.tensors.len(), 12); + parsed + .ensure_runtime_compatible(&SpdHeadRuntimeProfile { + base_model_path: Some("Qwen/Qwen3-0.6B"), + hidden_size: 1024, + vocab_size: 10, + num_stages: 3, + }) + .unwrap(); + } + + #[test] + fn rejects_serving_checkpoint_shape_mismatch() { + let temp = tempfile::tempdir().unwrap(); + let checkpoint = temp.path().join("spd-head.safetensors"); + write_test_safetensors( + &checkpoint, + &[ + ("stage_projs.0.weight", "F32", &[1024, 2048]), + ("stage_projs.1.weight", "F32", &[1024, 2048]), + ("g0_proj.weight", "F32", &[1024, 1024]), + ("lm_head.weight", "F32", &[3, 1024]), + ("spec_layers.0.input_layernorm.weight", "F32", &[1024]), + ( + "spec_layers.0.post_attention_layernorm.weight", + "F32", + &[1024], + ), + ], + ); + let sha256 = file_sha256(&checkpoint).unwrap(); + + let mut manifest = valid_manifest_with_serving_checkpoint(); + let serving_checkpoint = manifest.serving_checkpoint.as_mut().unwrap(); + serving_checkpoint.sha256 = sha256; + serving_checkpoint.bytes = fs::metadata(&checkpoint).unwrap().len(); + let manifest_path = temp.path().join("skippy-spd-head.json"); + fs::write(&manifest_path, serde_json::to_vec(&manifest).unwrap()).unwrap(); + + let parsed = SpdHeadManifest::from_path(&manifest_path).unwrap(); + let error = parsed + .ensure_serving_checkpoint_for_runtime(&manifest_path) + .unwrap_err() + .to_string(); + assert!(error.contains("stage_projs.0.weight shape mismatch")); + } + + #[test] + fn rejects_safetensors_byte_length_mismatch() { + let header = serde_json::json!({ + "bad.weight": { + "dtype": "F32", + "shape": [2_u64, 2_u64], + "data_offsets": [0_u64, 12_u64], + } + }); + let header = serde_json::to_vec(&header).unwrap(); + let mut bytes = Vec::new(); + bytes.extend_from_slice(&(header.len() as u64).to_le_bytes()); + bytes.extend_from_slice(&header); + bytes.extend_from_slice(&[0_u8; 12]); + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("bad.safetensors"); + fs::write(&path, bytes).unwrap(); + + let error = SpdSafetensorsIndex::from_path(&path) + .unwrap_err() + .to_string(); + assert!(error.contains("byte length mismatch")); + } + + #[test] + fn validates_external_manifest_when_skippy_spd_manifest_is_set() { + let Ok(manifest_path) = std::env::var("SKIPPY_SPD_MANIFEST") else { + return; + }; + let manifest_path = PathBuf::from(manifest_path); + let manifest = SpdHeadManifest::from_path(&manifest_path).unwrap(); + let index = manifest + .ensure_serving_checkpoint_for_runtime(&manifest_path) + .unwrap(); + if manifest.topology.head_kind() == SPD_HEAD_KIND_GENERIC_LAYER_TAP_V1 { + assert!(index.tensors.contains_key("tap_proj.weight")); + } else { + assert!(index.tensors.contains_key("lm_head.weight")); + } + } + + #[test] + fn checks_runtime_compatibility_profile() { + let manifest = valid_manifest(); + manifest + .ensure_runtime_compatible(&SpdHeadRuntimeProfile { + base_model_path: Some("Qwen/Qwen3-0.6B"), + hidden_size: 1024, + vocab_size: 10, + num_stages: 2, + }) + .unwrap(); + } + + #[test] + fn rejects_runtime_profile_mismatch() { + let manifest = valid_manifest(); + let error = manifest + .ensure_runtime_compatible(&SpdHeadRuntimeProfile { + base_model_path: Some("Qwen/Qwen3-1.7B"), + hidden_size: 1024, + vocab_size: 10, + num_stages: 2, + }) + .unwrap_err() + .to_string(); + assert!(error.contains("trained for base model")); + } +} diff --git a/crates/skippy-server/src/binary_transport.rs b/crates/skippy-server/src/binary_transport.rs index 37baa73395..211593a06d 100644 --- a/crates/skippy-server/src/binary_transport.rs +++ b/crates/skippy-server/src/binary_transport.rs @@ -1,7 +1,8 @@ use std::{ collections::{BTreeMap, VecDeque}, + env, future::Future, - io, + io::{self, Write}, net::{IpAddr, SocketAddr, TcpListener, TcpStream, ToSocketAddrs}, sync::{ Arc, Mutex, @@ -29,15 +30,19 @@ use skippy_protocol::{ StageReply, StageReplyStats, StageSamplingConfig, StageStateHeader, StageWireMessage, WireActivationDType, WireMessageKind, WireReplyKind, activation_frame_flags_from_state_flags, read_stage_message, recv_reply, send_ready, - send_reply_ack, send_reply_ack_with_stats, state_flags, + send_reply_ack, send_reply_ack_with_stats, send_reply_predicted_tokens_with_stats, + send_reply_predicted_with_tokens_and_stats, state_flags, }, }; use skippy_runtime::{ - ActivationDesc, ActivationFrame, LogitBias, MAX_LOGIT_BIAS, RuntimeActivationDType, - RuntimeActivationLayout, SamplingConfig, + ActivationDesc, ActivationFrame, LogitBias, MAX_LOGIT_BIAS, NativeMtpDraft, + RuntimeActivationDType, RuntimeActivationLayout, SamplingConfig, }; use socket2::{Domain, Protocol, SockAddr, Socket, Type}; +const AUTO_ALIGN_SESSION_ENV: &str = "SKIPPY_STAGE_AUTO_ALIGN_SESSION"; + +mod decode_batcher; pub(crate) mod direct_return; pub(crate) mod forwarding; mod kv_eviction; @@ -45,14 +50,11 @@ mod options; mod socket; mod wire; +pub(crate) use self::decode_batcher::DecodeFrameBatcher; pub use self::direct_return::PredictionReturnHub; -pub use self::direct_return::PredictionReturnListener; -pub(crate) use self::direct_return::PredictionReturnReceiver; pub(crate) use self::forwarding::{forwarded_stage_message, forwarded_stage_message_timed}; -#[cfg(test)] -use self::kv_eviction::BinaryProactiveEviction; use self::kv_eviction::{ - BinaryProactiveEvictionPlan, binary_proactive_eviction_plan, + BinaryProactiveEviction, BinaryProactiveEvictionPlan, binary_proactive_eviction_plan, evict_binary_resident_prefix_for_decode, }; pub use self::options::{BinaryStageOptions, EmbeddedOpenAiStageOptions, parse_wire_dtype}; @@ -61,6 +63,7 @@ pub use self::wire::WireCondition; pub(crate) use self::wire::write_stage_message_conditioned; static BINARY_SESSION_COUNTER: AtomicU64 = AtomicU64::new(1); +const NATIVE_MTP_ENABLED_ENV: &str = "SKIPPY_NATIVE_MTP_ENABLED"; #[derive(Default)] struct BinaryKvLookupResult { @@ -203,6 +206,7 @@ fn run_binary_stage(options: BinaryStageOptions, shutdown: Arc) -> R ); telemetry.emit("stage.binary_server_start", lifecycle_attrs(&config)); let runtime = load_runtime(&config)?.context("binary stage server requires model_path")?; + let decode_frame_batcher = DecodeFrameBatcher::new(runtime.clone(), max_inflight); if max_inflight > 0 { let timer = Instant::now(); let sessions = runtime @@ -231,7 +235,7 @@ fn run_binary_stage(options: BinaryStageOptions, shutdown: Arc) -> R telemetry.emit("stage.binary_runtime_prewarm", attrs); } let kv = KvStageIntegration::from_config(&config)?.map(Arc::new); - let prediction_returns = Arc::new(PredictionReturnHub::default()); + let prediction_returns = Arc::new(PredictionReturnHub); let listener = TcpListener::bind(bind_addr)?; listener.set_nonblocking(true)?; if let Some(openai_options) = openai { @@ -241,7 +245,6 @@ fn run_binary_stage(options: BinaryStageOptions, shutdown: Arc) -> R let openai_config = config.clone(); let openai_runtime = runtime.clone(); let openai_telemetry = telemetry.clone(); - let openai_prediction_returns = prediction_returns.clone(); tokio::spawn(async move { if let Err(error) = frontend::serve_embedded_openai(EmbeddedOpenAiArgs { bind_addr: openai_options.bind_addr, @@ -266,7 +269,6 @@ fn run_binary_stage(options: BinaryStageOptions, shutdown: Arc) -> R reply_credit_limit, downstream_connect_timeout_secs, downstream_wire_condition, - prediction_returns: Some(openai_prediction_returns), telemetry: openai_telemetry, hook_policy: None, openai_guardrails: Some(frontend::OpenAiGuardrailsConfig::disabled_for_skippy()), @@ -298,15 +300,29 @@ fn run_binary_stage(options: BinaryStageOptions, shutdown: Arc) -> R }; prepare_binary_stage_connection(&upstream)?; let peer_addr = upstream.peer_addr().ok(); + eprintln!( + "binary accepted connection: stage_id={} peer={peer_addr:?}", + config.stage_id + ); let config = config.clone(); let topology = topology.clone(); let runtime = runtime.clone(); + let decode_frame_batcher = decode_frame_batcher.clone(); let kv = kv.clone(); let telemetry = telemetry.clone(); let prediction_returns = prediction_returns.clone(); thread::spawn(move || { let connection_result = (|| -> Result<()> { + eprintln!( + "binary sending ready: stage_id={} peer={peer_addr:?}", + config.stage_id + ); send_ready(&mut upstream).context("failed to send binary ready")?; + upstream.flush().ok(); + eprintln!( + "binary sent ready: stage_id={} peer={peer_addr:?}", + config.stage_id + ); let first_message = match read_stage_message(&mut upstream, activation_width) { Ok(message) => message, Err(error) if error.kind() == io::ErrorKind::UnexpectedEof => return Ok(()), @@ -321,6 +337,7 @@ fn run_binary_stage(options: BinaryStageOptions, shutdown: Arc) -> R &config, topology.as_ref(), &runtime, + &decode_frame_batcher, kv.as_ref(), &telemetry, &mut upstream, @@ -331,7 +348,6 @@ fn run_binary_stage(options: BinaryStageOptions, shutdown: Arc) -> R reply_credit_limit, async_prefill_forward, downstream_wire_condition, - downstream_connect_timeout_secs, first_message, ) })() @@ -363,6 +379,7 @@ fn handle_binary_connection( config: &StageConfig, topology: Option<&StageTopology>, runtime: &Arc>, + decode_frame_batcher: &DecodeFrameBatcher, kv: Option<&Arc>, telemetry: &Telemetry, upstream: &mut TcpStream, @@ -373,7 +390,6 @@ fn handle_binary_connection( reply_credit_limit: Option, async_prefill_forward: bool, downstream_wire_condition: WireCondition, - downstream_connect_timeout_secs: u64, first_message: StageWireMessage, ) -> Result<()> { if let Some(downstream) = downstream.as_mut() { @@ -388,7 +404,6 @@ fn handle_binary_connection( let mut pending_reply_stats = StageReplyStats::default(); let mut request_summary = BinaryRequestSummary::default(); let mut accumulated_prefill_tokens: BTreeMap> = BTreeMap::new(); - let mut prediction_return_streams: BTreeMap<(u64, u64), TcpStream> = BTreeMap::new(); let mut next_message = Some(first_message); let mut async_forwarder = if async_prefill_forward { downstream @@ -435,6 +450,10 @@ fn handle_binary_connection( json!(recv_end_unix_nanos), ); recv_attrs.insert("llama_stage.recv_read_ms".to_string(), json!(recv_read_ms)); + recv_attrs.insert( + "skippy.upstream_message_wait_ms".to_string(), + json!(recv_read_ms), + ); recv_attrs.insert( "llama_stage.source_stage_index".to_string(), json!(message.state.source_stage_index), @@ -667,15 +686,6 @@ fn handle_binary_connection( ) .context("configure binary stage generation")?; } - let stream = direct_return::open_prediction_return_stream( - config, - topology, - message.request_id, - message.session_id, - wire_dtype, - downstream_connect_timeout_secs, - )?; - prediction_return_streams.insert((message.request_id, message.session_id), stream); } send_reply_ack_with_stats(&mut *upstream, generation_stats) .context("generation config ack")?; @@ -712,8 +722,7 @@ fn handle_binary_connection( activation_width, control_started, control_stats, - &mut prediction_return_streams, - downstream_connect_timeout_secs, + upstream, ) .context("handle restore-prefill-decode control")?; continue; @@ -793,6 +802,40 @@ fn handle_binary_connection( } let token_ids = token_sideband_or_fill(&message)?; + let mut session_auto_align_count = 0usize; + let mut session_auto_align_ms = 0.0; + let mut session_auto_align_trimmed_tokens = 0u64; + if binary_auto_align_session_enabled() + && message_allows_session_auto_align(&message) + && let Some(target_token_count) = message_pos_start_as_token_count(&message) + { + let align_started = Instant::now(); + let align = { + let mut runtime = runtime.lock().expect("runtime lock poisoned"); + runtime + .align_session_to_token_count_if_ahead(&session_key, target_token_count) + .context("auto-align binary stage session")? + }; + if let Some(align) = align { + let align_ms = elapsed_ms(align_started); + session_auto_align_count = 1; + session_auto_align_ms = align_ms; + session_auto_align_trimmed_tokens = align + .before_token_count + .saturating_sub(align.after_token_count); + let mut attrs = binary_message_attrs(config, session_id, &message); + attrs.insert( + "llama_stage.session_auto_align_before_tokens".to_string(), + json!(align.before_token_count), + ); + attrs.insert( + "llama_stage.session_auto_align_after_tokens".to_string(), + json!(align.after_token_count), + ); + attrs.insert("llama_stage.elapsed_ms".to_string(), json!(align_ms)); + telemetry.emit_debug("stage.binary_session_auto_align", attrs); + } + } if message.kind.is_prefill() { accumulate_prefill_tokens( &mut accumulated_prefill_tokens, @@ -832,6 +875,8 @@ fn handle_binary_connection( let mut runtime_lock_acquires = 0usize; let mut runtime_sessions_before = None; let mut runtime_sessions_after = None; + let mut decode_batch_size = 1usize; + let mut decode_batch_wait_ms = 0.0; let input_activation_bytes = message.activation.len(); let mut proactive_eviction = None; let (predicted_token, predicted_tokens, output, compute_ms) = if restored_prefill { @@ -882,7 +927,24 @@ fn handle_binary_connection( } compute_start_unix_nanos = now_unix_nanos() as u64; let compute_started = Instant::now(); - let result = { + let use_decode_frame_batch = + is_decode_frame_batch_candidate(config, &message, executable_token_ids); + let result = if use_decode_frame_batch { + let token_id = executable_token_ids + .first() + .copied() + .unwrap_or(message.state.current_token); + let sampling = runtime_sampling_config(message.sampling.as_ref()); + let outcome = decode_frame_batcher + .decode(&session_key, token_id, sampling.as_ref(), input) + .context("execute batched binary decode frame")?; + runtime_lock_wait_ms = outcome.runtime_lock_wait_ms; + runtime_lock_hold_ms = outcome.runtime_lock_hold_ms; + runtime_lock_acquires = 1; + decode_batch_size = outcome.batch_size; + decode_batch_wait_ms = outcome.batch_wait_ms; + (outcome.predicted, Vec::new(), outcome.output) + } else { let lock_started = Instant::now(); let mut runtime = runtime.lock().expect("runtime lock poisoned"); runtime_lock_wait_ms = elapsed_ms(lock_started); @@ -947,6 +1009,14 @@ fn handle_binary_connection( "llama_stage.runtime_lock_acquires".to_string(), json!(runtime_lock_acquires), ); + decode_attrs.insert( + "llama_stage.decode_batch_size".to_string(), + json!(decode_batch_size), + ); + decode_attrs.insert( + "llama_stage.decode_batch_wait_ms".to_string(), + json!(decode_batch_wait_ms), + ); if let Some(stats) = runtime_sessions_before.as_ref() { insert_runtime_session_stats( &mut decode_attrs, @@ -984,7 +1054,7 @@ fn handle_binary_connection( ); } if let Some(eviction) = proactive_eviction { - telemetry.emit("stage.binary_kv_record_decision", eviction.attrs()); + emit_binary_proactive_eviction(telemetry, &eviction); } if message.kind.is_prefill() && !restored_prefill { @@ -1181,6 +1251,68 @@ fn handle_binary_connection( pending_prefill_replies -= 1; deferred_prefill_replies_drained += 1; } + let wait_start_unix_nanos = now_unix_nanos() as u64; + downstream_wait_start_unix_nanos.get_or_insert(wait_start_unix_nanos); + let wait_started = Instant::now(); + let reply = recv_reply(&mut *downstream).context("downstream predicted reply")?; + downstream_wait_end_unix_nanos = Some(now_unix_nanos() as u64); + downstream_wait_ms += elapsed_ms(wait_started); + let expected = predicted_reply_kind(message.kind); + if reply.kind != expected { + bail!( + "expected downstream {expected:?} for {:?}, got {:?}", + message.kind, + reply.kind + ); + } + record_prefill_edge_transport( + &mut message_reply_stats, + config, + &message, + forward_write_ms, + downstream_wait_ms, + forward_activation_bytes, + ); + message_reply_stats.merge(pending_reply_stats); + pending_reply_stats = StageReplyStats::default(); + message_reply_stats.merge(reply.stats); + record_verify_span_timing( + &mut message_reply_stats, + &message, + compute_ms, + forward_write_ms, + downstream_wait_ms, + ); + let reply_start_unix_nanos = now_unix_nanos() as u64; + upstream_reply_start_unix_nanos.get_or_insert(reply_start_unix_nanos); + let reply_started = Instant::now(); + let predicted_token_count = + predicted_reply_token_count(message.kind, &reply.predicted_tokens); + send_stage_reply( + &mut *upstream, + StageReply { + stats: message_reply_stats, + ..reply + }, + ) + .context("relay predicted reply")?; + upstream_reply_end_unix_nanos = Some(now_unix_nanos() as u64); + let reply_write_ms = elapsed_ms(reply_started); + upstream_reply_ms += reply_write_ms; + emit_upstream_reply_write_span( + telemetry, + config, + session_id, + &message, + UpstreamReplyWriteSpan { + reply_kind: expected, + predicted_token_count, + start_unix_nanos: reply_start_unix_nanos, + end_unix_nanos: upstream_reply_end_unix_nanos + .unwrap_or(reply_start_unix_nanos), + write_ms: reply_write_ms, + }, + ); } else if max_deferred_prefill_replies == 0 { let wait_start_unix_nanos = now_unix_nanos() as u64; downstream_wait_start_unix_nanos.get_or_insert(wait_start_unix_nanos); @@ -1289,23 +1421,21 @@ fn handle_binary_connection( let predicted_token_count = if message.kind == WireMessageKind::VerifySpan { predicted_tokens.len() } else { - 1 + predicted_tokens.len().max(1) }; - let return_stream = prediction_return_streams - .get_mut(&(message.request_id, message.session_id)) - .ok_or_else(|| anyhow!("missing direct prediction return stream"))?; let reply_start_unix_nanos = now_unix_nanos() as u64; upstream_reply_start_unix_nanos.get_or_insert(reply_start_unix_nanos); let reply_started = Instant::now(); - direct_return::send_direct_prediction_return( - return_stream, + send_stage_reply( + &mut *upstream, StageReply { kind: reply_kind, predicted: predicted_token, predicted_tokens, stats: message_reply_stats, }, - )?; + ) + .context("send predicted reply")?; upstream_reply_end_unix_nanos = Some(now_unix_nanos() as u64); let reply_write_ms = elapsed_ms(reply_started); upstream_reply_ms += reply_write_ms; @@ -1365,6 +1495,30 @@ fn handle_binary_connection( let message_end_unix_nanos = now_unix_nanos() as u64; let message_elapsed_ms = elapsed_ms(message_started); + let verify_span_pre_compute_ms = if message.kind == WireMessageKind::VerifySpan { + nanos_delta_ms(message_start_unix_nanos, compute_start_unix_nanos) + } else { + 0.0 + }; + let verify_span_post_compute_ms = if message.kind == WireMessageKind::VerifySpan { + nanos_delta_ms(compute_end_unix_nanos, message_end_unix_nanos) + } else { + 0.0 + }; + let verify_span_pre_reply_ms = if message.kind == WireMessageKind::VerifySpan { + upstream_reply_start_unix_nanos + .map(|reply_start| nanos_delta_ms(compute_end_unix_nanos, reply_start)) + .unwrap_or(0.0) + } else { + 0.0 + }; + let verify_span_after_reply_ms = if message.kind == WireMessageKind::VerifySpan { + upstream_reply_end_unix_nanos + .map(|reply_end| nanos_delta_ms(reply_end, message_end_unix_nanos)) + .unwrap_or(0.0) + } else { + 0.0 + }; request_summary.observe(BinaryMessageObservation { config, message: &message, @@ -1384,6 +1538,14 @@ fn handle_binary_connection( pending_prefill_replies_after: pending_prefill_replies, credit_wait_count, deferred_prefill_replies_drained, + session_auto_align_count, + session_auto_align_ms, + session_auto_align_trimmed_tokens, + verify_span_pre_compute_ms, + verify_span_post_compute_ms, + verify_span_pre_reply_ms, + verify_span_after_reply_ms, + upstream_message_wait_ms: recv_read_ms, }); if telemetry.is_debug_enabled() { @@ -1405,6 +1567,11 @@ fn handle_binary_connection( json!(compute_end_unix_nanos), ); timing_attrs.insert("llama_stage.compute_ms".to_string(), json!(compute_ms)); + timing_attrs.insert("llama_stage.recv_read_ms".to_string(), json!(recv_read_ms)); + timing_attrs.insert( + "skippy.upstream_message_wait_ms".to_string(), + json!(recv_read_ms), + ); timing_attrs.insert( "llama_stage.input_activation_decode_ms".to_string(), json!(input_activation_decode_ms), @@ -1539,6 +1706,57 @@ fn insert_optional_unix_nanos(attrs: &mut BTreeMap, key: &str, va } } +fn native_mtp_prediction_tokens(predicted: i32, draft: Option) -> Vec { + let Some(draft) = draft else { + return vec![predicted]; + }; + vec![ + predicted, + draft.token_id, + draft.proposal_compute_us.clamp(0, i64::from(i32::MAX)) as i32, + ] +} + +fn native_mtp_enabled() -> bool { + native_mtp_enabled_from(env::var(NATIVE_MTP_ENABLED_ENV).ok().as_deref()) +} + +fn binary_auto_align_session_enabled() -> bool { + truthy_env(env::var(AUTO_ALIGN_SESSION_ENV).ok().as_deref()) +} + +fn truthy_env(value: Option<&str>) -> bool { + matches!( + value.map(|value| value.trim().to_ascii_lowercase()), + Some(value) + if matches!( + value.as_str(), + "1" | "true" | "on" | "enable" | "enabled" | "yes" + ) + ) +} + +fn message_allows_session_auto_align(message: &StageWireMessage) -> bool { + matches!( + message.kind, + WireMessageKind::DecodeEmbd + | WireMessageKind::DecodeReadout + | WireMessageKind::DecodeLightCtx + | WireMessageKind::VerifySpan + ) +} + +fn message_pos_start_as_token_count(message: &StageWireMessage) -> Option { + u64::try_from(message.pos_start).ok() +} + +fn native_mtp_enabled_from(value: Option<&str>) -> bool { + !matches!( + value.map(str::trim).map(str::to_ascii_lowercase).as_deref(), + Some("0" | "false" | "off" | "disable" | "disabled" | "no") + ) +} + pub(crate) fn stage_output_activation_capacity( config: &StageConfig, token_count: i32, @@ -1788,6 +2006,10 @@ fn elapsed_ms(started: Instant) -> f64 { started.elapsed().as_secs_f64() * 1000.0 } +fn nanos_delta_ms(start_unix_nanos: u64, end_unix_nanos: u64) -> f64 { + end_unix_nanos.saturating_sub(start_unix_nanos) as f64 / 1_000_000.0 +} + fn elapsed_us(started: Instant) -> i64 { let micros = started.elapsed().as_micros(); micros.min(i64::MAX as u128) as i64 @@ -2072,8 +2294,7 @@ fn handle_binary_restore_prefill_decode_control( activation_width: i32, control_started: Instant, mut control_stats: StageReplyStats, - prediction_return_streams: &mut BTreeMap<(u64, u64), TcpStream>, - downstream_connect_timeout_secs: u64, + upstream: &mut TcpStream, ) -> Result<()> { let (prefix_tokens, current_token) = restore_decode_sideband(&message)?; let local = maybe_prefix_cache_control( @@ -2094,19 +2315,16 @@ fn handle_binary_restore_prefill_decode_control( json!(elapsed_ms(control_started)), ); telemetry.emit_debug("stage.binary_prefix_cache_decode_control", attrs); - send_one_off_direct_return( - config, - topology, - &message, - wire_dtype, - downstream_connect_timeout_secs, + send_stage_reply( + upstream, StageReply { kind: WireReplyKind::Ack, predicted: 0, predicted_tokens: Vec::new(), stats: control_stats, }, - )?; + ) + .context("send restore-decode miss ACK")?; return Ok(()); } @@ -2157,10 +2375,7 @@ fn handle_binary_restore_prefill_decode_control( ) }; let compute_ms = elapsed_ms(compute_started); - telemetry.emit( - "stage.binary_kv_record_decision", - proactive_eviction.attrs(), - ); + emit_binary_proactive_eviction(telemetry, &proactive_eviction); if let Some(downstream) = downstream { let forwarded = @@ -2198,6 +2413,23 @@ fn handle_binary_restore_prefill_decode_control( json!(forwarded.activation_encode_ms), ); telemetry.emit_debug("stage.binary_prefix_cache_decode_control", attrs); + let downstream_reply = + recv_reply(&mut *downstream).context("restore-decode downstream reply")?; + if downstream_reply.kind != WireReplyKind::PredictedToken { + bail!( + "restore-decode expected downstream PredictedToken, got {:?}", + downstream_reply.kind + ); + } + control_stats.merge(downstream_reply.stats); + send_stage_reply( + upstream, + StageReply { + stats: control_stats, + ..downstream_reply + }, + ) + .context("relay restore-decode predicted reply")?; return Ok(()); } @@ -2231,38 +2463,68 @@ fn handle_binary_restore_prefill_decode_control( ); proactive_eviction.insert_attrs(&mut attrs); telemetry.emit_debug("stage.binary_prefix_cache_decode_control", attrs); - let return_stream = prediction_return_streams - .get_mut(&(message.request_id, message.session_id)) - .ok_or_else(|| anyhow!("missing direct prediction return stream"))?; - direct_return::send_direct_prediction_return( - return_stream, + send_stage_reply( + upstream, StageReply { kind: WireReplyKind::PredictedToken, predicted: predicted_token, predicted_tokens: vec![predicted_token], stats: control_stats, }, - )?; + ) + .context("send restore-decode predicted reply")?; Ok(()) } -fn send_one_off_direct_return( - config: &StageConfig, - topology: Option<&StageTopology>, - message: &StageWireMessage, - wire_dtype: WireActivationDType, - downstream_connect_timeout_secs: u64, - reply: StageReply, -) -> Result<()> { - let mut stream = direct_return::open_prediction_return_stream( - config, - topology, - message.request_id, - message.session_id, - wire_dtype, - downstream_connect_timeout_secs, - )?; - direct_return::send_direct_prediction_return(&mut stream, reply) +fn predicted_reply_kind(kind: WireMessageKind) -> WireReplyKind { + if kind == WireMessageKind::VerifySpan { + WireReplyKind::PredictedTokens + } else { + WireReplyKind::PredictedToken + } +} + +fn predicted_reply_token_count(kind: WireMessageKind, predicted_tokens: &[i32]) -> usize { + if kind == WireMessageKind::VerifySpan { + predicted_tokens.len() + } else { + predicted_tokens.len().max(1) + } +} + +fn send_stage_reply(stream: &mut TcpStream, reply: StageReply) -> Result<()> { + match reply.kind { + WireReplyKind::Ack => { + send_reply_ack_with_stats(stream, reply.stats).context("send stage ACK reply") + } + WireReplyKind::PredictedToken => send_reply_predicted_with_tokens_and_stats( + stream, + reply.predicted, + &stage_reply_prediction_tokens(&reply), + reply.stats, + ) + .context("send stage predicted-token reply"), + WireReplyKind::PredictedTokens => { + send_reply_predicted_tokens_with_stats(stream, &reply.predicted_tokens, reply.stats) + .context("send stage predicted-tokens reply") + } + } +} + +fn emit_binary_proactive_eviction(telemetry: &Telemetry, eviction: &BinaryProactiveEviction) { + if eviction.should_emit_summary() { + telemetry.emit("stage.binary_kv_record_decision", eviction.attrs()); + } else { + telemetry.emit_debug("stage.binary_kv_record_decision", eviction.attrs()); + } +} + +fn stage_reply_prediction_tokens(reply: &StageReply) -> Vec { + if reply.predicted_tokens.is_empty() { + vec![reply.predicted] + } else { + reply.predicted_tokens.clone() + } } fn restore_decode_sideband(message: &StageWireMessage) -> Result<(&[i32], i32)> { @@ -2890,6 +3152,24 @@ struct BinaryRequestSummary { prefill_credit_wait_count: usize, prefill_deferred_replies_drained: usize, prefill_pending_replies_max: usize, + session_auto_align_count: usize, + session_auto_align_ms: f64, + session_auto_align_trimmed_tokens: u64, + verify_span_count: usize, + verify_span_session_auto_align_count: usize, + verify_span_session_auto_align_ms: f64, + verify_span_session_auto_align_trimmed_tokens: u64, + verify_span_token_count: u64, + verify_span_max_tokens: u64, + verify_span_compute_ms: f64, + verify_span_input_activation_decode_ms: f64, + verify_span_runtime_lock_hold_ms: f64, + verify_span_upstream_reply_ms: f64, + verify_span_pre_compute_ms: f64, + verify_span_post_compute_ms: f64, + verify_span_pre_reply_ms: f64, + verify_span_after_reply_ms: f64, + verify_span_upstream_message_wait_ms: f64, reply_stats: StageReplyStats, } @@ -2912,6 +3192,14 @@ struct BinaryMessageObservation<'a> { pending_prefill_replies_after: usize, credit_wait_count: usize, deferred_prefill_replies_drained: usize, + session_auto_align_count: usize, + session_auto_align_ms: f64, + session_auto_align_trimmed_tokens: u64, + verify_span_pre_compute_ms: f64, + verify_span_post_compute_ms: f64, + verify_span_pre_reply_ms: f64, + verify_span_after_reply_ms: f64, + upstream_message_wait_ms: f64, } #[derive(Clone, Copy)] @@ -3060,6 +3348,31 @@ impl BinaryRequestSummary { .prefill_pending_replies_max .max(observation.pending_prefill_replies_before) .max(observation.pending_prefill_replies_after); + self.session_auto_align_count += observation.session_auto_align_count; + self.session_auto_align_ms += observation.session_auto_align_ms; + self.session_auto_align_trimmed_tokens = self + .session_auto_align_trimmed_tokens + .saturating_add(observation.session_auto_align_trimmed_tokens); + if message.kind == WireMessageKind::VerifySpan { + let token_count = message.token_count.max(0) as u64; + self.verify_span_count += 1; + self.verify_span_token_count = self.verify_span_token_count.saturating_add(token_count); + self.verify_span_max_tokens = self.verify_span_max_tokens.max(token_count); + self.verify_span_session_auto_align_count += observation.session_auto_align_count; + self.verify_span_session_auto_align_ms += observation.session_auto_align_ms; + self.verify_span_session_auto_align_trimmed_tokens = self + .verify_span_session_auto_align_trimmed_tokens + .saturating_add(observation.session_auto_align_trimmed_tokens); + self.verify_span_compute_ms += observation.compute_ms; + self.verify_span_input_activation_decode_ms += observation.input_activation_decode_ms; + self.verify_span_runtime_lock_hold_ms += observation.runtime_lock_hold_ms; + self.verify_span_upstream_reply_ms += observation.upstream_reply_ms; + self.verify_span_pre_compute_ms += observation.verify_span_pre_compute_ms; + self.verify_span_post_compute_ms += observation.verify_span_post_compute_ms; + self.verify_span_pre_reply_ms += observation.verify_span_pre_reply_ms; + self.verify_span_after_reply_ms += observation.verify_span_after_reply_ms; + self.verify_span_upstream_message_wait_ms += observation.upstream_message_wait_ms; + } self.reply_stats.merge(observation.reply_stats); } @@ -3166,6 +3479,136 @@ impl BinaryRequestSummary { "skippy.prefill_pending_replies_max".to_string(), json!(self.prefill_pending_replies_max), ); + attrs.insert( + "skippy.session_auto_align_count".to_string(), + json!(self.session_auto_align_count), + ); + attrs.insert( + "skippy.session_auto_align_ms".to_string(), + json!(self.session_auto_align_ms), + ); + attrs.insert( + "skippy.session_auto_align_trimmed_tokens".to_string(), + json!(self.session_auto_align_trimmed_tokens), + ); + if self.session_auto_align_count > 0 { + attrs.insert( + "skippy.session_auto_align_ms_avg".to_string(), + json!(self.session_auto_align_ms / self.session_auto_align_count as f64), + ); + } + attrs.insert( + "skippy.verify_span_count".to_string(), + json!(self.verify_span_count), + ); + attrs.insert( + "skippy.verify_span_token_count".to_string(), + json!(self.verify_span_token_count), + ); + attrs.insert( + "skippy.verify_span_max_tokens".to_string(), + json!(self.verify_span_max_tokens), + ); + attrs.insert( + "skippy.verify_span_session_auto_align_count".to_string(), + json!(self.verify_span_session_auto_align_count), + ); + attrs.insert( + "skippy.verify_span_session_auto_align_ms".to_string(), + json!(self.verify_span_session_auto_align_ms), + ); + attrs.insert( + "skippy.verify_span_session_auto_align_trimmed_tokens".to_string(), + json!(self.verify_span_session_auto_align_trimmed_tokens), + ); + if self.verify_span_session_auto_align_count > 0 { + attrs.insert( + "skippy.verify_span_session_auto_align_ms_avg".to_string(), + json!( + self.verify_span_session_auto_align_ms + / self.verify_span_session_auto_align_count as f64 + ), + ); + } + attrs.insert( + "skippy.verify_span_pre_compute_ms".to_string(), + json!(self.verify_span_pre_compute_ms), + ); + attrs.insert( + "skippy.verify_span_compute_ms".to_string(), + json!(self.verify_span_compute_ms), + ); + attrs.insert( + "skippy.verify_span_input_activation_decode_ms".to_string(), + json!(self.verify_span_input_activation_decode_ms), + ); + attrs.insert( + "skippy.verify_span_runtime_lock_hold_ms".to_string(), + json!(self.verify_span_runtime_lock_hold_ms), + ); + attrs.insert( + "skippy.verify_span_upstream_reply_ms".to_string(), + json!(self.verify_span_upstream_reply_ms), + ); + attrs.insert( + "skippy.verify_span_post_compute_ms".to_string(), + json!(self.verify_span_post_compute_ms), + ); + attrs.insert( + "skippy.verify_span_pre_reply_ms".to_string(), + json!(self.verify_span_pre_reply_ms), + ); + attrs.insert( + "skippy.verify_span_after_reply_ms".to_string(), + json!(self.verify_span_after_reply_ms), + ); + attrs.insert( + "skippy.verify_span_upstream_message_wait_ms".to_string(), + json!(self.verify_span_upstream_message_wait_ms), + ); + if self.verify_span_count > 0 { + let verify_span_count = self.verify_span_count as f64; + attrs.insert( + "skippy.verify_span_pre_compute_ms_avg".to_string(), + json!(self.verify_span_pre_compute_ms / verify_span_count), + ); + attrs.insert( + "skippy.verify_span_compute_ms_avg".to_string(), + json!(self.verify_span_compute_ms / verify_span_count), + ); + attrs.insert( + "skippy.verify_span_input_activation_decode_ms_avg".to_string(), + json!(self.verify_span_input_activation_decode_ms / verify_span_count), + ); + attrs.insert( + "skippy.verify_span_runtime_lock_hold_ms_avg".to_string(), + json!(self.verify_span_runtime_lock_hold_ms / verify_span_count), + ); + attrs.insert( + "skippy.verify_span_upstream_reply_ms_avg".to_string(), + json!(self.verify_span_upstream_reply_ms / verify_span_count), + ); + attrs.insert( + "skippy.verify_span_tokens_avg".to_string(), + json!(self.verify_span_token_count as f64 / verify_span_count), + ); + attrs.insert( + "skippy.verify_span_post_compute_ms_avg".to_string(), + json!(self.verify_span_post_compute_ms / verify_span_count), + ); + attrs.insert( + "skippy.verify_span_pre_reply_ms_avg".to_string(), + json!(self.verify_span_pre_reply_ms / verify_span_count), + ); + attrs.insert( + "skippy.verify_span_after_reply_ms_avg".to_string(), + json!(self.verify_span_after_reply_ms / verify_span_count), + ); + attrs.insert( + "skippy.verify_span_upstream_message_wait_ms_avg".to_string(), + json!(self.verify_span_upstream_message_wait_ms / verify_span_count), + ); + } let lookups = self.reply_stats.kv_lookup_hits + self.reply_stats.kv_lookup_misses; let hit_rate = if lookups > 0 { self.reply_stats.kv_lookup_hits as f64 / lookups as f64 @@ -3300,18 +3743,38 @@ pub(crate) fn run_binary_stage_message( .copied() .unwrap_or(message.state.current_token); let sampling = runtime_sampling_config(message.sampling.as_ref()); - let (predicted, output) = runtime.decode_frame_sampled( + if !native_mtp_enabled() { + let (predicted, output) = runtime.decode_frame_sampled( + session_id, + token_id, + sampling.as_ref(), + input, + output_capacity, + )?; + return Ok((predicted, vec![predicted], output)); + } + let (predicted, native_mtp, output) = runtime.decode_frame_sampled_mtp_n1( session_id, token_id, sampling.as_ref(), input, output_capacity, )?; - Ok((predicted, Vec::new(), output)) + Ok(( + predicted, + native_mtp_prediction_tokens(predicted, native_mtp), + output, + )) } WireMessageKind::VerifySpan => { - let (predicted_tokens, output) = - runtime.verify_frame(session_id, token_ids, input, output_capacity)?; + let sampling = runtime_sampling_config(message.sampling.as_ref()); + let (predicted_tokens, output) = runtime.verify_frame_sampled( + session_id, + token_ids, + sampling.as_ref(), + input, + output_capacity, + )?; let predicted = predicted_tokens.first().copied().unwrap_or(0); Ok((predicted, predicted_tokens, output)) } @@ -3332,6 +3795,26 @@ pub(crate) fn run_binary_stage_message( } } +fn is_decode_frame_batch_candidate( + config: &StageConfig, + message: &StageWireMessage, + token_ids: &[i32], +) -> bool { + if config.downstream.is_none() { + return false; + } + + matches!( + message.kind, + WireMessageKind::DecodeEmbd + | WireMessageKind::DecodeReadout + | WireMessageKind::DecodeLightCtx + | WireMessageKind::DecodeReplayEmbd + | WireMessageKind::DecodeReplayFinalEmbd + ) && message.token_count == 1 + && token_ids.len() == 1 +} + fn runtime_sampling_config(sampling: Option<&StageSamplingConfig>) -> Option { let sampling = sampling?; let mut config = SamplingConfig { diff --git a/crates/skippy-server/src/binary_transport/decode_batcher.rs b/crates/skippy-server/src/binary_transport/decode_batcher.rs new file mode 100644 index 0000000000..88de959857 --- /dev/null +++ b/crates/skippy-server/src/binary_transport/decode_batcher.rs @@ -0,0 +1,215 @@ +use std::{ + collections::VecDeque, + sync::{ + Arc, Condvar, Mutex, Weak, + atomic::{AtomicUsize, Ordering}, + mpsc as std_mpsc, + }, + thread, + time::Instant, +}; + +use anyhow::{Result, anyhow}; +use skippy_runtime::{ActivationFrame, SamplingConfig}; + +use crate::runtime_state::{RuntimeDecodeFrameBatchRequest, RuntimeState}; + +pub(crate) struct DecodeFrameBatcher { + shared: Arc, +} + +struct DecodeFrameBatcherShared { + runtime: Arc>, + state: Mutex, + ready: Condvar, + max_batch_size: usize, + owner_count: AtomicUsize, +} + +#[derive(Default)] +struct DecodeFrameBatcherState { + pending: VecDeque, + stopping: bool, +} + +struct PendingDecodeFrame { + session_id: String, + token_id: i32, + sampling: Option, + input: Option, + enqueued_at: Instant, + reply: std_mpsc::SyncSender>, +} + +pub(crate) struct DecodeFrameBatchOutcome { + pub(crate) predicted: i32, + pub(crate) output: ActivationFrame, + pub(crate) batch_size: usize, + pub(crate) batch_wait_ms: f64, + pub(crate) runtime_lock_wait_ms: f64, + pub(crate) runtime_lock_hold_ms: f64, +} + +impl DecodeFrameBatcher { + pub(crate) fn new(runtime: Arc>, max_batch_size: usize) -> Self { + let shared = Arc::new(DecodeFrameBatcherShared { + runtime, + state: Mutex::new(DecodeFrameBatcherState::default()), + ready: Condvar::new(), + max_batch_size: max_batch_size.max(1), + owner_count: AtomicUsize::new(1), + }); + let worker = Arc::downgrade(&shared); + thread::spawn(move || DecodeFrameBatcherShared::run_worker(worker)); + Self { shared } + } + + pub(crate) fn decode( + &self, + session_id: &str, + token_id: i32, + sampling: Option<&SamplingConfig>, + input: Option, + ) -> Result { + let (reply, receiver) = std_mpsc::sync_channel(1); + self.shared.enqueue(PendingDecodeFrame { + session_id: session_id.to_string(), + token_id, + sampling: sampling.cloned(), + input, + enqueued_at: Instant::now(), + reply, + })?; + receiver + .recv() + .map_err(|error| anyhow!("decode frame batcher stopped: {error}"))? + } +} + +impl Clone for DecodeFrameBatcher { + fn clone(&self) -> Self { + self.shared.owner_count.fetch_add(1, Ordering::Relaxed); + Self { + shared: self.shared.clone(), + } + } +} + +impl Drop for DecodeFrameBatcher { + fn drop(&mut self) { + if self.shared.owner_count.fetch_sub(1, Ordering::AcqRel) != 1 { + return; + } + if let Ok(mut state) = self.shared.state.lock() { + state.stopping = true; + self.shared.ready.notify_all(); + } + } +} + +impl DecodeFrameBatcherShared { + fn enqueue(&self, pending: PendingDecodeFrame) -> Result<()> { + let mut state = self + .state + .lock() + .map_err(|_| anyhow!("decode frame batcher lock poisoned"))?; + state.pending.push_back(pending); + self.ready.notify_one(); + Ok(()) + } + + fn run_worker(shared: Weak) { + while let Some(shared) = shared.upgrade() { + let Some(batch) = shared.wait_for_batch() else { + break; + }; + shared.run_batch(batch); + } + } + + fn wait_for_batch(&self) -> Option> { + let mut state = self + .state + .lock() + .expect("decode frame batcher lock poisoned"); + while state.pending.is_empty() && !state.stopping { + state = self + .ready + .wait(state) + .expect("decode frame batcher lock poisoned"); + } + if state.pending.is_empty() && state.stopping { + return None; + } + let batch_size = self.max_batch_size.min(state.pending.len()); + Some(state.pending.drain(..batch_size).collect()) + } + + fn run_batch(&self, batch: Vec) { + let batch_size = batch.len(); + let batch_wait_ms = batch + .iter() + .map(|pending| pending.enqueued_at.elapsed().as_secs_f64() * 1000.0) + .fold(0.0, f64::max); + let lock_started = Instant::now(); + let runtime_result = self + .runtime + .lock() + .map_err(|_| anyhow!("runtime lock poisoned while running decode frame batch")); + let runtime_lock_wait_ms = elapsed_ms(lock_started); + let result = runtime_result.and_then(|mut runtime| { + let hold_started = Instant::now(); + let requests = batch + .iter() + .map(|pending| RuntimeDecodeFrameBatchRequest { + session_id: pending.session_id.as_str(), + token_id: pending.token_id, + sampling: pending.sampling.as_ref(), + input: pending.input.as_ref(), + }) + .collect::>(); + let outputs = runtime.decode_frame_batch_sampled(&requests)?; + Ok((outputs, elapsed_ms(hold_started))) + }); + Self::send_batch_replies( + batch, + batch_size, + batch_wait_ms, + runtime_lock_wait_ms, + result, + ); + } + + fn send_batch_replies( + batch: Vec, + batch_size: usize, + batch_wait_ms: f64, + runtime_lock_wait_ms: f64, + result: Result<(Vec, f64)>, + ) { + match result { + Ok((outputs, runtime_lock_hold_ms)) => { + for (pending, output) in batch.into_iter().zip(outputs) { + let _ = pending.reply.send(Ok(DecodeFrameBatchOutcome { + predicted: output.predicted_token, + output: output.output, + batch_size, + batch_wait_ms, + runtime_lock_wait_ms, + runtime_lock_hold_ms, + })); + } + } + Err(error) => { + let error = error.to_string(); + for pending in batch { + let _ = pending.reply.send(Err(anyhow!(error.clone()))); + } + } + } + } +} + +fn elapsed_ms(started: Instant) -> f64 { + started.elapsed().as_secs_f64() * 1000.0 +} diff --git a/crates/skippy-server/src/binary_transport/direct_return.rs b/crates/skippy-server/src/binary_transport/direct_return.rs index 6a9c917ecc..dcdbdeb883 100644 --- a/crates/skippy-server/src/binary_transport/direct_return.rs +++ b/crates/skippy-server/src/binary_transport/direct_return.rs @@ -1,329 +1,21 @@ -use std::{ - collections::HashMap, - io, - net::{SocketAddr, TcpListener, TcpStream}, - sync::{ - Arc, Mutex, - atomic::{AtomicBool, Ordering}, - mpsc, - }, - thread::{self, JoinHandle}, - time::Duration, -}; +use std::net::TcpStream; -use anyhow::{Context, Result, anyhow, bail}; -use skippy_protocol::{ - StageConfig, StageTopology, - binary::{ - StageReply, StageStateHeader, StageWireMessage, WireActivationDType, WireMessageKind, - WireReplyKind, read_stage_message, recv_ready, recv_reply, send_ready, - send_reply_ack_with_stats, send_reply_predicted_tokens_with_stats, - send_reply_predicted_with_stats, write_stage_message, - }, -}; +use anyhow::{Result, bail}; +use skippy_protocol::binary::{WireMessageKind, recv_reply}; -use super::socket::{connect_downstream_socket, downstream_source_ip, resolve_downstream_endpoint}; - -#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] -pub(crate) struct PredictionReturnKey { - request_id: u64, - session_id: u64, -} - -impl PredictionReturnKey { - pub(crate) fn new(request_id: u64, session_id: u64) -> Self { - Self { - request_id, - session_id, - } - } -} - -pub struct PredictionReturnHub { - waiters: Mutex>>>, -} - -impl Default for PredictionReturnHub { - fn default() -> Self { - Self { - waiters: Mutex::new(HashMap::new()), - } - } -} - -pub struct PredictionReturnListener { - shutdown: Arc, - thread: Option>, - hub: Arc, -} - -impl PredictionReturnListener { - pub fn start(bind_addr: SocketAddr) -> Result { - let listener = TcpListener::bind(bind_addr) - .with_context(|| format!("bind direct prediction return listener {bind_addr}"))?; - listener - .set_nonblocking(true) - .context("set direct prediction return listener nonblocking")?; - let shutdown = Arc::new(AtomicBool::new(false)); - let thread_shutdown = shutdown.clone(); - let hub = Arc::new(PredictionReturnHub::default()); - let thread_hub = hub.clone(); - let thread = thread::spawn(move || { - while !thread_shutdown.load(Ordering::SeqCst) { - match listener.accept() { - Ok((stream, _)) => { - if let Err(error) = stream.set_nonblocking(false) { - eprintln!( - "direct prediction return connection failed: set blocking: {error}" - ); - continue; - } - let hub = thread_hub.clone(); - thread::spawn(move || { - if let Err(error) = handle_prediction_return_connection(hub, stream) { - eprintln!("direct prediction return connection failed: {error:#}"); - } - }); - } - Err(error) if error.kind() == io::ErrorKind::WouldBlock => { - thread::sleep(Duration::from_millis(50)); - } - Err(error) if error.kind() == io::ErrorKind::Interrupted => {} - Err(error) => { - eprintln!("direct prediction return listener failed: {error}"); - break; - } - } - } - }); - Ok(Self { - shutdown, - thread: Some(thread), - hub, - }) - } - - pub fn hub(&self) -> Arc { - self.hub.clone() - } -} - -impl Drop for PredictionReturnListener { - fn drop(&mut self) { - self.shutdown.store(true, Ordering::SeqCst); - if let Some(thread) = self.thread.take() { - let _ = thread.join(); - } - } -} - -fn handle_prediction_return_connection( - hub: Arc, - mut stream: TcpStream, -) -> Result<()> { - send_ready(&mut stream).context("send direct prediction return ready")?; - let open = read_stage_message(&mut stream, 0).context("read direct prediction return open")?; - hub.handle_return_connection(open, stream) -} +#[derive(Default)] +pub struct PredictionReturnHub; impl PredictionReturnHub { - pub(crate) fn register( - self: &Arc, - request_id: u64, - session_id: u64, - ) -> Result { - let key = PredictionReturnKey::new(request_id, session_id); - let (sender, receiver) = mpsc::channel(); - self.waiters - .lock() - .map_err(|_| anyhow!("prediction return hub lock poisoned"))? - .insert(key, sender); - Ok(PredictionReturnReceiver { - key, - hub: self.clone(), - receiver, - timeout: Duration::from_secs(300), - }) - } - - pub(crate) fn unregister(&self, key: PredictionReturnKey) { - if let Ok(mut waiters) = self.waiters.lock() { - waiters.remove(&key); - } - } - pub(crate) fn handle_return_connection( &self, - open: StageWireMessage, + open: skippy_protocol::binary::StageWireMessage, mut stream: TcpStream, ) -> Result<()> { if open.kind != WireMessageKind::PredictionReturnOpen { bail!("expected prediction return open message"); } - let key = PredictionReturnKey::new(open.request_id, open.session_id); - let sender = self - .waiters - .lock() - .map_err(|_| anyhow!("prediction return hub lock poisoned"))? - .get(&key) - .cloned() - .ok_or_else(|| anyhow!("no prediction return waiter for request {}", key.request_id))?; - loop { - match recv_reply(&mut stream) { - Ok(reply) => { - if sender.send(Ok(reply)).is_err() { - return Ok(()); - } - } - Err(error) if error.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(()), - Err(error) => { - let _ = sender.send(Err(error.to_string())); - return Err(error).context("read direct prediction return"); - } - } - } - } -} - -pub(crate) struct PredictionReturnReceiver { - key: PredictionReturnKey, - hub: Arc, - receiver: mpsc::Receiver>, - timeout: Duration, -} - -impl PredictionReturnReceiver { - pub(crate) fn recv_expected(&self, expected: WireReplyKind) -> Result { - let reply = self.recv()?; - if reply.kind != expected { - bail!( - "expected {expected:?} direct prediction return, got {:?}", - reply.kind - ); - } - Ok(reply) - } - - pub(crate) fn recv(&self) -> Result { - let reply = self - .receiver - .recv_timeout(self.timeout) - .context("timed out waiting for direct prediction return")? - .map_err(|error| anyhow!(error))?; - Ok(reply) - } -} - -impl Drop for PredictionReturnReceiver { - fn drop(&mut self) { - self.hub.unregister(self.key); - } -} - -pub(crate) fn open_prediction_return_stream( - config: &StageConfig, - topology: Option<&StageTopology>, - request_id: u64, - session_id: u64, - wire_dtype: WireActivationDType, - timeout_secs: u64, -) -> Result { - let endpoint = driver_stage_endpoint(config, topology)?; - let return_addr = resolve_downstream_endpoint(endpoint)?; - let source_ip = downstream_source_ip(config)?; - let attempts = timeout_secs.saturating_mul(2).max(1); - let mut last_error = None; - for _ in 0..attempts { - match connect_downstream_socket(return_addr, source_ip, Duration::from_secs(2)) { - Ok(mut stream) => { - stream.set_nodelay(true).ok(); - recv_ready(&mut stream).context("prediction return sink did not become ready")?; - write_stage_message( - &mut stream, - &prediction_return_open_message(request_id, session_id), - wire_dtype, - ) - .context("open direct prediction return stream")?; - return Ok(stream); - } - Err(error) => { - last_error = Some(anyhow!(error)); - std::thread::sleep(Duration::from_millis(500)); - } - } - } - Err(last_error - .unwrap_or_else(|| anyhow!("timed out")) - .context(format!( - "connect direct prediction return sink at {endpoint}" - ))) -} - -pub(crate) fn send_direct_prediction_return( - stream: &mut TcpStream, - reply: StageReply, -) -> Result<()> { - match reply.kind { - WireReplyKind::PredictedToken => { - send_reply_predicted_with_stats(stream, reply.predicted, reply.stats) - .context("send direct predicted-token return") - } - WireReplyKind::PredictedTokens => { - send_reply_predicted_tokens_with_stats(stream, &reply.predicted_tokens, reply.stats) - .context("send direct predicted-tokens return") - } - WireReplyKind::Ack => { - send_reply_ack_with_stats(stream, reply.stats).context("send direct ACK return") - } - } -} - -fn driver_stage_endpoint<'a>( - config: &'a StageConfig, - topology: Option<&'a StageTopology>, -) -> Result<&'a str> { - if let Some(topology) = topology { - return driver_stage_endpoint_from_topology(topology); - } - if let Some(upstream) = config - .upstream - .as_ref() - .filter(|upstream| upstream.stage_index == 0) - { - return Ok(strip_tcp_prefix(&upstream.endpoint)); - } - Err(anyhow!("direct prediction return requires topology")) -} - -fn driver_stage_endpoint_from_topology(topology: &StageTopology) -> Result<&str> { - topology - .stages - .iter() - .find(|stage| stage.stage_index == 0) - .map(|stage| strip_tcp_prefix(&stage.endpoint)) - .ok_or_else(|| anyhow!("topology does not contain driver-facing stage 0")) -} - -fn strip_tcp_prefix(endpoint: &str) -> &str { - endpoint.strip_prefix("tcp://").unwrap_or(endpoint) -} - -fn prediction_return_open_message(request_id: u64, session_id: u64) -> StageWireMessage { - StageWireMessage { - kind: WireMessageKind::PredictionReturnOpen, - pos_start: 0, - token_count: 0, - state: StageStateHeader::new( - WireMessageKind::PredictionReturnOpen, - WireActivationDType::F32, - ), - request_id, - session_id, - sampling: None, - chat_sampling_metadata: None, - tokens: Vec::new(), - positions: Vec::new(), - activation: Vec::new(), - raw_bytes: Vec::new(), + while recv_reply(&mut stream).is_ok() {} + Ok(()) } } diff --git a/crates/skippy-server/src/binary_transport/forwarding.rs b/crates/skippy-server/src/binary_transport/forwarding.rs index fc87b37664..82432abdf3 100644 --- a/crates/skippy-server/src/binary_transport/forwarding.rs +++ b/crates/skippy-server/src/binary_transport/forwarding.rs @@ -1,11 +1,11 @@ use std::time::Instant; -use anyhow::{Context, Result}; +use anyhow::{Context, Result, bail}; use skippy_protocol::{ StageConfig, binary::{StageWireMessage, WireActivationDType, activation_state_flags_from_frame_flags}, }; -use skippy_runtime::ActivationFrame; +use skippy_runtime::{ActivationFrame, RuntimeActivationDType}; pub(crate) fn forwarded_stage_message( config: &StageConfig, @@ -37,24 +37,20 @@ pub(crate) fn forwarded_stage_message_timed( state.reserved = wire_dtype as i32; state.flags |= activation_state_flags_from_frame_flags(output.desc.flags); let encode_started = Instant::now(); - let activation = skippy_protocol::binary::encode_f32_activation_payload_with_state_flags( - wire_dtype, - incoming.token_count, - activation_width, - &output.payload, - state.flags, - ) - .with_context(|| { - format!( - "encode output activation payload; wire_dtype={wire_dtype:?} incoming_tokens={} output_tokens={} activation_width={} payload_bytes={} frame_payload_bytes={} state_flags={}", - incoming.token_count, - output.desc.token_count, - activation_width, - output.payload.len(), - output.desc.payload_bytes, - state.flags, - ) - })?; + let activation = + encode_output_activation_payload(wire_dtype, incoming, output, activation_width, state.flags) + .with_context(|| { + format!( + "encode output activation payload; wire_dtype={wire_dtype:?} frame_dtype={:?} incoming_tokens={} output_tokens={} activation_width={} payload_bytes={} frame_payload_bytes={} state_flags={}", + output.desc.dtype, + incoming.token_count, + output.desc.token_count, + activation_width, + output.payload.len(), + output.desc.payload_bytes, + state.flags, + ) + })?; Ok(ForwardedStageMessage { message: StageWireMessage { kind: incoming.kind, @@ -74,6 +70,55 @@ pub(crate) fn forwarded_stage_message_timed( }) } +fn encode_output_activation_payload( + wire_dtype: WireActivationDType, + incoming: &StageWireMessage, + output: &ActivationFrame, + activation_width: i32, + state_flags: i32, +) -> Result> { + match (output.desc.dtype, wire_dtype) { + (RuntimeActivationDType::F32, _) => Ok( + skippy_protocol::binary::encode_f32_activation_payload_with_state_flags( + wire_dtype, + incoming.token_count, + activation_width, + &output.payload, + state_flags, + )?, + ), + (RuntimeActivationDType::F16, WireActivationDType::F16) => { + validate_f16_passthrough_payload(incoming, output, activation_width, state_flags)?; + Ok(output.payload.clone()) + } + (dtype, wire_dtype) => { + bail!("unsupported activation dtype conversion: {dtype:?} to {wire_dtype:?}") + } + } +} + +fn validate_f16_passthrough_payload( + incoming: &StageWireMessage, + output: &ActivationFrame, + activation_width: i32, + state_flags: i32, +) -> Result<()> { + if output.payload.len() as u64 != output.desc.payload_bytes { + bail!("F16 activation payload length does not match frame descriptor"); + } + let expected = skippy_protocol::binary::activation_wire_bytes_with_state_flags( + WireActivationDType::F16, + incoming.token_count, + activation_width, + state_flags, + ) + .context("compute expected F16 activation payload size")?; + if output.payload.len() != expected { + bail!("F16 activation payload size mismatch"); + } + Ok(()) +} + #[cfg(test)] mod tests { use super::*; @@ -142,9 +187,9 @@ mod tests { } } - fn rwkv7_sideband_frame() -> ActivationFrame { + fn f32_frame(flags: u64, token_count: u32, values: &[f32]) -> ActivationFrame { let mut payload = Vec::new(); - for value in [1.0_f32, 2.0, 3.0, 4.0] { + for value in values { payload.extend_from_slice(&value.to_le_bytes()); } ActivationFrame { @@ -155,15 +200,41 @@ mod tests { producer_stage_index: 1, layer_start: 4, layer_end: 8, - token_count: 1, + token_count, sequence_count: 1, payload_bytes: payload.len() as u64, - flags: skippy_protocol::binary::ACTIVATION_FLAG_RWKV7_V_FIRST, + flags, }, payload, } } + fn rwkv7_sideband_frame() -> ActivationFrame { + f32_frame( + skippy_protocol::binary::ACTIVATION_FLAG_RWKV7_V_FIRST, + 1, + &[1.0_f32, 2.0, 3.0, 4.0], + ) + } + + fn f16_frame() -> ActivationFrame { + ActivationFrame { + desc: ActivationDesc { + version: 1, + dtype: RuntimeActivationDType::F16, + layout: RuntimeActivationLayout::TokenMajor, + producer_stage_index: 1, + layer_start: 4, + layer_end: 8, + token_count: 2, + sequence_count: 1, + payload_bytes: 8, + flags: 0, + }, + payload: vec![0, 1, 2, 3, 4, 5, 6, 7], + } + } + #[test] fn forwarded_stage_message_preserves_rwkv7_sideband_shape() { let forwarded = forwarded_stage_message_timed( @@ -203,4 +274,47 @@ mod tests { 0 ); } + + #[test] + fn forwarded_stage_message_passes_through_f16_activation_for_f16_wire() { + let mut incoming = incoming_message(); + incoming.token_count = 2; + + let forwarded = forwarded_stage_message_timed( + &stage_config(), + &incoming, + &f16_frame(), + WireActivationDType::F16, + 2, + ) + .unwrap(); + + assert_eq!(forwarded.message.activation, vec![0, 1, 2, 3, 4, 5, 6, 7]); + assert_eq!( + forwarded.message.state.reserved, + WireActivationDType::F16 as i32 + ); + } + + #[test] + fn forwarded_stage_message_rejects_bad_f16_passthrough_size() { + let mut incoming = incoming_message(); + incoming.token_count = 2; + let mut output = f16_frame(); + output.payload.pop(); + output.desc.payload_bytes = output.payload.len() as u64; + + let error = match forwarded_stage_message_timed( + &stage_config(), + &incoming, + &output, + WireActivationDType::F16, + 2, + ) { + Ok(_) => panic!("expected bad F16 passthrough payload to fail"), + Err(error) => error, + }; + + assert!(format!("{error:#}").contains("F16 activation payload size mismatch")); + } } diff --git a/crates/skippy-server/src/binary_transport/kv_eviction.rs b/crates/skippy-server/src/binary_transport/kv_eviction.rs index f10a11a0c1..e3da6f95cd 100644 --- a/crates/skippy-server/src/binary_transport/kv_eviction.rs +++ b/crates/skippy-server/src/binary_transport/kv_eviction.rs @@ -42,6 +42,10 @@ impl BinaryProactiveEviction { pub(super) fn insert_attrs(&self, attrs: &mut BTreeMap) { attrs.extend(self.attrs()); } + + pub(super) fn should_emit_summary(&self) -> bool { + self.error_kind.is_some() || self.evicted_entries > 0 || self.evicted_tokens > 0 + } } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -116,3 +120,47 @@ pub(super) fn evict_binary_resident_prefix_for_decode( evicted_tokens: eviction.evicted_tokens, }) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn disabled_and_noop_evictions_are_debug_only() { + assert!(!BinaryProactiveEviction::disabled().should_emit_summary()); + assert!( + !BinaryProactiveEviction { + status: "noop", + error_kind: None, + target_tokens: 1024, + evicted_entries: 0, + evicted_tokens: 0, + } + .should_emit_summary() + ); + } + + #[test] + fn actionable_evictions_stay_summary_visible() { + assert!( + BinaryProactiveEviction { + status: "evicted", + error_kind: None, + target_tokens: 1024, + evicted_entries: 1, + evicted_tokens: 512, + } + .should_emit_summary() + ); + assert!( + BinaryProactiveEviction { + status: "error", + error_kind: Some("runtime"), + target_tokens: 1024, + evicted_entries: 0, + evicted_tokens: 0, + } + .should_emit_summary() + ); + } +} diff --git a/crates/skippy-server/src/binary_transport/tests.rs b/crates/skippy-server/src/binary_transport/tests.rs index 4977c71133..77473ebf12 100644 --- a/crates/skippy-server/src/binary_transport/tests.rs +++ b/crates/skippy-server/src/binary_transport/tests.rs @@ -1,7 +1,7 @@ use super::{ binary_full_prefill_record_identities, decode_record_tokens_sideband, - prepare_binary_stage_connection, restore_prefill_decode_as_decode_message, - token_sideband_or_fill, + is_decode_frame_batch_candidate, native_mtp_enabled_from, prepare_binary_stage_connection, + restore_prefill_decode_as_decode_message, token_sideband_or_fill, }; use std::{ io, @@ -14,7 +14,8 @@ use std::{ use crate::kv_integration::KvStageIntegration; use crate::runtime_state::RuntimeState; use skippy_protocol::binary::{ - StageSamplingConfig, StageStateHeader, StageWireMessage, WireActivationDType, WireMessageKind, + StageReplyStats, StageSamplingConfig, StageStateHeader, StageWireMessage, WireActivationDType, + WireMessageKind, }; use skippy_protocol::{ LoadMode, PeerConfig, StageConfig, StageKvCacheConfig, StageKvCacheMode, StageKvCachePayload, @@ -52,6 +53,66 @@ fn accepted_binary_stage_connection_is_blocking() { drop(client.join().unwrap()); } +#[test] +fn native_mtp_enabled_flag_defaults_on_and_accepts_false_values() { + assert!(native_mtp_enabled_from(None)); + assert!(native_mtp_enabled_from(Some("1"))); + assert!(native_mtp_enabled_from(Some("true"))); + assert!(!native_mtp_enabled_from(Some("0"))); + assert!(!native_mtp_enabled_from(Some("false"))); + assert!(!native_mtp_enabled_from(Some(" disabled "))); +} + +#[test] +fn request_summary_tracks_verify_span_compute_ms() { + let config = prefix_cache_test_config(); + let mut summary = super::BinaryRequestSummary::default(); + let verify = test_message(WireMessageKind::VerifySpan, 2); + let decode = test_message(WireMessageKind::DecodeEmbd, 1); + + summary.observe(summary_observation(&config, &verify, 12.5)); + summary.observe(summary_observation(&config, &decode, 7.0)); + + assert_eq!(summary.verify_span_count, 1); + assert_eq!(summary.verify_span_token_count, 2); + assert_eq!(summary.verify_span_max_tokens, 2); + assert_eq!(summary.verify_span_compute_ms, 12.5); + assert_eq!(summary.verify_span_input_activation_decode_ms, 1.25); + assert_eq!(summary.verify_span_runtime_lock_hold_ms, 2.5); + assert_eq!(summary.verify_span_upstream_reply_ms, 0.75); + assert_eq!(summary.compute_ms, 19.5); + assert_eq!(summary.input_activation_decode_ms, 2.5); + assert_eq!(summary.runtime_lock_hold_ms, 5.0); + assert_eq!(summary.upstream_reply_ms, 1.5); +} + +#[test] +fn request_summary_tracks_auto_align_totals() { + let config = prefix_cache_test_config(); + let mut summary = super::BinaryRequestSummary::default(); + let verify = test_message(WireMessageKind::VerifySpan, 2); + let decode = test_message(WireMessageKind::DecodeEmbd, 1); + + let mut verify_observation = summary_observation(&config, &verify, 12.5); + verify_observation.session_auto_align_count = 1; + verify_observation.session_auto_align_ms = 0.75; + verify_observation.session_auto_align_trimmed_tokens = 1; + summary.observe(verify_observation); + + let mut decode_observation = summary_observation(&config, &decode, 7.0); + decode_observation.session_auto_align_count = 1; + decode_observation.session_auto_align_ms = 1.25; + decode_observation.session_auto_align_trimmed_tokens = 2; + summary.observe(decode_observation); + + assert_eq!(summary.session_auto_align_count, 2); + assert_eq!(summary.session_auto_align_ms, 2.0); + assert_eq!(summary.session_auto_align_trimmed_tokens, 3); + assert_eq!(summary.verify_span_session_auto_align_count, 1); + assert_eq!(summary.verify_span_session_auto_align_ms, 0.75); + assert_eq!(summary.verify_span_session_auto_align_trimmed_tokens, 1); +} + #[test] fn restore_prefill_decode_as_decode_preserves_chat_metadata() { let metadata = r#"{"grammar":"chat"}"#; @@ -179,6 +240,58 @@ fn prefix_cache_test_config() -> StageConfig { } } +fn test_message(kind: WireMessageKind, token_count: i32) -> StageWireMessage { + StageWireMessage { + kind, + pos_start: 0, + token_count, + state: StageStateHeader::new(kind, WireActivationDType::F16), + request_id: 11, + session_id: 13, + sampling: None, + chat_sampling_metadata: None, + tokens: Vec::new(), + positions: Vec::new(), + activation: Vec::new(), + raw_bytes: Vec::new(), + } +} + +fn summary_observation<'a>( + config: &'a StageConfig, + message: &'a StageWireMessage, + compute_ms: f64, +) -> super::BinaryMessageObservation<'a> { + super::BinaryMessageObservation { + config, + message, + reply_stats: StageReplyStats::default(), + compute_ms, + forward_write_ms: 0.0, + downstream_wait_ms: 0.0, + upstream_reply_ms: 0.75, + message_elapsed_ms: compute_ms, + input_activation_decode_ms: 1.25, + forward_activation_encode_ms: 0.0, + runtime_lock_hold_ms: 2.5, + input_activation_bytes: 0, + output_activation_bytes: 0, + prefill_credit_limit: 0, + pending_prefill_replies_before: 0, + pending_prefill_replies_after: 0, + credit_wait_count: 0, + deferred_prefill_replies_drained: 0, + session_auto_align_count: 0, + session_auto_align_ms: 0.0, + session_auto_align_trimmed_tokens: 0, + verify_span_pre_compute_ms: 0.25, + verify_span_post_compute_ms: 0.5, + verify_span_pre_reply_ms: 0.0, + verify_span_after_reply_ms: 0.0, + upstream_message_wait_ms: 0.0, + } +} + fn prefill_message() -> StageWireMessage { StageWireMessage { kind: WireMessageKind::PrefillEmbd, @@ -251,6 +364,23 @@ fn decode_record_tokens_sideband_rejects_wrong_checkpoint_len() { assert_eq!(token_sideband_or_fill(&message).unwrap(), vec![104]); } +#[test] +fn decode_frame_batch_candidate_keeps_intermediate_decode_batching() { + let config = prefix_cache_test_config(); + let message = first_decode_message_with_full_prompt_sideband(); + + assert!(is_decode_frame_batch_candidate(&config, &message, &[104])); +} + +#[test] +fn decode_frame_batch_candidate_skips_final_output_stage() { + let mut config = prefix_cache_test_config(); + config.downstream = None; + let message = first_decode_message_with_full_prompt_sideband(); + + assert!(!is_decode_frame_batch_candidate(&config, &message, &[104])); +} + #[test] fn binary_full_prefill_record_plan_includes_shared_prefix_candidate() { let config = prefix_cache_test_config(); diff --git a/crates/skippy-server/src/frontend.rs b/crates/skippy-server/src/frontend.rs index 46360455fc..02eb9925e4 100644 --- a/crates/skippy-server/src/frontend.rs +++ b/crates/skippy-server/src/frontend.rs @@ -57,9 +57,9 @@ use tokio::{ use crate::{ binary_transport::{ - PredictionReturnHub, PredictionReturnReceiver, WireCondition, connect_binary_downstream, - forwarded_stage_message, forwarded_stage_message_timed, run_binary_stage_message, - stage_output_activation_capacity, write_stage_message_conditioned, + DecodeFrameBatcher, WireCondition, connect_binary_downstream, forwarded_stage_message, + forwarded_stage_message_timed, run_binary_stage_message, stage_output_activation_capacity, + write_stage_message_conditioned, }, cli::ServeOpenAiArgs, config::{load_json, validate_config}, @@ -70,10 +70,12 @@ use crate::{ mod admission; mod backend; +mod decode_batcher; mod embedded_execution; mod embedded_generation; mod generation_flow; mod local_generation; +mod native_mtp; mod prefill; mod prefix_cache; mod prompting; @@ -84,6 +86,8 @@ mod wire_messages; use self::{ admission::{GenerationTokenBudget, GenerationTokenBudgetRequest}, + decode_batcher::DecodeBatcher, + native_mtp::*, prefill::*, request::*, speculative::*, @@ -174,6 +178,9 @@ pub async fn serve_openai(args: ServeOpenAiArgs) -> Result<()> { } let kv = KvStageIntegration::from_config(&config)?.map(Arc::new); let ctx_size = usize::try_from(config.ctx_size).unwrap_or(usize::MAX); + let decode_batcher = DecodeBatcher::new(runtime.clone(), args.generation_concurrency); + let decode_frame_batcher = + DecodeFrameBatcher::new(runtime.clone(), args.generation_concurrency); let backend = Arc::new(StageOpenAiBackend { runtime, config, @@ -192,6 +199,8 @@ pub async fn serve_openai(args: ServeOpenAiArgs) -> Result<()> { generation_token_budget: Arc::new(GenerationTokenBudget::new(ctx_size)), hook_policy: None, kv, + decode_batcher, + decode_frame_batcher, }); let app: Router = instrumented_openai_router(backend, telemetry.clone()); @@ -229,7 +238,6 @@ pub struct EmbeddedOpenAiArgs { pub reply_credit_limit: Option, pub downstream_connect_timeout_secs: u64, pub downstream_wire_condition: WireCondition, - pub prediction_returns: Option>, pub telemetry: Telemetry, pub hook_policy: Option>, pub openai_guardrails: Option, @@ -505,7 +513,6 @@ pub fn embedded_openai_backend(args: EmbeddedOpenAiArgs) -> Result Result = Arc::new(StageOpenAiBackend { runtime: args.runtime, config: args.config.clone(), @@ -537,6 +547,8 @@ pub fn embedded_openai_backend(args: EmbeddedOpenAiArgs) -> Result, hook_policy: Option>, kv: Option>, + decode_batcher: DecodeBatcher, + decode_frame_batcher: DecodeFrameBatcher, } struct GenerationQueueReservation { @@ -812,7 +826,6 @@ enum OpenAiBackendMode { downstream_wire_condition: WireCondition, prefill_reply_credit_limit: usize, lane_pool: Option>, - prediction_returns: Option>, }, } @@ -1045,7 +1058,17 @@ impl PersistentStageLanePool { let timer = PhaseTimer::start(); let mut stream = connect_binary_downstream(&self.config, self.timeout_secs)? .ok_or_else(|| anyhow!("embedded stage0 has no downstream"))?; + let local_addr = stream.local_addr().ok(); + let peer_addr = stream.peer_addr().ok(); + eprintln!( + "openai downstream lane waiting ready: stage_id={} lane_id={lane_id} local={local_addr:?} peer={peer_addr:?}", + self.config.stage_id + ); recv_ready(&mut stream).context("persistent downstream lane did not become ready")?; + eprintln!( + "openai downstream lane received ready: stage_id={} lane_id={lane_id} local={local_addr:?} peer={peer_addr:?}", + self.config.stage_id + ); let mut attrs = lifecycle_attrs(&self.config); attrs.insert( "llama_stage.openai_downstream_lane_id".to_string(), @@ -1486,6 +1509,13 @@ fn tool_calls_requested(request: &ChatCompletionRequest) -> bool { .is_some_and(|choice| matches!(choice.as_str(), Some("none"))) } +fn chat_output_parser_required( + request: &ChatCompletionRequest, + template_options: &ChatTemplateOptions, +) -> bool { + tool_calls_requested(request) || template_options.enable_thinking == Some(true) +} + fn chat_response_from_generated_text( model: String, output: &GeneratedText, @@ -1708,7 +1738,6 @@ struct EmbeddedStageZeroGeneration<'a> { downstream_wire_condition: WireCondition, prefill_reply_credit_limit: usize, lane_pool: Option>, - prediction_return: Option, draft: Option>>, speculative_window: usize, adaptive_speculative_window: bool, @@ -1734,7 +1763,6 @@ struct SplitMultimodalGeneration<'a> { activation_width: i32, downstream_wire_condition: WireCondition, lane_pool: Arc, - prediction_return: Option, } struct EmbeddedLocalOutput { @@ -1764,6 +1792,7 @@ struct EmbeddedStageExecution { struct EmbeddedFusedFirstDecode { predicted: i32, predicted_tokens: Vec, + native_mtp_draft: Option, reply_stats: StageReplyStats, execution: EmbeddedExecutionStats, elapsed_ms: f64, @@ -1844,9 +1873,8 @@ impl ChatOutputStreamParser { if let Some(delta) = suffix_delta(parsed.content.as_deref(), &mut self.emitted_content) { events.push(GenerationStreamEvent::Delta(delta)); } - if !is_partial - && !self.emitted_tool_calls - && let Some(tool_calls) = parsed.tool_calls + if let (true, Some(tool_calls)) = + (!is_partial && !self.emitted_tool_calls, parsed.tool_calls) { self.emitted_tool_calls = true; events.push(GenerationStreamEvent::ToolCalls(tool_calls)); diff --git a/crates/skippy-server/src/frontend/backend.rs b/crates/skippy-server/src/frontend/backend.rs index fb559a03d2..349b56ec6f 100644 --- a/crates/skippy-server/src/frontend/backend.rs +++ b/crates/skippy-server/src/frontend/backend.rs @@ -18,6 +18,7 @@ impl OpenAiBackend for StageOpenAiBackend { ensure_chat_runtime_features_supported(&request)?; let sampling = chat_sampling_config(&request)?; let template_options = chat_template_options(&request)?; + let parse_chat_output = chat_output_parser_required(&request, &template_options); let template_timer = PhaseTimer::start(); let prompt = self.prepare_chat_prompt(&request, template_options)?; let mut template_attrs = self.openai_attrs(&ids); @@ -54,12 +55,16 @@ impl OpenAiBackend for StageOpenAiBackend { ) .await?; let response_timer = PhaseTimer::start(); - let parsed_message = self.parse_chat_output( - &output.text, - &request, - chat_parse_metadata.as_deref(), - false, - )?; + let parsed_message = if parse_chat_output { + self.parse_chat_output( + &output.text, + &request, + chat_parse_metadata.as_deref(), + false, + )? + } else { + None + }; let response = chat_response_from_generated_text(request.model.clone(), &output, parsed_message); let mut response_attrs = self.openai_attrs(&ids); @@ -111,6 +116,7 @@ impl OpenAiBackend for StageOpenAiBackend { let sampling = chat_sampling_config(&request)?; let include_usage = request.include_usage(); let template_options = chat_template_options(&request)?; + let parse_chat_output = chat_output_parser_required(&request, &template_options); let template_timer = PhaseTimer::start(); let prompt = self.prepare_chat_prompt(&request, template_options)?; let mut template_attrs = self.openai_attrs(&ids); @@ -144,6 +150,7 @@ impl OpenAiBackend for StageOpenAiBackend { sampling, include_usage, Some(request.clone()), + parse_chat_output, context, ids, ) @@ -261,6 +268,7 @@ impl OpenAiBackend for StageOpenAiBackend { sampling, include_usage, None, + false, context, ids, ) @@ -425,6 +433,7 @@ impl StageOpenAiBackend { sampling: SamplingConfig, include_usage: bool, hook_request: Option, + parse_chat_output: bool, context: OpenAiRequestContext, ids: OpenAiGenerationIds, ) -> OpenAiResult { @@ -440,16 +449,17 @@ impl StageOpenAiBackend { let chat_parse_metadata = prompt.chat_parse_metadata.clone(); let (tx, rx) = mpsc::channel(16); let hook_runtime = Some(tokio::runtime::Handle::current()); - let mut chat_stream_parser = - if let (Some(request), Some(metadata)) = (hook_request.clone(), chat_parse_metadata) { - Some(ChatOutputStreamParser::new( - backend.clone(), - request, - metadata, - )) - } else { - None - }; + let mut chat_stream_parser = if let (true, Some(request), Some(metadata)) = + (parse_chat_output, hook_request.clone(), chat_parse_metadata) + { + Some(ChatOutputStreamParser::new( + backend.clone(), + request, + metadata, + )) + } else { + None + }; task::spawn_blocking(move || { let _permit = permit; let result = backend.generate_text( diff --git a/crates/skippy-server/src/frontend/decode_batcher.rs b/crates/skippy-server/src/frontend/decode_batcher.rs new file mode 100644 index 0000000000..00469595db --- /dev/null +++ b/crates/skippy-server/src/frontend/decode_batcher.rs @@ -0,0 +1,202 @@ +use std::{ + collections::VecDeque, + sync::{ + Arc, Condvar, Mutex, Weak, + atomic::{AtomicUsize, Ordering}, + mpsc as std_mpsc, + }, + thread, + time::Instant, +}; + +use super::*; +use crate::runtime_state::RuntimeDecodeBatchRequest; + +pub(super) struct DecodeBatcher { + shared: Arc, +} + +struct DecodeBatcherShared { + runtime: Arc>, + state: Mutex, + ready: Condvar, + max_batch_size: usize, + owner_count: AtomicUsize, +} + +#[derive(Default)] +struct DecodeBatcherState { + pending: VecDeque, + stopping: bool, +} + +struct PendingDecode { + session_id: String, + token_id: i32, + sampling: Option, + enqueued_at: Instant, + reply: std_mpsc::SyncSender>, +} + +#[derive(Debug, Clone, Copy)] +pub(super) struct DecodeBatchOutcome { + pub predicted: i32, + pub batch_size: usize, + pub batch_wait_ms: f64, + pub runtime_lock_wait_ms: f64, + pub runtime_lock_hold_ms: f64, +} + +impl DecodeBatcher { + pub(super) fn new(runtime: Arc>, max_batch_size: usize) -> Self { + let shared = Arc::new(DecodeBatcherShared { + runtime, + state: Mutex::new(DecodeBatcherState::default()), + ready: Condvar::new(), + max_batch_size: max_batch_size.max(1), + owner_count: AtomicUsize::new(1), + }); + let worker = Arc::downgrade(&shared); + thread::spawn(move || DecodeBatcherShared::run_worker(worker)); + Self { shared } + } + + pub(super) fn decode( + &self, + session_id: &str, + token_id: i32, + sampling: Option<&SamplingConfig>, + ) -> OpenAiResult { + let (reply, receiver) = std_mpsc::sync_channel(1); + self.shared.enqueue(PendingDecode { + session_id: session_id.to_string(), + token_id, + sampling: sampling.cloned(), + enqueued_at: Instant::now(), + reply, + })?; + receiver + .recv() + .map_err(|error| OpenAiError::backend(format!("decode batcher stopped: {error}")))? + } +} + +impl Clone for DecodeBatcher { + fn clone(&self) -> Self { + self.shared.owner_count.fetch_add(1, Ordering::Relaxed); + Self { + shared: self.shared.clone(), + } + } +} + +impl Drop for DecodeBatcher { + fn drop(&mut self) { + if self.shared.owner_count.fetch_sub(1, Ordering::AcqRel) != 1 { + return; + } + if let Ok(mut state) = self.shared.state.lock() { + state.stopping = true; + self.shared.ready.notify_all(); + } + } +} + +impl DecodeBatcherShared { + fn enqueue(&self, pending: PendingDecode) -> OpenAiResult<()> { + let mut state = self + .state + .lock() + .map_err(|_| OpenAiError::backend("decode batcher lock poisoned"))?; + state.pending.push_back(pending); + self.ready.notify_one(); + Ok(()) + } + + fn run_worker(shared: Weak) { + while let Some(shared) = shared.upgrade() { + let Some(batch) = shared.wait_for_batch() else { + break; + }; + shared.run_batch(batch); + } + } + + fn wait_for_batch(&self) -> Option> { + let mut state = self.state.lock().expect("decode batcher lock poisoned"); + while state.pending.is_empty() && !state.stopping { + state = self + .ready + .wait(state) + .expect("decode batcher lock poisoned"); + } + if state.pending.is_empty() && state.stopping { + return None; + } + let batch_size = self.max_batch_size.min(state.pending.len()); + Some(state.pending.drain(..batch_size).collect()) + } + + fn run_batch(&self, batch: Vec) { + let batch_size = batch.len(); + let batch_wait_ms = batch + .iter() + .map(|pending| pending.enqueued_at.elapsed().as_secs_f64() * 1000.0) + .fold(0.0, f64::max); + let lock_timer = PhaseTimer::start(); + let runtime_result = self + .runtime + .lock() + .map_err(|_| OpenAiError::backend("runtime lock poisoned while running decode batch")); + let runtime_lock_wait_ms = lock_timer.elapsed_ms(); + let result = runtime_result.and_then(|mut runtime| { + let hold_timer = PhaseTimer::start(); + let requests = batch + .iter() + .map(|pending| RuntimeDecodeBatchRequest { + session_id: pending.session_id.as_str(), + token_id: pending.token_id, + sampling: pending.sampling.as_ref(), + }) + .collect::>(); + let predicted = runtime + .decode_batch_sampled(&requests) + .map_err(openai_backend_error)?; + Ok((predicted, hold_timer.elapsed_ms())) + }); + Self::send_batch_replies( + batch, + batch_size, + batch_wait_ms, + runtime_lock_wait_ms, + result, + ); + } + + fn send_batch_replies( + batch: Vec, + batch_size: usize, + batch_wait_ms: f64, + runtime_lock_wait_ms: f64, + result: OpenAiResult<(Vec, f64)>, + ) { + match result { + Ok((predicted, runtime_lock_hold_ms)) => { + for (pending, predicted) in batch.into_iter().zip(predicted) { + let _ = pending.reply.send(Ok(DecodeBatchOutcome { + predicted, + batch_size, + batch_wait_ms, + runtime_lock_wait_ms, + runtime_lock_hold_ms, + })); + } + } + Err(error) => { + for pending in batch { + let _ = pending.reply.send(Err(error.clone())); + } + } + } + } +} diff --git a/crates/skippy-server/src/frontend/embedded_execution.rs b/crates/skippy-server/src/frontend/embedded_execution.rs index 07193d6985..195d91ac28 100644 --- a/crates/skippy-server/src/frontend/embedded_execution.rs +++ b/crates/skippy-server/src/frontend/embedded_execution.rs @@ -77,13 +77,14 @@ impl StageOpenAiBackend { .map_err(openai_io_error)?; let forward_write_ms = write_timer.elapsed_ms(); let wait_timer = PhaseTimer::start(); - let reply = request - .prediction_return - .as_ref() - .ok_or_else(|| OpenAiError::backend("missing direct prediction return receiver"))? - .recv_expected(expected_reply) - .map_err(openai_backend_error)?; + let reply = recv_reply(&mut *downstream).map_err(openai_io_error)?; let downstream_wait_ms = wait_timer.elapsed_ms(); + if reply.kind != expected_reply { + return Err(OpenAiError::backend(format!( + "expected embedded stage {expected_reply:?} reply from downstream, got {:?}", + reply.kind + ))); + } stats.merge(reply.stats); if message.kind == WireMessageKind::VerifySpan { stats.verify_span_compute_us += ms_to_us(stage0_compute_ms); @@ -164,4 +165,77 @@ impl StageOpenAiBackend { downstream_wait_ms, }) } + + pub(super) fn trim_embedded_stage_session( + &self, + request: &EmbeddedStageZeroGeneration<'_>, + downstream: &mut TcpStream, + session_key: &str, + request_id: u64, + session_id: u64, + token_count: usize, + ) -> OpenAiResult { + let timer = PhaseTimer::start(); + let local_timer = PhaseTimer::start(); + { + let mut runtime = self + .runtime + .lock() + .map_err(|_| OpenAiError::backend("runtime lock poisoned"))?; + runtime + .trim_session(session_key, token_count as u64) + .map_err(openai_backend_error)?; + } + let local_ms = local_timer.elapsed_ms(); + let message = + embedded_trim_session_message(request.wire_dtype, request_id, session_id, token_count)?; + let write_timer = PhaseTimer::start(); + write_stage_message_conditioned( + &mut *downstream, + &message, + request.wire_dtype, + request.downstream_wire_condition, + ) + .map_err(openai_io_error)?; + let downstream_write_ms = write_timer.elapsed_ms(); + let wait_timer = PhaseTimer::start(); + let reply = recv_reply(&mut *downstream).map_err(openai_io_error)?; + let downstream_wait_ms = wait_timer.elapsed_ms(); + if reply.kind != WireReplyKind::Ack { + return Err(OpenAiError::backend(format!( + "trim expected ACK from downstream, got {:?}", + reply.kind + ))); + } + Ok(EmbeddedSessionControl { + elapsed_ms: timer.elapsed_ms(), + local_ms, + downstream_write_ms, + downstream_wait_ms, + }) + } + + pub(super) fn trim_embedded_stage_session_local( + &self, + session_key: &str, + token_count: usize, + ) -> OpenAiResult { + let timer = PhaseTimer::start(); + let local_timer = PhaseTimer::start(); + { + let mut runtime = self + .runtime + .lock() + .map_err(|_| OpenAiError::backend("runtime lock poisoned"))?; + runtime + .trim_session(session_key, token_count as u64) + .map_err(openai_backend_error)?; + } + Ok(EmbeddedSessionControl { + elapsed_ms: timer.elapsed_ms(), + local_ms: local_timer.elapsed_ms(), + downstream_write_ms: 0.0, + downstream_wait_ms: 0.0, + }) + } } diff --git a/crates/skippy-server/src/frontend/embedded_generation.rs b/crates/skippy-server/src/frontend/embedded_generation.rs index 7fb8bc30c4..020a3aae97 100644 --- a/crates/skippy-server/src/frontend/embedded_generation.rs +++ b/crates/skippy-server/src/frontend/embedded_generation.rs @@ -134,14 +134,17 @@ impl StageOpenAiBackend { fused_first_decode = Some(fused); } } - if !prefill_chain_cache_restored - && let Some(restore) = self.try_restore_embedded_split_prefill( + let split_prefill_restore = if prefill_chain_cache_restored { + None + } else { + self.try_restore_embedded_split_prefill( &request, &session_key, downstream, prefill_tokens, )? - { + }; + if let Some(restore) = split_prefill_restore { prefill_chain_restored_tokens = restore.restored_tokens; prefill_chain_cache_restored = prefill_chain_restored_tokens >= prefill_tokens.len(); @@ -593,8 +596,8 @@ impl StageOpenAiBackend { let mut decode_runtime_lock_hold_ms = 0.0; let mut decode_runtime_lock_hold_max_ms = 0.0_f64; let mut decode_runtime_lock_acquires = 0usize; - let mut decode_runtime_sessions_before = None; - let mut decode_runtime_sessions_after = None; + let mut decode_batch_size_max = 1usize; + let mut decode_batch_wait_ms = 0.0; let mut decode_forward_write_ms = 0.0; let mut decode_forward_activation_encode_ms = 0.0; let mut decode_output_activation_bytes = 0usize; @@ -618,6 +621,27 @@ impl StageOpenAiBackend { }, )?; let mut fused_reached_stop = false; + let mut native_mtp = NativeMtpN1Verifier::default(); + let native_mtp_batched_verify = native_mtp_batched_verify_enabled(); + let native_mtp_reject_cooldown_tokens = native_mtp_reject_cooldown_tokens(); + let native_mtp_defer_reject_trim = native_mtp_defer_reject_trim_enabled(); + let native_mtp_suppress_cooldown_drafts = native_mtp_suppress_cooldown_drafts_enabled(); + let native_mtp_suppress_cooldown_draft_limit = + native_mtp_suppress_cooldown_draft_limit(); + let mut native_mtp_reject_cooldown_remaining = 0usize; + let mut native_mtp_suppress_cooldown_drafts_remaining = 0usize; + let mut native_mtp_deferred_reject_trim_count = 0usize; + let mut native_mtp_deferred_reject_trim_local_ms = 0.0_f64; + let mut native_mtp_suppressed_cooldown_draft_count = 0usize; + let mut native_mtp_batched_verification_count = 0usize; + let mut native_mtp_initial_serial_verification_count = 0usize; + let mut native_mtp_initial_serial_accepted_count = 0usize; + let mut native_mtp_serial_after_gap_verification_count = 0usize; + let mut native_mtp_serial_after_gap_accepted_count = 0usize; + let mut native_mtp_verify_next_verification_count = 0usize; + let mut native_mtp_verify_next_accepted_count = 0usize; + let mut native_mtp_verify_next_draft_available_count = 0usize; + let mut native_mtp_verify_next_draft_adopted_count = 0usize; if let Some(fused) = fused_first_decode.take() { current = fused.predicted; decoded_tokens = fused.predicted_tokens.len(); @@ -640,6 +664,16 @@ impl StageOpenAiBackend { current = token; exact_replay_tokens.push(current); context_tokens.push(current); + let native_mtp_decision = if index == 0 { + native_mtp.observe_target_token( + current, + ms_to_us(fused.execution.downstream_wait_ms), + fused.native_mtp_draft, + NativeMtpDraftOrigin::InitialSerial, + ) + } else { + NativeMtpVerification::NoPending + }; if self.telemetry.is_debug_enabled() { let mut token_attrs = self.openai_attrs(request.ids); token_attrs.insert("llama_stage.decode_step".to_string(), json!(index)); @@ -729,6 +763,10 @@ impl StageOpenAiBackend { ); token_attrs .insert("llama_stage.predicted_token".to_string(), json!(current)); + token_attrs.insert( + "llama_stage.native_mtp.verification".to_string(), + json!(native_mtp_decision.label()), + ); self.telemetry .emit_debug("stage.openai_decode_token", token_attrs); } @@ -789,6 +827,9 @@ impl StageOpenAiBackend { if fused_reached_stop { break; } + if decoded_tokens >= request.max_tokens as usize { + break; + } if request .cancellation .is_some_and(openai_frontend::CancellationToken::is_cancelled) @@ -796,8 +837,300 @@ impl StageOpenAiBackend { break; } let token_timer = PhaseTimer::start(); + let native_mtp_remaining = + (request.max_tokens as usize).saturating_sub(decoded_tokens); + let can_run_native_mtp_batched_verify = native_mtp_batched_verify + && native_mtp_reject_cooldown_remaining == 0 + && draft_guard.is_none() + && native_mtp_remaining >= 2; + let pending_native_mtp_draft = can_run_native_mtp_batched_verify + .then(|| native_mtp.take_pending_draft()) + .flatten(); + if let Some(pending_native_mtp_draft) = pending_native_mtp_draft { + let batched_token_timer = + self.telemetry.is_debug_enabled().then(PhaseTimer::start); + let native_mtp_draft_token = pending_native_mtp_draft.token; + let native_mtp_draft_origin = pending_native_mtp_draft.origin; + let verify_inputs = [current, native_mtp_draft_token]; + let message = embedded_verify_message( + request.wire_dtype, + VerifySpanMessageArgs { + request_id, + session_id, + prompt_token_count: request.prompt_token_ids.len(), + pos_start: prefill_token_count + decoded_tokens, + decode_step: decoded_tokens, + tokens: &verify_inputs, + sampling: wire_sampling.clone(), + checkpoint: false, + }, + )?; + let verify = self.execute_embedded_stage_message( + &request, + downstream, + &session_key, + &message, + &verify_inputs, + WireReplyKind::PredictedTokens, + )?; + if verify.reply.predicted_tokens.len() < verify_inputs.len() { + return Err(OpenAiError::backend(format!( + "native MTP verify span returned too few tokens: got {} expected {}", + verify.reply.predicted_tokens.len(), + verify_inputs.len() + ))); + } + let target_token = verify.reply.predicted_tokens[0]; + let after_draft_token = verify.reply.predicted_tokens[1]; + let verify_next_mtp_draft = NativeMtpDraft::from_verify_prediction_tokens( + &verify.reply.predicted_tokens, + verify_inputs.len(), + ); + let native_mtp_decision = native_mtp.observe_taken_draft_verification( + native_mtp_draft_token, + target_token, + ms_to_us(verify.elapsed_ms), + ); + let accepted = + matches!(native_mtp_decision, NativeMtpVerification::Accepted { .. }); + native_mtp_batched_verification_count += 1; + match native_mtp_draft_origin { + NativeMtpDraftOrigin::InitialSerial => { + native_mtp_initial_serial_verification_count += 1; + if accepted { + native_mtp_initial_serial_accepted_count += 1; + } + } + NativeMtpDraftOrigin::SerialAfterGap => { + native_mtp_serial_after_gap_verification_count += 1; + if accepted { + native_mtp_serial_after_gap_accepted_count += 1; + } + } + NativeMtpDraftOrigin::VerifyNext => { + native_mtp_verify_next_verification_count += 1; + if accepted { + native_mtp_verify_next_accepted_count += 1; + } + } + } + let commit_tokens = [target_token, after_draft_token]; + let commit_token_count = if accepted { 2 } else { 1 }; + let consumed_positions = verify_inputs.len(); + let mut committed_positions = 0usize; + let mut reached_stop = false; + for token in commit_tokens.into_iter().take(commit_token_count) { + current = token; + decoded_tokens += 1; + committed_positions += 1; + exact_replay_tokens.push(current); + context_tokens.push(current); + if on_token(current)? == TokenControl::Stop { + reached_stop = true; + break; + } + if decoded_tokens >= request.max_tokens as usize { + break; + } + } + if !accepted && native_mtp_reject_cooldown_tokens > 0 { + native_mtp_reject_cooldown_remaining = native_mtp_reject_cooldown_tokens; + native_mtp_suppress_cooldown_drafts_remaining = + native_mtp_suppress_cooldown_draft_limit; + native_mtp.clear_pending_draft(); + } + let verify_next_mtp_draft_available = verify_next_mtp_draft.is_some(); + let verify_next_mtp_draft_adopted = accepted + && committed_positions == consumed_positions + && !reached_stop + && decoded_tokens < request.max_tokens as usize + && verify_next_mtp_draft.is_some(); + if verify_next_mtp_draft_available { + native_mtp_verify_next_draft_available_count += 1; + } + if verify_next_mtp_draft_adopted { + native_mtp_verify_next_draft_adopted_count += 1; + } + if verify_next_mtp_draft_adopted { + native_mtp.observe_next_draft( + verify_next_mtp_draft, + NativeMtpDraftOrigin::VerifyNext, + ); + } + let mut trim_control = None; + if committed_positions < consumed_positions { + let target_token_count = prefill_token_count + decoded_tokens; + let defer_trim = native_mtp_defer_reject_trim && !accepted && !reached_stop; + let trim = if defer_trim { + let trim = self.trim_embedded_stage_session_local( + &session_key, + target_token_count, + )?; + native_mtp_deferred_reject_trim_count += 1; + native_mtp_deferred_reject_trim_local_ms += trim.local_ms; + trim + } else { + self.trim_embedded_stage_session( + &request, + downstream, + &session_key, + request_id, + session_id, + target_token_count, + )? + }; + trim_control = Some(trim); + } + decode_stage0_compute_ms += verify.stats.stage0_compute_ms; + decode_runtime_lock_wait_ms += verify.stats.runtime_lock_wait_ms; + decode_runtime_lock_wait_max_ms = + decode_runtime_lock_wait_max_ms.max(verify.stats.runtime_lock_wait_ms); + decode_runtime_lock_hold_ms += verify.stats.runtime_lock_hold_ms; + decode_runtime_lock_hold_max_ms = + decode_runtime_lock_hold_max_ms.max(verify.stats.runtime_lock_hold_ms); + decode_runtime_lock_acquires += 1; + decode_forward_activation_encode_ms += verify.stats.activation_encode_ms; + decode_output_activation_bytes = decode_output_activation_bytes + .saturating_add(verify.stats.output_activation_bytes); + decode_forward_activation_bytes = decode_forward_activation_bytes + .saturating_add(verify.stats.forward_activation_bytes); + decode_forward_write_ms += verify.stats.forward_write_ms; + decode_downstream_wait_ms += verify.stats.downstream_wait_ms; + if let Some(batched_token_timer) = batched_token_timer { + let mut token_attrs = self.openai_attrs(request.ids); + token_attrs + .insert("llama_stage.decode_step".to_string(), json!(decode_step)); + token_attrs + .insert("llama_stage.message_kind".to_string(), json!("VerifySpan")); + token_attrs.insert( + "llama_stage.native_mtp.batched_verification".to_string(), + json!(true), + ); + token_attrs.insert( + "llama_stage.native_mtp.verification".to_string(), + json!(native_mtp_decision.label()), + ); + token_attrs.insert( + "llama_stage.native_mtp.verify_elapsed_ms".to_string(), + json!(verify.elapsed_ms), + ); + token_attrs.insert( + "llama_stage.native_mtp.draft_token".to_string(), + json!(native_mtp_draft_token), + ); + token_attrs.insert( + "llama_stage.native_mtp.pending_origin".to_string(), + json!(native_mtp_draft_origin.label()), + ); + token_attrs.insert( + "llama_stage.native_mtp.target_token".to_string(), + json!(target_token), + ); + token_attrs.insert( + "llama_stage.native_mtp.after_draft_token".to_string(), + json!(after_draft_token), + ); + token_attrs.insert( + "llama_stage.native_mtp.verify_next_draft_available".to_string(), + json!(verify_next_mtp_draft_available), + ); + token_attrs.insert( + "llama_stage.native_mtp.verify_next_draft_adopted".to_string(), + json!(verify_next_mtp_draft_adopted), + ); + if let Some(next_draft) = verify_next_mtp_draft { + token_attrs.insert( + "llama_stage.native_mtp.verify_next_draft_token".to_string(), + json!(next_draft.token), + ); + token_attrs.insert( + "llama_stage.native_mtp.verify_next_draft_compute_us".to_string(), + json!(next_draft.proposal_compute_us), + ); + } + token_attrs.insert( + "llama_stage.native_mtp.consumed_positions".to_string(), + json!(consumed_positions), + ); + token_attrs.insert( + "llama_stage.native_mtp.committed_positions".to_string(), + json!(committed_positions), + ); + token_attrs.insert( + "llama_stage.native_mtp.reject_cooldown_tokens".to_string(), + json!(native_mtp_reject_cooldown_tokens), + ); + token_attrs.insert( + "llama_stage.native_mtp.reject_cooldown_remaining".to_string(), + json!(native_mtp_reject_cooldown_remaining), + ); + token_attrs.insert( + "llama_stage.native_mtp.defer_reject_trim".to_string(), + json!(native_mtp_defer_reject_trim), + ); + if let Some(trim) = trim_control.as_ref() { + token_attrs.insert( + "llama_stage.native_mtp.trim_ms".to_string(), + json!(trim.elapsed_ms), + ); + token_attrs.insert( + "llama_stage.native_mtp.trim_local_ms".to_string(), + json!(trim.local_ms), + ); + token_attrs.insert( + "llama_stage.native_mtp.trim_downstream_write_ms".to_string(), + json!(trim.downstream_write_ms), + ); + token_attrs.insert( + "llama_stage.native_mtp.trim_downstream_wait_ms".to_string(), + json!(trim.downstream_wait_ms), + ); + } + token_attrs.insert( + "llama_stage.stage0_compute_ms".to_string(), + json!(verify.stats.stage0_compute_ms), + ); + token_attrs.insert( + "llama_stage.runtime_lock_wait_ms".to_string(), + json!(verify.stats.runtime_lock_wait_ms), + ); + token_attrs.insert( + "llama_stage.runtime_lock_hold_ms".to_string(), + json!(verify.stats.runtime_lock_hold_ms), + ); + token_attrs.insert( + "llama_stage.activation_encode_ms".to_string(), + json!(verify.stats.activation_encode_ms), + ); + token_attrs.insert( + "llama_stage.forward_write_ms".to_string(), + json!(verify.stats.forward_write_ms), + ); + token_attrs.insert( + "llama_stage.downstream_wait_ms".to_string(), + json!(verify.stats.downstream_wait_ms), + ); + token_attrs.insert( + "llama_stage.output_activation_bytes".to_string(), + json!(verify.stats.output_activation_bytes), + ); + token_attrs.insert( + "llama_stage.forward_activation_bytes".to_string(), + json!(verify.stats.forward_activation_bytes), + ); + self.emit_openai_phase( + "stage.openai_native_mtp_verify", + batched_token_timer, + token_attrs, + ); + } + if reached_stop { + break; + } + continue; + } if draft_guard.is_some() { - let remaining = request.max_tokens as usize - decoded_tokens; + let remaining = (request.max_tokens as usize).saturating_sub(decoded_tokens); if remaining == 0 { break; } @@ -805,8 +1138,8 @@ impl StageOpenAiBackend { let proposal_limit = remaining.min(adaptive_window); let propose_timer = PhaseTimer::start(); let mut draft_tokens = Vec::new(); - if draft_tokens.is_empty() - && let Some(draft) = draft_guard.as_deref_mut() + if let (true, Some(draft)) = + (draft_tokens.is_empty(), draft_guard.as_deref_mut()) { let proposal_limit = proposal_limit.min(draft.window); draft_tokens = draft @@ -829,6 +1162,7 @@ impl StageOpenAiBackend { pos_start: prefill_token_count + decoded_tokens, decode_step: decoded_tokens, tokens: &verify_inputs, + sampling: wire_sampling.clone(), checkpoint: true, }, )?; @@ -968,6 +1302,7 @@ impl StageOpenAiBackend { pos_start: prefill_token_count + decoded_tokens, decode_step: decoded_tokens, tokens: repair_inputs, + sampling: wire_sampling.clone(), checkpoint: false, }, )?; @@ -1127,45 +1462,27 @@ impl StageOpenAiBackend { decode_message.update(decode_step_index, current)? }; let stage0_timer = PhaseTimer::start(); - let token_runtime_lock_wait_ms; - let token_runtime_lock_hold_ms; - let output = { - let lock_timer = PhaseTimer::start(); - let mut runtime = self - .runtime - .lock() - .map_err(|_| OpenAiError::backend("runtime lock poisoned"))?; - let lock_wait_ms = lock_timer.elapsed_ms(); - token_runtime_lock_wait_ms = lock_wait_ms; - decode_runtime_lock_wait_ms += lock_wait_ms; - decode_runtime_lock_wait_max_ms = - decode_runtime_lock_wait_max_ms.max(lock_wait_ms); - decode_runtime_lock_acquires += 1; - let lock_hold_timer = PhaseTimer::start(); - decode_runtime_sessions_before.get_or_insert_with(|| runtime.session_stats()); - let output = run_binary_stage_message( - &mut runtime, + let batch_outcome = self + .decode_frame_batcher + .decode( &session_key, - message, - &[current], + current, + request.sampling.enabled.then_some(request.sampling), None, - false, - stage_output_activation_capacity( - request.config, - message.token_count, - request.activation_width, - ) - .map_err(openai_backend_error)?, ) - .map_err(openai_backend_error)? - .2; - decode_runtime_sessions_after = Some(runtime.session_stats()); - token_runtime_lock_hold_ms = lock_hold_timer.elapsed_ms(); - decode_runtime_lock_hold_ms += token_runtime_lock_hold_ms; - decode_runtime_lock_hold_max_ms = - decode_runtime_lock_hold_max_ms.max(token_runtime_lock_hold_ms); - output - }; + .map_err(openai_backend_error)?; + let token_runtime_lock_wait_ms = batch_outcome.runtime_lock_wait_ms; + let token_runtime_lock_hold_ms = batch_outcome.runtime_lock_hold_ms; + decode_runtime_lock_wait_ms += token_runtime_lock_wait_ms; + decode_runtime_lock_wait_max_ms = + decode_runtime_lock_wait_max_ms.max(token_runtime_lock_wait_ms); + decode_runtime_lock_hold_ms += token_runtime_lock_hold_ms; + decode_runtime_lock_hold_max_ms = + decode_runtime_lock_hold_max_ms.max(token_runtime_lock_hold_ms); + decode_runtime_lock_acquires += 1; + decode_batch_size_max = decode_batch_size_max.max(batch_outcome.batch_size); + decode_batch_wait_ms += batch_outcome.batch_wait_ms; + let output = batch_outcome.output; let stage0_compute_ms = stage0_timer.elapsed_ms(); decode_stage0_compute_ms += stage0_compute_ms; let forwarded = forwarded_stage_message_timed( @@ -1192,15 +1509,14 @@ impl StageOpenAiBackend { let forward_write_ms = write_timer.elapsed_ms(); decode_forward_write_ms += forward_write_ms; let wait_timer = PhaseTimer::start(); - let reply = request - .prediction_return - .as_ref() - .ok_or_else(|| { - OpenAiError::backend("missing direct prediction return receiver") - })? - .recv_expected(WireReplyKind::PredictedToken) - .map_err(openai_backend_error)?; + let reply = recv_reply(&mut *downstream).map_err(openai_io_error)?; let downstream_wait_ms = wait_timer.elapsed_ms(); + if reply.kind != WireReplyKind::PredictedToken { + return Err(OpenAiError::backend(format!( + "expected decode PredictedToken reply from downstream, got {:?}", + reply.kind + ))); + } decode_downstream_wait_ms += downstream_wait_ms; if records_replay_checkpoint && super::prefix_cache::request_allows_exact_replay(&request) @@ -1226,6 +1542,35 @@ impl StageOpenAiBackend { )?; } current = reply.predicted; + let suppress_cooldown_draft_broad = + native_mtp_suppress_cooldown_drafts && native_mtp_reject_cooldown_remaining > 0; + let suppress_cooldown_draft_limited = native_mtp_reject_cooldown_remaining > 0 + && native_mtp_suppress_cooldown_drafts_remaining > 0; + let suppress_cooldown_draft = + suppress_cooldown_draft_broad || suppress_cooldown_draft_limited; + let native_mtp_draft = if suppress_cooldown_draft { + None + } else { + NativeMtpDraft::from_prediction_tokens(&reply.predicted_tokens) + }; + if suppress_cooldown_draft { + native_mtp.clear_pending_draft(); + native_mtp_suppressed_cooldown_draft_count += 1; + native_mtp_suppress_cooldown_drafts_remaining = + native_mtp_suppress_cooldown_drafts_remaining.saturating_sub(1); + } + let native_mtp_decision = native_mtp.observe_target_token( + current, + ms_to_us(downstream_wait_ms), + native_mtp_draft, + if native_mtp_batched_verification_count == 0 { + NativeMtpDraftOrigin::InitialSerial + } else { + NativeMtpDraftOrigin::SerialAfterGap + }, + ); + native_mtp_reject_cooldown_remaining = + native_mtp_reject_cooldown_remaining.saturating_sub(1); decoded_tokens += 1; exact_replay_tokens.push(current); context_tokens.push(current); @@ -1248,6 +1593,14 @@ impl StageOpenAiBackend { "llama_stage.runtime_lock_hold_ms".to_string(), json!(token_runtime_lock_hold_ms), ); + token_attrs.insert( + "llama_stage.decode_batch_size".to_string(), + json!(batch_outcome.batch_size), + ); + token_attrs.insert( + "llama_stage.decode_batch_wait_ms".to_string(), + json!(batch_outcome.batch_wait_ms), + ); token_attrs.insert( "llama_stage.output_activation_bytes".to_string(), json!(output.payload.len()), @@ -1270,6 +1623,22 @@ impl StageOpenAiBackend { ); token_attrs.insert("llama_stage.predicted_token".to_string(), json!(current)); token_attrs.insert("llama_stage.message_kind".to_string(), json!("DecodeEmbd")); + token_attrs.insert( + "llama_stage.native_mtp.verification".to_string(), + json!(native_mtp_decision.label()), + ); + token_attrs.insert( + "llama_stage.native_mtp.suppress_cooldown_drafts".to_string(), + json!(native_mtp_suppress_cooldown_drafts), + ); + token_attrs.insert( + "llama_stage.native_mtp.suppress_cooldown_draft_limit".to_string(), + json!(native_mtp_suppress_cooldown_draft_limit), + ); + token_attrs.insert( + "llama_stage.native_mtp.cooldown_draft_suppressed".to_string(), + json!(suppress_cooldown_draft), + ); self.emit_openai_phase("stage.openai_decode_token", token_timer, token_attrs); } if on_token(current)? == TokenControl::Stop { @@ -1305,20 +1674,14 @@ impl StageOpenAiBackend { "llama_stage.runtime_lock_acquires".to_string(), json!(decode_runtime_lock_acquires), ); - if let Some(stats) = decode_runtime_sessions_before.as_ref() { - Self::insert_runtime_session_stats( - &mut decode_attrs, - "llama_stage.runtime_sessions_before", - stats, - ); - } - if let Some(stats) = decode_runtime_sessions_after.as_ref() { - Self::insert_runtime_session_stats( - &mut decode_attrs, - "llama_stage.runtime_sessions_after", - stats, - ); - } + decode_attrs.insert( + "llama_stage.decode_batch_size_max".to_string(), + json!(decode_batch_size_max), + ); + decode_attrs.insert( + "llama_stage.decode_batch_wait_ms".to_string(), + json!(decode_batch_wait_ms), + ); decode_attrs.insert( "llama_stage.forward_write_ms".to_string(), json!(decode_forward_write_ms), @@ -1340,7 +1703,72 @@ impl StageOpenAiBackend { json!(decode_downstream_wait_ms), ); speculative_stats.insert_attrs(&mut decode_attrs); - self.emit_openai_phase("stage.openai_decode", decode_timer, decode_attrs); + native_mtp.stats().insert_attrs(&mut decode_attrs); + decode_attrs.insert( + "llama_stage.native_mtp.reject_cooldown_tokens".to_string(), + json!(native_mtp_reject_cooldown_tokens), + ); + decode_attrs.insert( + "llama_stage.native_mtp.defer_reject_trim".to_string(), + json!(native_mtp_defer_reject_trim), + ); + decode_attrs.insert( + "llama_stage.native_mtp.suppress_cooldown_drafts".to_string(), + json!(native_mtp_suppress_cooldown_drafts), + ); + decode_attrs.insert( + "llama_stage.native_mtp.suppress_cooldown_draft_limit".to_string(), + json!(native_mtp_suppress_cooldown_draft_limit), + ); + decode_attrs.insert( + "llama_stage.native_mtp.suppressed_cooldown_draft_count".to_string(), + json!(native_mtp_suppressed_cooldown_draft_count), + ); + decode_attrs.insert( + "llama_stage.native_mtp.batched_verification_count".to_string(), + json!(native_mtp_batched_verification_count), + ); + decode_attrs.insert( + "llama_stage.native_mtp.initial_serial_verification_count".to_string(), + json!(native_mtp_initial_serial_verification_count), + ); + decode_attrs.insert( + "llama_stage.native_mtp.initial_serial_accepted_count".to_string(), + json!(native_mtp_initial_serial_accepted_count), + ); + decode_attrs.insert( + "llama_stage.native_mtp.serial_after_gap_verification_count".to_string(), + json!(native_mtp_serial_after_gap_verification_count), + ); + decode_attrs.insert( + "llama_stage.native_mtp.serial_after_gap_accepted_count".to_string(), + json!(native_mtp_serial_after_gap_accepted_count), + ); + decode_attrs.insert( + "llama_stage.native_mtp.verify_next_verification_count".to_string(), + json!(native_mtp_verify_next_verification_count), + ); + decode_attrs.insert( + "llama_stage.native_mtp.verify_next_accepted_count".to_string(), + json!(native_mtp_verify_next_accepted_count), + ); + decode_attrs.insert( + "llama_stage.native_mtp.deferred_reject_trim_count".to_string(), + json!(native_mtp_deferred_reject_trim_count), + ); + decode_attrs.insert( + "llama_stage.native_mtp.deferred_reject_trim_local_ms".to_string(), + json!(native_mtp_deferred_reject_trim_local_ms), + ); + decode_attrs.insert( + "llama_stage.native_mtp.verify_next_draft_available_count".to_string(), + json!(native_mtp_verify_next_draft_available_count), + ); + decode_attrs.insert( + "llama_stage.native_mtp.verify_next_draft_adopted_count".to_string(), + json!(native_mtp_verify_next_draft_adopted_count), + ); + self.emit_openai_summary("stage.openai_decode", decode_timer, decode_attrs); Ok(()) })(); diff --git a/crates/skippy-server/src/frontend/generation_flow.rs b/crates/skippy-server/src/frontend/generation_flow.rs index 80aad67b79..cd17745b92 100644 --- a/crates/skippy-server/src/frontend/generation_flow.rs +++ b/crates/skippy-server/src/frontend/generation_flow.rs @@ -33,7 +33,12 @@ impl StageOpenAiBackend { on_text_chunk, ); } - let stop_values = stop.map(|stop| stop.values()).unwrap_or_default(); + let stop_value_storage = + generation_stop_values(stop, prompt.chat_parse_metadata.as_deref()); + let stop_values = stop_value_storage + .iter() + .map(String::as_str) + .collect::>(); let tokenize_timer = PhaseTimer::start(); let prompt_token_ids = self.tokenize(&prompt.text)?; let mut tokenize_attrs = self.openai_attrs(&ids); @@ -109,7 +114,6 @@ impl StageOpenAiBackend { downstream_wire_condition, prefill_reply_credit_limit, lane_pool, - prediction_returns, } => self.generate_embedded_stage_zero_tokens( EmbeddedStageZeroGeneration { config: &config, @@ -119,11 +123,6 @@ impl StageOpenAiBackend { downstream_wire_condition, prefill_reply_credit_limit, lane_pool, - prediction_return: prediction_returns - .as_ref() - .map(|hub| hub.register(ids.request_id, ids.session_id)) - .transpose() - .map_err(openai_backend_error)?, draft: self.draft.clone(), speculative_window: self.speculative_window, adaptive_speculative_window: self.adaptive_speculative_window, @@ -200,42 +199,36 @@ impl StageOpenAiBackend { ids: OpenAiGenerationIds, on_text_chunk: impl FnMut(&str) -> OpenAiResult<()>, ) -> OpenAiResult { - if let OpenAiBackendMode::EmbeddedStageZero { - config, - wire_dtype, - activation_width, - downstream_wire_condition, - lane_pool, - prediction_returns, - .. - } = self.mode.clone() - && config.downstream.is_some() - { - let lane_pool = lane_pool.ok_or_else(|| { - OpenAiError::backend("embedded stage 0 has no downstream lane pool") - })?; - let prediction_return = prediction_returns - .as_ref() - .map(|hub| hub.register(ids.request_id, ids.session_id)) - .transpose() - .map_err(openai_backend_error)?; - return self.generate_split_multimodal_text( - SplitMultimodalGeneration { - prompt, - max_tokens, - stop, - sampling, - cancellation, - ids, - config, - wire_dtype, - activation_width, - downstream_wire_condition, - lane_pool, - prediction_return, - }, - on_text_chunk, - ); + match self.mode.clone() { + OpenAiBackendMode::EmbeddedStageZero { + config, + wire_dtype, + activation_width, + downstream_wire_condition, + lane_pool, + .. + } if config.downstream.is_some() => { + let lane_pool = lane_pool.ok_or_else(|| { + OpenAiError::backend("embedded stage 0 has no downstream lane pool") + })?; + return self.generate_split_multimodal_text( + SplitMultimodalGeneration { + prompt, + max_tokens, + stop, + sampling, + cancellation, + ids, + config, + wire_dtype, + activation_width, + downstream_wire_condition, + lane_pool, + }, + on_text_chunk, + ); + } + _ => {} } match &self.mode { @@ -248,7 +241,12 @@ impl StageOpenAiBackend { } } - let stop_values = stop.map(|stop| stop.values()).unwrap_or_default(); + let stop_value_storage = + generation_stop_values(stop, prompt.chat_parse_metadata.as_deref()); + let stop_values = stop_value_storage + .iter() + .map(String::as_str) + .collect::>(); let session_id = ids.session_label.clone(); let prefill_timer = PhaseTimer::start(); let (prefill, mut token_signal, mut signal_window) = { @@ -502,7 +500,7 @@ impl StageOpenAiBackend { stats, ); } - self.emit_openai_phase("stage.openai_decode", decode_timer, attrs); + self.emit_openai_summary("stage.openai_decode", decode_timer, attrs); Ok(()) })(); let lock_timer = PhaseTimer::start(); @@ -547,7 +545,12 @@ impl StageOpenAiBackend { request: SplitMultimodalGeneration<'_>, on_text_chunk: impl FnMut(&str) -> OpenAiResult<()>, ) -> OpenAiResult { - let stop_values = request.stop.map(|stop| stop.values()).unwrap_or_default(); + let stop_value_storage = + generation_stop_values(request.stop, request.prompt.chat_parse_metadata.as_deref()); + let stop_values = stop_value_storage + .iter() + .map(String::as_str) + .collect::>(); let mut collector = TextGenerationCollector::new(self.runtime.clone(), stop_values, on_text_chunk); let wire_sampling = wire_sampling_config(&request.sampling); @@ -692,18 +695,7 @@ impl StageOpenAiBackend { .map_err(openai_io_error)?; forward_write_ms += write_timer.elapsed_ms(); let wait_timer = PhaseTimer::start(); - let reply = if is_final_chunk { - request - .prediction_return - .as_ref() - .ok_or_else(|| { - OpenAiError::backend("missing direct prediction return receiver") - })? - .recv_expected(WireReplyKind::PredictedToken) - .map_err(openai_backend_error)? - } else { - recv_reply(&mut lane.stream).map_err(openai_io_error)? - }; + let reply = recv_reply(&mut lane.stream).map_err(openai_io_error)?; downstream_wait_ms += wait_timer.elapsed_ms(); let expected = if is_final_chunk { WireReplyKind::PredictedToken @@ -760,6 +752,8 @@ impl StageOpenAiBackend { let mut decode_runtime_lock_wait_ms = 0.0; let mut decode_runtime_lock_hold_ms = 0.0; let mut decode_runtime_lock_acquires = 0usize; + let mut decode_batch_size_max = 1usize; + let mut decode_batch_wait_ms = 0.0; let mut decode_forward_write_ms = 0.0; let mut decode_downstream_wait_ms = 0.0; let mut decode_output_activation_bytes = 0usize; @@ -796,35 +790,16 @@ impl StageOpenAiBackend { let message = decode_message.update(decode_input_index, current)?; let token_timer = PhaseTimer::start(); let stage0_timer = PhaseTimer::start(); - let output = { - let lock_timer = PhaseTimer::start(); - let mut runtime = self - .runtime - .lock() - .map_err(|_| OpenAiError::backend("runtime lock poisoned"))?; - let lock_wait_ms = lock_timer.elapsed_ms(); - decode_runtime_lock_wait_ms += lock_wait_ms; - decode_runtime_lock_acquires += 1; - let hold_timer = PhaseTimer::start(); - let output = run_binary_stage_message( - &mut runtime, - &session_key, - message, - &[current], - None, - false, - stage_output_activation_capacity( - &request.config, - message.token_count, - request.activation_width, - ) - .map_err(openai_backend_error)?, - ) - .map_err(openai_backend_error)? - .2; - decode_runtime_lock_hold_ms += hold_timer.elapsed_ms(); - output - }; + let batch_outcome = self + .decode_frame_batcher + .decode(&session_key, current, Some(&request.sampling), None) + .map_err(openai_backend_error)?; + decode_runtime_lock_wait_ms += batch_outcome.runtime_lock_wait_ms; + decode_runtime_lock_hold_ms += batch_outcome.runtime_lock_hold_ms; + decode_runtime_lock_acquires += 1; + decode_batch_size_max = decode_batch_size_max.max(batch_outcome.batch_size); + decode_batch_wait_ms += batch_outcome.batch_wait_ms; + let output = batch_outcome.output; let stage0_compute_ms = stage0_timer.elapsed_ms(); decode_stage0_compute_ms += stage0_compute_ms; let forwarded = forwarded_stage_message_timed( @@ -850,15 +825,14 @@ impl StageOpenAiBackend { let forward_write_ms = write_timer.elapsed_ms(); decode_forward_write_ms += forward_write_ms; let wait_timer = PhaseTimer::start(); - let reply = request - .prediction_return - .as_ref() - .ok_or_else(|| { - OpenAiError::backend("missing direct prediction return receiver") - })? - .recv_expected(WireReplyKind::PredictedToken) - .map_err(openai_backend_error)?; + let reply = recv_reply(&mut lane.stream).map_err(openai_io_error)?; let downstream_wait_ms = wait_timer.elapsed_ms(); + if reply.kind != WireReplyKind::PredictedToken { + return Err(OpenAiError::backend(format!( + "expected split multimodal decode PredictedToken reply from downstream, got {:?}", + reply.kind + ))); + } decode_downstream_wait_ms += downstream_wait_ms; current = reply.predicted; if self.telemetry.is_debug_enabled() { @@ -879,6 +853,14 @@ impl StageOpenAiBackend { "llama_stage.downstream_wait_ms".to_string(), json!(downstream_wait_ms), ); + token_attrs.insert( + "llama_stage.decode_batch_size".to_string(), + json!(batch_outcome.batch_size), + ); + token_attrs.insert( + "llama_stage.decode_batch_wait_ms".to_string(), + json!(batch_outcome.batch_wait_ms), + ); token_attrs.insert("llama_stage.predicted_token".to_string(), json!(current)); token_attrs.insert("llama_stage.message_kind".to_string(), json!("DecodeEmbd")); self.emit_openai_phase("stage.openai_decode_token", token_timer, token_attrs); @@ -906,6 +888,14 @@ impl StageOpenAiBackend { "llama_stage.runtime_lock_acquires".to_string(), json!(decode_runtime_lock_acquires), ); + decode_attrs.insert( + "llama_stage.decode_batch_size_max".to_string(), + json!(decode_batch_size_max), + ); + decode_attrs.insert( + "llama_stage.decode_batch_wait_ms".to_string(), + json!(decode_batch_wait_ms), + ); decode_attrs.insert( "llama_stage.forward_write_ms".to_string(), json!(decode_forward_write_ms), @@ -922,7 +912,7 @@ impl StageOpenAiBackend { "llama_stage.forward_activation_bytes".to_string(), json!(decode_forward_activation_bytes), ); - self.emit_openai_phase("stage.openai_decode", decode_timer, decode_attrs); + self.emit_openai_summary("stage.openai_decode", decode_timer, decode_attrs); Ok(()) })(); diff --git a/crates/skippy-server/src/frontend/local_generation.rs b/crates/skippy-server/src/frontend/local_generation.rs index 52749d9306..4046db356a 100644 --- a/crates/skippy-server/src/frontend/local_generation.rs +++ b/crates/skippy-server/src/frontend/local_generation.rs @@ -9,7 +9,80 @@ impl StageOpenAiBackend { let session_id = request.ids.session_label.clone(); let mut cache_stats = GenerationCacheStats::default(); let result = (|| { - if request.prompt_token_ids.len() > 1 { + let mut prompt_prefill_sample = None; + let mut chat_sampling_configured = false; + if request.max_tokens > 0 && request.prompt_token_ids.len() > 1 && self.kv.is_none() { + if let Some(metadata) = request.chat_sampling_metadata { + let mut runtime = self + .runtime + .lock() + .map_err(|_| OpenAiError::backend("runtime lock poisoned"))?; + runtime + .configure_chat_sampling( + &session_id, + metadata, + request.prompt_token_ids.len() as u64, + request.sampling.enabled.then_some(request.sampling), + ) + .map_err(openai_backend_error)?; + chat_sampling_configured = true; + } + let prefill_timer = PhaseTimer::start(); + let lock_timer = PhaseTimer::start(); + let mut runtime = self + .runtime + .lock() + .map_err(|_| OpenAiError::backend("runtime lock poisoned"))?; + let runtime_lock_wait_ms = lock_timer.elapsed_ms(); + let runtime_lock_hold_timer = PhaseTimer::start(); + let runtime_sessions_before = runtime.session_stats(); + let (predicted, _) = runtime + .prefill_final_frame_sampled( + &session_id, + request.prompt_token_ids, + &[], + request.sampling.enabled.then_some(request.sampling), + None, + ) + .map_err(openai_backend_error)?; + prompt_prefill_sample = Some(predicted); + cache_stats.suffix_prefill_tokens = saturating_u32(request.prompt_token_ids.len()); + let runtime_sessions_after = runtime.session_stats(); + let runtime_lock_hold_ms = runtime_lock_hold_timer.elapsed_ms(); + let mut attrs = self.openai_attrs(request.ids); + attrs.insert( + "llama_stage.prefill_token_count".to_string(), + json!(request.prompt_token_ids.len()), + ); + attrs.insert("llama_stage.prefill_chunk_count".to_string(), json!(1)); + attrs.insert("skippy.kv.restored_prefill".to_string(), json!(false)); + attrs.insert("skippy.kv.restored_prefill_tokens".to_string(), json!(0)); + attrs.insert( + "skippy.kv.prefill_suffix_tokens".to_string(), + json!(request.prompt_token_ids.len()), + ); + attrs.insert("skippy.kv.recorded_pages".to_string(), json!(0)); + attrs.insert( + "llama_stage.runtime_lock_wait_ms".to_string(), + json!(runtime_lock_wait_ms), + ); + attrs.insert( + "llama_stage.runtime_lock_hold_ms".to_string(), + json!(runtime_lock_hold_ms), + ); + attrs.insert("llama_stage.runtime_lock_acquires".to_string(), json!(1)); + Self::insert_runtime_session_stats( + &mut attrs, + "llama_stage.runtime_sessions_before", + &runtime_sessions_before, + ); + Self::insert_runtime_session_stats( + &mut attrs, + "llama_stage.runtime_sessions_after", + &runtime_sessions_after, + ); + self.emit_openai_phase("stage.openai_prefill", prefill_timer, attrs); + } else if request.prompt_token_ids.len() > 1 { let prefill_timer = PhaseTimer::start(); let prefill_tokens = &request.prompt_token_ids[..request.prompt_token_ids.len() - 1]; @@ -187,9 +260,10 @@ impl StageOpenAiBackend { cache_stats.matched_prefix_tokens = saturating_u32(restored_prefill_tokens); cache_stats.suffix_prefill_tokens = saturating_u32(prefill_tokens.len().saturating_sub(restored_prefill_tokens)); - if (!restored_prefill || decoded_prefill_suffix) - && let Some(kv) = self.kv.as_ref() - { + if let (true, Some(kv)) = ( + !restored_prefill || decoded_prefill_suffix, + self.kv.as_ref(), + ) { let base = self.local_kv_message_base(&session_id, request.ids); let exact_identity = kv.prefill_identity(&self.config, &base, 0, prefill_tokens); @@ -377,7 +451,10 @@ impl StageOpenAiBackend { return Err(openai_backend_error(error)); } } - if let Some(metadata) = request.chat_sampling_metadata { + let chat_sampling_metadata = (!chat_sampling_configured) + .then_some(request.chat_sampling_metadata) + .flatten(); + if let Some(metadata) = chat_sampling_metadata { let mut runtime = self .runtime .lock() @@ -404,6 +481,12 @@ impl StageOpenAiBackend { .prompt_token_ids .last() .expect("checked non-empty prompt"); + let mut stopped = false; + if let Some(predicted) = prompt_prefill_sample { + current = predicted; + decoded_tokens += 1; + stopped = on_token(current)? == TokenControl::Stop; + } let mut hook_request = request.hook_request; let hook_runtime = request.hook_runtime; let generation_hooks_active = @@ -411,7 +494,7 @@ impl StageOpenAiBackend { let emit_token_debug = self.telemetry.is_debug_enabled(); let mut post_prefill_hook_checked = false; let mut last_mid_generation_hook_at = None; - while decoded_tokens < request.max_tokens as usize { + while !stopped && decoded_tokens < request.max_tokens as usize { if request .cancellation .is_some_and(openai_frontend::CancellationToken::is_cancelled) @@ -420,61 +503,48 @@ impl StageOpenAiBackend { } let decode_step = decoded_tokens; let token_timer = PhaseTimer::start(); - let token_runtime_lock_wait_ms; - let token_runtime_lock_hold_ms; - let token_decode_ms; let token_signal_ms; let token_signal; let signal_window; - current = { - let lock_timer = PhaseTimer::start(); + let decode_call_timer = PhaseTimer::start(); + let outcome = self.decode_batcher.decode( + &session_id, + current, + request.sampling.enabled.then_some(request.sampling), + )?; + current = outcome.predicted; + let token_batch_size = outcome.batch_size; + let token_batch_wait_ms = outcome.batch_wait_ms; + let token_runtime_lock_wait_ms = outcome.runtime_lock_wait_ms; + let token_runtime_lock_hold_ms = outcome.runtime_lock_hold_ms; + runtime_lock_wait_ms += token_runtime_lock_wait_ms; + runtime_lock_wait_max_ms = runtime_lock_wait_max_ms.max(token_runtime_lock_wait_ms); + runtime_lock_hold_ms += token_runtime_lock_hold_ms; + runtime_lock_hold_max_ms = runtime_lock_hold_max_ms.max(token_runtime_lock_hold_ms); + runtime_lock_acquires += 1; + let token_decode_ms = if emit_token_debug { + decode_call_timer.elapsed_ms() + } else { + 0.0 + }; + if generation_hooks_active { + let signal_timer = PhaseTimer::start(); let mut runtime = self .runtime .lock() .map_err(|_| OpenAiError::backend("runtime lock poisoned"))?; - let lock_wait_ms = lock_timer.elapsed_ms(); - token_runtime_lock_wait_ms = lock_wait_ms; - runtime_lock_wait_ms += lock_wait_ms; - runtime_lock_wait_max_ms = runtime_lock_wait_max_ms.max(lock_wait_ms); - runtime_lock_acquires += 1; - let hold_timer = PhaseTimer::start(); runtime_sessions_before.get_or_insert_with(|| runtime.session_stats()); - let decode_call_timer = PhaseTimer::start(); - let predicted = runtime - .decode_sampled( - &session_id, - current, - request.sampling.enabled.then_some(request.sampling), - ) - .map_err(openai_backend_error)?; - token_decode_ms = if emit_token_debug { - decode_call_timer.elapsed_ms() - } else { - 0.0 - }; - if generation_hooks_active { - let signal_timer = PhaseTimer::start(); - token_signal = runtime.last_token_signal(&session_id).ok(); - signal_window = runtime.signal_window(&session_id, 16).ok(); - token_signal_ms = signal_timer.elapsed_ms(); - } else { - token_signal = None; - signal_window = None; - token_signal_ms = 0.0; - } + token_signal = runtime.last_token_signal(&session_id).ok(); + signal_window = runtime.signal_window(&session_id, 16).ok(); runtime_sessions_after = Some(runtime.session_stats()); - token_runtime_lock_hold_ms = if emit_token_debug { - hold_timer.elapsed_ms() - } else { - 0.0 - }; - runtime_lock_hold_ms += token_runtime_lock_hold_ms; - runtime_lock_hold_max_ms = - runtime_lock_hold_max_ms.max(token_runtime_lock_hold_ms); - predicted - }; - if generation_hooks_active - && let Some(injected_current) = self.maybe_run_generation_hooks( + token_signal_ms = signal_timer.elapsed_ms(); + } else { + token_signal = None; + signal_window = None; + token_signal_ms = 0.0; + } + let injected_current = if generation_hooks_active { + self.maybe_run_generation_hooks( &session_id, &mut hook_request, hook_runtime.as_ref(), @@ -484,7 +554,10 @@ impl StageOpenAiBackend { token_signal, signal_window, )? - { + } else { + None + }; + if let Some(injected_current) = injected_current { current = injected_current; continue; } @@ -506,6 +579,14 @@ impl StageOpenAiBackend { "llama_stage.decode_call_ms".to_string(), json!(token_decode_ms), ); + token_attrs.insert( + "llama_stage.decode_batch_size".to_string(), + json!(token_batch_size), + ); + token_attrs.insert( + "llama_stage.decode_batch_wait_ms".to_string(), + json!(token_batch_wait_ms), + ); token_attrs.insert("llama_stage.signal_ms".to_string(), json!(token_signal_ms)); token_attrs.insert( "llama_stage.runtime_lock_wait_ms".to_string(), @@ -524,48 +605,46 @@ impl StageOpenAiBackend { break; } } - if emit_token_debug { - let mut attrs = self.openai_attrs(request.ids); - attrs.insert( - "llama_stage.decode_token_count".to_string(), - json!(decoded_tokens), - ); - attrs.insert( - "llama_stage.runtime_lock_wait_ms".to_string(), - json!(runtime_lock_wait_ms), - ); - attrs.insert( - "llama_stage.runtime_lock_wait_max_ms".to_string(), - json!(runtime_lock_wait_max_ms), - ); - attrs.insert( - "llama_stage.runtime_lock_hold_ms".to_string(), - json!(runtime_lock_hold_ms), - ); - attrs.insert( - "llama_stage.runtime_lock_hold_max_ms".to_string(), - json!(runtime_lock_hold_max_ms), + let mut attrs = self.openai_attrs(request.ids); + attrs.insert( + "llama_stage.decode_token_count".to_string(), + json!(decoded_tokens), + ); + attrs.insert( + "llama_stage.runtime_lock_wait_ms".to_string(), + json!(runtime_lock_wait_ms), + ); + attrs.insert( + "llama_stage.runtime_lock_wait_max_ms".to_string(), + json!(runtime_lock_wait_max_ms), + ); + attrs.insert( + "llama_stage.runtime_lock_hold_ms".to_string(), + json!(runtime_lock_hold_ms), + ); + attrs.insert( + "llama_stage.runtime_lock_hold_max_ms".to_string(), + json!(runtime_lock_hold_max_ms), + ); + attrs.insert( + "llama_stage.runtime_lock_acquires".to_string(), + json!(runtime_lock_acquires), + ); + if let Some(stats) = runtime_sessions_before.as_ref() { + Self::insert_runtime_session_stats( + &mut attrs, + "llama_stage.runtime_sessions_before", + stats, ); - attrs.insert( - "llama_stage.runtime_lock_acquires".to_string(), - json!(runtime_lock_acquires), + } + if let Some(stats) = runtime_sessions_after.as_ref() { + Self::insert_runtime_session_stats( + &mut attrs, + "llama_stage.runtime_sessions_after", + stats, ); - if let Some(stats) = runtime_sessions_before.as_ref() { - Self::insert_runtime_session_stats( - &mut attrs, - "llama_stage.runtime_sessions_before", - stats, - ); - } - if let Some(stats) = runtime_sessions_after.as_ref() { - Self::insert_runtime_session_stats( - &mut attrs, - "llama_stage.runtime_sessions_after", - stats, - ); - } - self.emit_openai_phase("stage.openai_decode", decode_timer, attrs); } + self.emit_openai_summary("stage.openai_decode", decode_timer, attrs); Ok(()) })(); let lock_timer = PhaseTimer::start(); diff --git a/crates/skippy-server/src/frontend/native_mtp/draft.rs b/crates/skippy-server/src/frontend/native_mtp/draft.rs new file mode 100644 index 0000000000..388ffcc744 --- /dev/null +++ b/crates/skippy-server/src/frontend/native_mtp/draft.rs @@ -0,0 +1,118 @@ +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(in crate::frontend) struct NativeMtpDraft { + pub(in crate::frontend) token: i32, + pub(in crate::frontend) proposal_compute_us: i64, +} + +impl NativeMtpDraft { + pub(in crate::frontend) fn from_prediction_tokens(tokens: &[i32]) -> Option { + let token = *tokens.get(1)?; + let proposal_compute_us = tokens.get(2).copied().unwrap_or_default(); + Some(Self { + token, + proposal_compute_us: i64::from(proposal_compute_us.max(0)), + }) + } + + pub(in crate::frontend) fn from_verify_prediction_tokens( + tokens: &[i32], + verified_token_count: usize, + ) -> Option { + let token = *tokens.get(verified_token_count)?; + let proposal_compute_us = tokens + .get(verified_token_count.saturating_add(1)) + .copied() + .unwrap_or_default(); + Some(Self { + token, + proposal_compute_us: i64::from(proposal_compute_us.max(0)), + }) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(in crate::frontend) struct PendingNativeMtpDraft { + pub(in crate::frontend) token: i32, + pub(in crate::frontend) origin: NativeMtpDraftOrigin, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(in crate::frontend) enum NativeMtpDraftOrigin { + InitialSerial, + SerialAfterGap, + VerifyNext, +} + +impl NativeMtpDraftOrigin { + pub(in crate::frontend) fn label(self) -> &'static str { + match self { + Self::InitialSerial => "initial_serial", + Self::SerialAfterGap => "serial_after_gap", + Self::VerifyNext => "verify_next", + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_prediction_token_sideband() { + assert_eq!( + NativeMtpDraft::from_prediction_tokens(&[11, 12, 34]), + Some(NativeMtpDraft { + token: 12, + proposal_compute_us: 34, + }) + ); + assert_eq!( + NativeMtpDraft::from_prediction_tokens(&[11, 12, 34, 567]), + Some(NativeMtpDraft { + token: 12, + proposal_compute_us: 34, + }) + ); + assert_eq!(NativeMtpDraft::from_prediction_tokens(&[11]), None); + } + + #[test] + fn parses_verify_prediction_token_sideband_after_verified_tokens() { + assert_eq!( + NativeMtpDraft::from_verify_prediction_tokens(&[10, 11, 12, 34], 2), + Some(NativeMtpDraft { + token: 12, + proposal_compute_us: 34, + }) + ); + assert_eq!( + NativeMtpDraft::from_verify_prediction_tokens(&[10, 11, 12, -3], 2), + Some(NativeMtpDraft { + token: 12, + proposal_compute_us: 0, + }) + ); + assert_eq!( + NativeMtpDraft::from_verify_prediction_tokens(&[10, 11, 12, 34, 567], 2), + Some(NativeMtpDraft { + token: 12, + proposal_compute_us: 34, + }) + ); + assert_eq!( + NativeMtpDraft::from_verify_prediction_tokens(&[10, 11], 2), + None + ); + } + + #[test] + fn pending_draft_keeps_origin_label() { + let pending = PendingNativeMtpDraft { + token: 12, + origin: NativeMtpDraftOrigin::VerifyNext, + }; + + assert_eq!(pending.token, 12); + assert_eq!(pending.origin.label(), "verify_next"); + } +} diff --git a/crates/skippy-server/src/frontend/native_mtp/env.rs b/crates/skippy-server/src/frontend/native_mtp/env.rs new file mode 100644 index 0000000000..54fbb698d0 --- /dev/null +++ b/crates/skippy-server/src/frontend/native_mtp/env.rs @@ -0,0 +1,94 @@ +const BATCHED_VERIFY_ENV: &str = "SKIPPY_NATIVE_MTP_BATCHED_VERIFY"; +const REJECT_COOLDOWN_TOKENS_ENV: &str = "SKIPPY_NATIVE_MTP_REJECT_COOLDOWN_TOKENS"; +const DEFER_REJECT_TRIM_ENV: &str = "SKIPPY_NATIVE_MTP_DEFER_REJECT_TRIM"; +const SUPPRESS_COOLDOWN_DRAFTS_ENV: &str = "SKIPPY_NATIVE_MTP_SUPPRESS_COOLDOWN_DRAFTS"; +const SUPPRESS_COOLDOWN_DRAFT_LIMIT_ENV: &str = "SKIPPY_NATIVE_MTP_SUPPRESS_COOLDOWN_DRAFT_LIMIT"; + +pub(in crate::frontend) fn native_mtp_batched_verify_enabled() -> bool { + native_mtp_batched_verify_enabled_from(std::env::var(BATCHED_VERIFY_ENV).ok().as_deref()) +} + +pub(in crate::frontend) fn native_mtp_reject_cooldown_tokens() -> usize { + parse_usize_env(REJECT_COOLDOWN_TOKENS_ENV, 0) +} + +pub(in crate::frontend) fn native_mtp_defer_reject_trim_enabled() -> bool { + truthy_env(std::env::var(DEFER_REJECT_TRIM_ENV).ok().as_deref()) +} + +pub(in crate::frontend) fn native_mtp_suppress_cooldown_drafts_enabled() -> bool { + truthy_env(std::env::var(SUPPRESS_COOLDOWN_DRAFTS_ENV).ok().as_deref()) +} + +pub(in crate::frontend) fn native_mtp_suppress_cooldown_draft_limit() -> usize { + parse_usize_env(SUPPRESS_COOLDOWN_DRAFT_LIMIT_ENV, 0) +} + +fn native_mtp_batched_verify_enabled_from(value: Option<&str>) -> bool { + !falsey_env(value) +} + +fn truthy_env(value: Option<&str>) -> bool { + matches!( + normalized_env(value).as_deref(), + Some("1" | "true" | "on" | "enable" | "enabled" | "yes") + ) +} + +fn falsey_env(value: Option<&str>) -> bool { + matches!( + normalized_env(value).as_deref(), + Some("0" | "false" | "off" | "disable" | "disabled" | "no") + ) +} + +fn normalized_env(value: Option<&str>) -> Option { + value.map(str::trim).map(str::to_ascii_lowercase) +} + +fn parse_usize_env(name: &str, default: usize) -> usize { + std::env::var(name) + .ok() + .and_then(|value| value.trim().parse::().ok()) + .unwrap_or(default) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn batched_verify_flag_defaults_on_and_accepts_false_values() { + assert!(native_mtp_batched_verify_enabled_from(None)); + assert!(native_mtp_batched_verify_enabled_from(Some("1"))); + assert!(native_mtp_batched_verify_enabled_from(Some("true"))); + assert!(!native_mtp_batched_verify_enabled_from(Some("0"))); + assert!(!native_mtp_batched_verify_enabled_from(Some("false"))); + assert!(!native_mtp_batched_verify_enabled_from(Some(" disabled "))); + } + + #[test] + fn truthy_env_accepts_enabled_aliases_only() { + for value in ["1", "true", " enabled ", "yes", "on"] { + assert!(truthy_env(Some(value)), "{value}"); + } + for value in [ + None, + Some("0"), + Some("false"), + Some("off"), + Some("disabled"), + ] { + assert!(!truthy_env(value), "{value:?}"); + } + } + + #[test] + fn numeric_options_default_when_absent() { + assert_eq!(parse_usize_env("SKIPPY_TEST_MISSING_REJECT_COOLDOWN", 0), 0); + assert_eq!( + parse_usize_env("SKIPPY_TEST_MISSING_SUPPRESS_COOLDOWN_LIMIT", 0), + 0 + ); + } +} diff --git a/crates/skippy-server/src/frontend/native_mtp/mod.rs b/crates/skippy-server/src/frontend/native_mtp/mod.rs new file mode 100644 index 0000000000..1262995fe2 --- /dev/null +++ b/crates/skippy-server/src/frontend/native_mtp/mod.rs @@ -0,0 +1,13 @@ +mod draft; +mod env; +mod stats; +mod verifier; + +pub(super) use draft::{NativeMtpDraft, NativeMtpDraftOrigin, PendingNativeMtpDraft}; +pub(super) use env::{ + native_mtp_batched_verify_enabled, native_mtp_defer_reject_trim_enabled, + native_mtp_reject_cooldown_tokens, native_mtp_suppress_cooldown_draft_limit, + native_mtp_suppress_cooldown_drafts_enabled, +}; +pub(super) use stats::{NativeMtpN1Stats, NativeMtpVerification}; +pub(super) use verifier::NativeMtpN1Verifier; diff --git a/crates/skippy-server/src/frontend/native_mtp/stats.rs b/crates/skippy-server/src/frontend/native_mtp/stats.rs new file mode 100644 index 0000000000..39fd1b6a14 --- /dev/null +++ b/crates/skippy-server/src/frontend/native_mtp/stats.rs @@ -0,0 +1,150 @@ +use std::collections::BTreeMap; + +use serde_json::{Value, json}; + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(in crate::frontend) enum NativeMtpVerification { + #[default] + NoPending, + Accepted { + draft: i32, + target: i32, + }, + Rejected { + draft: i32, + target: i32, + }, +} + +impl NativeMtpVerification { + pub(in crate::frontend) fn label(self) -> &'static str { + match self { + Self::NoPending => "none", + Self::Accepted { .. } => "accepted", + Self::Rejected { .. } => "rejected", + } + } +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(in crate::frontend) struct NativeMtpN1Stats { + pub(in crate::frontend) drafted_tokens: u64, + pub(in crate::frontend) accepted_tokens: u64, + pub(in crate::frontend) rejected_tokens: u64, + pub(in crate::frontend) pending_tokens: u64, + pub(in crate::frontend) verification_count: u64, + pub(in crate::frontend) proposal_compute_us: i64, + pub(in crate::frontend) verification_compute_us: i64, +} + +impl NativeMtpN1Stats { + pub(in crate::frontend) fn verified_tokens(self) -> u64 { + self.accepted_tokens + self.rejected_tokens + } + + pub(in crate::frontend) fn accept_rate(self) -> f64 { + let verified = self.verified_tokens(); + if verified == 0 { + 0.0 + } else { + self.accepted_tokens as f64 / verified as f64 + } + } + + pub(in crate::frontend) fn insert_attrs(self, attrs: &mut BTreeMap) { + if self.drafted_tokens == 0 && self.verified_tokens() == 0 { + attrs.insert("llama_stage.native_mtp.enabled".to_string(), json!(false)); + return; + } + + attrs.insert("llama_stage.native_mtp.enabled".to_string(), json!(true)); + attrs.insert( + "llama_stage.native_mtp.drafted".to_string(), + json!(self.drafted_tokens), + ); + attrs.insert( + "llama_stage.native_mtp.accepted".to_string(), + json!(self.accepted_tokens), + ); + attrs.insert( + "llama_stage.native_mtp.rejected".to_string(), + json!(self.rejected_tokens), + ); + attrs.insert( + "llama_stage.native_mtp.pending".to_string(), + json!(self.pending_tokens), + ); + attrs.insert( + "llama_stage.native_mtp.accept_rate".to_string(), + json!(self.accept_rate()), + ); + attrs.insert( + "llama_stage.native_mtp.proposal_compute_us".to_string(), + json!(self.proposal_compute_us), + ); + attrs.insert( + "llama_stage.native_mtp.verification_compute_us".to_string(), + json!(self.verification_compute_us), + ); + attrs.insert( + "llama_stage.native_mtp.verifications".to_string(), + json!(self.verification_count), + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn attrs_include_disabled_and_enabled_shapes() { + let mut attrs = BTreeMap::new(); + NativeMtpN1Stats::default().insert_attrs(&mut attrs); + assert_eq!( + attrs.get("llama_stage.native_mtp.enabled"), + Some(&json!(false)) + ); + + let stats = NativeMtpN1Stats { + drafted_tokens: 1, + accepted_tokens: 1, + verification_count: 1, + proposal_compute_us: 7, + verification_compute_us: 9, + ..NativeMtpN1Stats::default() + }; + + let mut attrs = BTreeMap::new(); + stats.insert_attrs(&mut attrs); + assert_eq!( + attrs.get("llama_stage.native_mtp.enabled"), + Some(&json!(true)) + ); + assert_eq!( + attrs.get("llama_stage.native_mtp.accept_rate"), + Some(&json!(1.0)) + ); + } + + #[test] + fn verification_labels_match_telemetry_values() { + assert_eq!(NativeMtpVerification::NoPending.label(), "none"); + assert_eq!( + NativeMtpVerification::Accepted { + draft: 1, + target: 1 + } + .label(), + "accepted" + ); + assert_eq!( + NativeMtpVerification::Rejected { + draft: 1, + target: 2 + } + .label(), + "rejected" + ); + } +} diff --git a/crates/skippy-server/src/frontend/native_mtp/verifier.rs b/crates/skippy-server/src/frontend/native_mtp/verifier.rs new file mode 100644 index 0000000000..be074c97fa --- /dev/null +++ b/crates/skippy-server/src/frontend/native_mtp/verifier.rs @@ -0,0 +1,340 @@ +use super::{ + NativeMtpDraft, NativeMtpDraftOrigin, NativeMtpN1Stats, NativeMtpVerification, + PendingNativeMtpDraft, +}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct PendingDraft { + token: i32, + origin: NativeMtpDraftOrigin, +} + +#[derive(Default)] +pub(in crate::frontend) struct NativeMtpN1Verifier { + pending: Option, + stats: NativeMtpN1Stats, +} + +impl NativeMtpN1Verifier { + pub(in crate::frontend) fn take_pending_draft(&mut self) -> Option { + self.pending.take().map(|pending| PendingNativeMtpDraft { + token: pending.token, + origin: pending.origin, + }) + } + + pub(in crate::frontend) fn clear_pending_draft(&mut self) { + self.pending = None; + } + + pub(in crate::frontend) fn observe_taken_draft_verification( + &mut self, + draft_token: i32, + target_token: i32, + verification_compute_us: i64, + ) -> NativeMtpVerification { + self.record_verification(draft_token, target_token, verification_compute_us) + } + + pub(in crate::frontend) fn observe_target_token( + &mut self, + target_token: i32, + verification_compute_us: i64, + next_draft: Option, + next_draft_origin: NativeMtpDraftOrigin, + ) -> NativeMtpVerification { + let verification = self.verify_pending(target_token, verification_compute_us); + self.observe_next_draft(next_draft, next_draft_origin); + verification + } + + pub(in crate::frontend) fn stats(&self) -> NativeMtpN1Stats { + let mut stats = self.stats; + stats.pending_tokens = u64::from(self.pending.is_some()); + stats + } + + fn verify_pending( + &mut self, + target_token: i32, + verification_compute_us: i64, + ) -> NativeMtpVerification { + let Some(pending) = self.pending.take() else { + return NativeMtpVerification::NoPending; + }; + + self.record_verification(pending.token, target_token, verification_compute_us) + } + + fn record_verification( + &mut self, + draft_token: i32, + target_token: i32, + verification_compute_us: i64, + ) -> NativeMtpVerification { + self.stats.verification_count += 1; + self.stats.verification_compute_us = self + .stats + .verification_compute_us + .saturating_add(verification_compute_us); + if draft_token == target_token { + self.stats.accepted_tokens += 1; + NativeMtpVerification::Accepted { + draft: draft_token, + target: target_token, + } + } else { + self.stats.rejected_tokens += 1; + NativeMtpVerification::Rejected { + draft: draft_token, + target: target_token, + } + } + } + + pub(in crate::frontend) fn observe_next_draft( + &mut self, + next_draft: Option, + origin: NativeMtpDraftOrigin, + ) { + let Some(next_draft) = next_draft else { + return; + }; + self.stats.drafted_tokens += 1; + self.stats.proposal_compute_us = self + .stats + .proposal_compute_us + .saturating_add(next_draft.proposal_compute_us); + self.pending = Some(PendingDraft { + token: next_draft.token, + origin, + }); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn draft(token: i32) -> NativeMtpDraft { + NativeMtpDraft { + token, + proposal_compute_us: 7, + } + } + + fn observe( + verifier: &mut NativeMtpN1Verifier, + target_token: i32, + verification_compute_us: i64, + next_draft: Option, + ) -> NativeMtpVerification { + verifier.observe_target_token( + target_token, + verification_compute_us, + next_draft, + NativeMtpDraftOrigin::InitialSerial, + ) + } + + #[test] + fn no_draft_behaves_like_baseline() { + let mut verifier = NativeMtpN1Verifier::default(); + + let decision = observe(&mut verifier, 11, 5, None); + + assert_eq!(decision, NativeMtpVerification::NoPending); + assert_eq!(verifier.stats(), NativeMtpN1Stats::default()); + } + + #[test] + fn first_draft_is_pending_until_next_target_decode() { + let mut verifier = NativeMtpN1Verifier::default(); + + let decision = observe(&mut verifier, 11, 5, Some(draft(12))); + + assert_eq!(decision, NativeMtpVerification::NoPending); + assert_eq!( + verifier.stats(), + NativeMtpN1Stats { + drafted_tokens: 1, + pending_tokens: 1, + proposal_compute_us: 7, + ..NativeMtpN1Stats::default() + } + ); + } + + #[test] + fn matching_next_target_accepts_pending_draft() { + let mut verifier = NativeMtpN1Verifier::default(); + observe(&mut verifier, 11, 5, Some(draft(12))); + + let decision = observe(&mut verifier, 12, 9, None); + + assert_eq!( + decision, + NativeMtpVerification::Accepted { + draft: 12, + target: 12, + } + ); + assert_eq!( + verifier.stats(), + NativeMtpN1Stats { + drafted_tokens: 1, + accepted_tokens: 1, + verification_count: 1, + proposal_compute_us: 7, + verification_compute_us: 9, + ..NativeMtpN1Stats::default() + } + ); + } + + #[test] + fn different_next_target_rejects_pending_draft() { + let mut verifier = NativeMtpN1Verifier::default(); + observe(&mut verifier, 11, 5, Some(draft(12))); + + let decision = observe(&mut verifier, 13, 9, None); + + assert_eq!( + decision, + NativeMtpVerification::Rejected { + draft: 12, + target: 13, + } + ); + assert_eq!( + verifier.stats(), + NativeMtpN1Stats { + drafted_tokens: 1, + rejected_tokens: 1, + verification_count: 1, + proposal_compute_us: 7, + verification_compute_us: 9, + ..NativeMtpN1Stats::default() + } + ); + } + + #[test] + fn verifies_previous_draft_before_storing_next_draft() { + let mut verifier = NativeMtpN1Verifier::default(); + observe(&mut verifier, 11, 5, Some(draft(12))); + + let decision = observe(&mut verifier, 12, 9, Some(draft(14))); + + assert_eq!( + decision, + NativeMtpVerification::Accepted { + draft: 12, + target: 12, + } + ); + assert_eq!( + verifier.stats(), + NativeMtpN1Stats { + drafted_tokens: 2, + accepted_tokens: 1, + pending_tokens: 1, + verification_count: 1, + proposal_compute_us: 14, + verification_compute_us: 9, + ..NativeMtpN1Stats::default() + } + ); + } + + #[test] + fn taken_pending_draft_can_be_recorded_as_batched_accept() { + let mut verifier = NativeMtpN1Verifier::default(); + observe(&mut verifier, 11, 5, Some(draft(12))); + + let pending = verifier.take_pending_draft().unwrap(); + assert_eq!(pending.origin, NativeMtpDraftOrigin::InitialSerial); + assert!(verifier.take_pending_draft().is_none()); + let decision = verifier.observe_taken_draft_verification(pending.token, 12, 9); + + assert_eq!( + decision, + NativeMtpVerification::Accepted { + draft: 12, + target: 12, + } + ); + assert_eq!( + verifier.stats(), + NativeMtpN1Stats { + drafted_tokens: 1, + accepted_tokens: 1, + verification_count: 1, + proposal_compute_us: 7, + verification_compute_us: 9, + ..NativeMtpN1Stats::default() + } + ); + } + + #[test] + fn taken_pending_draft_can_be_recorded_as_batched_reject() { + let mut verifier = NativeMtpN1Verifier::default(); + observe(&mut verifier, 11, 5, Some(draft(12))); + + let pending = verifier.take_pending_draft().unwrap(); + assert_eq!(pending.origin, NativeMtpDraftOrigin::InitialSerial); + let decision = verifier.observe_taken_draft_verification(pending.token, 13, 9); + + assert_eq!( + decision, + NativeMtpVerification::Rejected { + draft: 12, + target: 13, + } + ); + assert_eq!( + verifier.stats(), + NativeMtpN1Stats { + drafted_tokens: 1, + rejected_tokens: 1, + verification_count: 1, + proposal_compute_us: 7, + verification_compute_us: 9, + ..NativeMtpN1Stats::default() + } + ); + } + + #[test] + fn clear_pending_draft_drops_unverified_draft_without_changing_stats() { + let mut verifier = NativeMtpN1Verifier::default(); + observe(&mut verifier, 11, 5, Some(draft(12))); + + verifier.clear_pending_draft(); + + assert_eq!( + verifier.stats(), + NativeMtpN1Stats { + drafted_tokens: 1, + proposal_compute_us: 7, + ..NativeMtpN1Stats::default() + } + ); + assert_eq!( + observe(&mut verifier, 12, 9, None), + NativeMtpVerification::NoPending + ); + } + + #[test] + fn verification_compute_time_saturates_instead_of_overflowing() { + let mut verifier = NativeMtpN1Verifier::default(); + + observe(&mut verifier, 11, i64::MAX, Some(draft(12))); + observe(&mut verifier, 13, i64::MAX, Some(draft(14))); + observe(&mut verifier, 15, 1, None); + + assert_eq!(verifier.stats().verification_compute_us, i64::MAX); + } +} diff --git a/crates/skippy-server/src/frontend/prefix_cache.rs b/crates/skippy-server/src/frontend/prefix_cache.rs index 20c932eb58..9f37315aad 100644 --- a/crates/skippy-server/src/frontend/prefix_cache.rs +++ b/crates/skippy-server/src/frontend/prefix_cache.rs @@ -661,6 +661,7 @@ impl StageOpenAiBackend { return Ok(Some(EmbeddedFusedFirstDecode { predicted: *replay.last().expect("checked replay length"), predicted_tokens: replay, + native_mtp_draft: None, reply_stats: restore.stats, execution: EmbeddedExecutionStats::default(), elapsed_ms: timer.elapsed_ms(), @@ -734,6 +735,7 @@ impl StageOpenAiBackend { Ok(Some(EmbeddedFusedFirstDecode { predicted, predicted_tokens: vec![predicted], + native_mtp_draft: None, reply_stats: restore.stats, execution: EmbeddedExecutionStats::default(), elapsed_ms: timer.elapsed_ms(), @@ -978,12 +980,7 @@ impl StageOpenAiBackend { .map_err(openai_io_error)?; let forward_write_ms = write_timer.elapsed_ms(); let wait_timer = PhaseTimer::start(); - let downstream_reply = request - .prediction_return - .as_ref() - .ok_or_else(|| OpenAiError::backend("missing direct prediction return receiver"))? - .recv() - .map_err(openai_backend_error)?; + let downstream_reply = recv_reply(&mut *downstream).map_err(openai_io_error)?; let downstream_wait_ms = wait_timer.elapsed_ms(); let downstream_missed = downstream_reply.kind != WireReplyKind::PredictedToken || downstream_reply.stats.kv_lookup_errors > 0 @@ -1031,6 +1028,9 @@ impl StageOpenAiBackend { Ok(Some(EmbeddedFusedFirstDecode { predicted: downstream_reply.predicted, predicted_tokens: vec![downstream_reply.predicted], + native_mtp_draft: NativeMtpDraft::from_prediction_tokens( + &downstream_reply.predicted_tokens, + ), reply_stats, execution: EmbeddedExecutionStats { stage0_compute_ms, diff --git a/crates/skippy-server/src/frontend/request.rs b/crates/skippy-server/src/frontend/request.rs index e7c56edaae..de5038ac68 100644 --- a/crates/skippy-server/src/frontend/request.rs +++ b/crates/skippy-server/src/frontend/request.rs @@ -126,8 +126,11 @@ pub(super) fn media_url(part: &MessageContentPart) -> Option { pub(super) fn media_data(part: &MessageContentPart) -> Option { for key in ["input_audio", "audio", "image", "input_image", "image_url"] { - if let Some(value) = part.extra.get(key) - && let Some(data) = value.get("data").and_then(Value::as_str) + if let Some(data) = part + .extra + .get(key) + .and_then(|value| value.get("data")) + .and_then(Value::as_str) { return Some(data.to_string()); } @@ -136,11 +139,11 @@ pub(super) fn media_data(part: &MessageContentPart) -> Option { } pub(super) fn decode_media_url(url: &str) -> OpenAiResult> { - if let Some((prefix, payload)) = url.split_once(',') - && prefix.starts_with("data:") - && prefix.contains(";base64") - { - return decode_base64_payload(payload); + match url.split_once(',') { + Some((prefix, payload)) if prefix.starts_with("data:") && prefix.contains(";base64") => { + return decode_base64_payload(payload); + } + _ => {} } if url.starts_with("http://") || url.starts_with("https://") { return Err(OpenAiError::unsupported( @@ -197,25 +200,23 @@ fn apply_shared_request_defaults( .as_ref() .map(|values| stop_sequence_from_defaults(values.clone())); } - if extra_value_is_omitted(extra, "top_k") - && let Some(value) = defaults.top_k - { + if let (true, Some(value)) = (extra_value_is_omitted(extra, "top_k"), defaults.top_k) { extra.insert("top_k".to_string(), serde_json::json!(value)); } - if extra_value_is_omitted(extra, "min_p") - && let Some(value) = defaults.min_p - { + if let (true, Some(value)) = (extra_value_is_omitted(extra, "min_p"), defaults.min_p) { extra.insert("min_p".to_string(), serde_json::json!(value)); } - if extra_value_is_omitted(extra, "repeat_penalty") - && extra_value_is_omitted(extra, "repetition_penalty") - && let Some(value) = defaults.repeat_penalty - { + if let (true, Some(value)) = ( + extra_value_is_omitted(extra, "repeat_penalty") + && extra_value_is_omitted(extra, "repetition_penalty"), + defaults.repeat_penalty, + ) { extra.insert("repeat_penalty".to_string(), serde_json::json!(value)); } - if extra_value_is_omitted(extra, "repeat_last_n") - && let Some(value) = defaults.repeat_last_n - { + if let (true, Some(value)) = ( + extra_value_is_omitted(extra, "repeat_last_n"), + defaults.repeat_last_n, + ) { extra.insert("repeat_last_n".to_string(), serde_json::json!(value)); } apply_reasoning_defaults(reasoning, reasoning_effort, extra, defaults); diff --git a/crates/skippy-server/src/frontend/tests.rs b/crates/skippy-server/src/frontend/tests.rs index d9a341aae4..78da4d509c 100644 --- a/crates/skippy-server/src/frontend/tests.rs +++ b/crates/skippy-server/src/frontend/tests.rs @@ -965,6 +965,43 @@ fn tool_request() -> ChatCompletionRequest { .unwrap() } +#[test] +fn plain_chat_does_not_require_chat_output_parser() { + let request: ChatCompletionRequest = serde_json::from_value(json!({ + "model": "test", + "messages": [{"role": "user", "content": "hi"}] + })) + .unwrap(); + + assert!(!chat_output_parser_required( + &request, + &ChatTemplateOptions::default(), + )); +} + +#[test] +fn tools_and_enabled_thinking_require_chat_output_parser() { + assert!(chat_output_parser_required( + &tool_request(), + &ChatTemplateOptions::default(), + )); + + let request: ChatCompletionRequest = serde_json::from_value(json!({ + "model": "test", + "messages": [{"role": "user", "content": "think"}], + "reasoning": {"enabled": true} + })) + .unwrap(); + + assert!(chat_output_parser_required( + &request, + &ChatTemplateOptions { + enable_thinking: Some(true), + ..ChatTemplateOptions::default() + }, + )); +} + #[test] fn parses_llama_message_tool_calls() { let request = tool_request(); @@ -1239,6 +1276,8 @@ fn multimodal_stage_config( fn local_openai_backend(config: StageConfig) -> Result { let runtime = load_runtime(&config)?.context("load smoke runtime")?; let ctx_size = usize::try_from(config.ctx_size).unwrap_or(usize::MAX); + let decode_batcher = DecodeBatcher::new(runtime.clone(), 1); + let decode_frame_batcher = DecodeFrameBatcher::new(runtime.clone(), 1); Ok(StageOpenAiBackend { runtime, telemetry: Telemetry::new( @@ -1262,6 +1301,8 @@ fn local_openai_backend(config: StageConfig) -> Result { generation_token_budget: Arc::new(GenerationTokenBudget::new(ctx_size)), hook_policy: None, kv: None, + decode_batcher, + decode_frame_batcher, }) } @@ -1419,6 +1460,8 @@ async fn real_multimodal_split_smoke_when_fixture_is_set() -> Result<()> { .context("create split smoke lane pool")?; let runtime = load_runtime(&stage0_config)?.context("load stage-0 smoke runtime")?; let ctx_size = usize::try_from(stage0_config.ctx_size).unwrap_or(usize::MAX); + let decode_batcher = DecodeBatcher::new(runtime.clone(), 1); + let decode_frame_batcher = DecodeFrameBatcher::new(runtime.clone(), 1); let backend = StageOpenAiBackend { runtime, telemetry, @@ -1435,7 +1478,6 @@ async fn real_multimodal_split_smoke_when_fixture_is_set() -> Result<()> { downstream_wire_condition: WireCondition::new(0.0, None)?, prefill_reply_credit_limit: 0, lane_pool: Some(lane_pool), - prediction_returns: None, }, draft: None, speculative_window: 0, @@ -1446,6 +1488,8 @@ async fn real_multimodal_split_smoke_when_fixture_is_set() -> Result<()> { generation_token_budget: Arc::new(GenerationTokenBudget::new(ctx_size)), hook_policy: None, kv: None, + decode_batcher, + decode_frame_batcher, }; let response = backend .chat_completion(multimodal_chat_request(&fixture)?) @@ -1463,6 +1507,19 @@ fn trims_at_first_stop_sequence() { assert_eq!(trim_at_stop("abc", &[""]), "abc"); } +#[test] +fn generation_stop_values_include_chat_template_stops() { + let request_stop = openai_frontend::StopSequence::One("".to_string()); + let metadata = json!({ + "additional_stops": ["<|user|>", "<|observation|>", ""], + }) + .to_string(); + + let stops = generation_stop_values(Some(&request_stop), Some(&metadata)); + + assert_eq!(stops, vec!["", "<|user|>", "<|observation|>"]); +} + #[test] fn valid_utf8_prefix_skips_incomplete_suffix() { assert_eq!(valid_utf8_prefix_len("hello".as_bytes()), 5); diff --git a/crates/skippy-server/src/frontend/util.rs b/crates/skippy-server/src/frontend/util.rs index 12a14c42c6..d208425333 100644 --- a/crates/skippy-server/src/frontend/util.rs +++ b/crates/skippy-server/src/frontend/util.rs @@ -12,6 +12,33 @@ pub(super) fn trim_at_stop<'a>(text: &'a str, stop_values: &[&str]) -> &'a str { } } +pub(super) fn generation_stop_values( + stop: Option<&openai_frontend::StopSequence>, + chat_metadata: Option<&str>, +) -> Vec { + let mut values: Vec = stop + .map(|stop| stop.values().into_iter().map(str::to_string).collect()) + .unwrap_or_default(); + let additional_stops = chat_metadata + .and_then(|metadata| serde_json::from_str::(metadata).ok()) + .and_then(|value| { + value + .get("additional_stops") + .and_then(serde_json::Value::as_array) + .cloned() + }); + if let Some(stops) = additional_stops { + values.extend( + stops + .iter() + .filter_map(serde_json::Value::as_str) + .filter(|value| !value.is_empty()) + .map(str::to_string), + ); + } + values +} + pub(super) fn valid_utf8_prefix_len(bytes: &[u8]) -> usize { match std::str::from_utf8(bytes) { Ok(_) => bytes.len(), diff --git a/crates/skippy-server/src/frontend/wire_messages.rs b/crates/skippy-server/src/frontend/wire_messages.rs index 5150b77079..e3bdb63f83 100644 --- a/crates/skippy-server/src/frontend/wire_messages.rs +++ b/crates/skippy-server/src/frontend/wire_messages.rs @@ -128,6 +128,7 @@ pub(super) struct VerifySpanMessageArgs<'a> { pub(super) pos_start: usize, pub(super) decode_step: usize, pub(super) tokens: &'a [i32], + pub(super) sampling: Option, pub(super) checkpoint: bool, } @@ -160,7 +161,7 @@ pub(super) fn embedded_verify_message( state, request_id: args.request_id, session_id: args.session_id, - sampling: None, + sampling: args.sampling, chat_sampling_metadata: None, tokens: args.tokens.to_vec(), positions: Vec::new(), @@ -191,6 +192,23 @@ pub(super) fn embedded_session_control_message( } } +pub(super) fn embedded_trim_session_message( + wire_dtype: WireActivationDType, + request_id: u64, + session_id: u64, + token_count: usize, +) -> OpenAiResult { + let mut message = embedded_session_control_message( + wire_dtype, + WireMessageKind::TrimSession, + request_id, + session_id, + ); + message.token_count = i32::try_from(token_count) + .map_err(|_| OpenAiError::backend("trim token count exceeds i32"))?; + Ok(message) +} + pub(super) fn generation_config_message( wire_dtype: WireActivationDType, request_id: u64, diff --git a/crates/skippy-server/src/runtime_state.rs b/crates/skippy-server/src/runtime_state.rs index f858b1071d..5d20c01caf 100644 --- a/crates/skippy-server/src/runtime_state.rs +++ b/crates/skippy-server/src/runtime_state.rs @@ -1,5 +1,5 @@ use std::{ - collections::BTreeMap, + collections::{BTreeMap, BTreeSet}, sync::{Arc, Mutex}, time::Instant, }; @@ -7,10 +7,11 @@ use std::{ use anyhow::{Context, Result, bail}; use skippy_protocol::{FlashAttentionType, LoadMode, StageConfig}; use skippy_runtime::{ - ActivationFrame, FlashAttentionType as RuntimeFlashAttentionType, GenerationSignalWindow, - MediaInput, MediaPrefill, MediaPrefillFrame, RuntimeConfig, RuntimeKvPage, RuntimeKvPageDesc, - RuntimeLoadMode, SamplingConfig, StageModel, StageSession, StageSessionCheckpoint, TokenSignal, - parse_cache_type, + ActivationFrame, DecodeBatchRequest, DecodeFrameBatchOutput, DecodeFrameBatchRequest, + FlashAttentionType as RuntimeFlashAttentionType, GenerationSignalWindow, MediaInput, + MediaPrefill, MediaPrefillFrame, NativeMtpDraft, RuntimeConfig, RuntimeKvPage, + RuntimeKvPageDesc, RuntimeLoadMode, SamplingConfig, StageModel, StageSession, + StageSessionCheckpoint, TokenSignal, parse_cache_type, }; use crate::package::select_package_parts; @@ -90,6 +91,25 @@ pub struct RuntimeSessionDropStats { pub stats_after: RuntimeSessionStats, } +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct RuntimeSessionAlignStats { + pub before_token_count: u64, + pub after_token_count: u64, +} + +pub struct RuntimeDecodeBatchRequest<'a> { + pub session_id: &'a str, + pub token_id: i32, + pub sampling: Option<&'a SamplingConfig>, +} + +pub struct RuntimeDecodeFrameBatchRequest<'a> { + pub session_id: &'a str, + pub token_id: i32, + pub sampling: Option<&'a SamplingConfig>, + pub input: Option<&'a ActivationFrame>, +} + #[derive(Debug, Clone)] struct ResidentLanePrefix { page_id: String, @@ -163,6 +183,53 @@ impl RuntimeState { Ok(token) } + pub fn decode_batch_sampled( + &mut self, + requests: &[RuntimeDecodeBatchRequest<'_>], + ) -> Result> { + if requests.is_empty() { + return Ok(Vec::new()); + } + Self::ensure_unique_batch_sessions(requests)?; + for request in requests { + self.session(request.session_id)?; + } + + let mut lane_sessions = Vec::with_capacity(requests.len()); + for request in requests { + let lane_session = self.sessions.remove(request.session_id).ok_or_else(|| { + anyhow::anyhow!( + "session {} was not active after admission", + request.session_id + ) + })?; + lane_sessions.push((request.session_id.to_string(), lane_session)); + } + + let result = { + let mut decode_requests = lane_sessions + .iter_mut() + .zip(requests.iter()) + .map(|((_, lane_session), request)| DecodeBatchRequest { + session: &mut lane_session.session, + token_id: request.token_id, + sampling: request.sampling, + }) + .collect::>(); + StageSession::decode_batch_sampled(&mut decode_requests) + }; + + for (session_id, lane_session) in lane_sessions { + self.sessions.insert(session_id, lane_session); + } + if result.is_ok() { + for request in requests { + self.add_session_tokens(request.session_id, 1); + } + } + result + } + pub fn session_batch_size(&mut self, session_id: &str) -> Result { self.active_session(session_id)?.batch_size() } @@ -259,19 +326,132 @@ impl RuntimeState { Ok(output) } + pub fn decode_frame_sampled_mtp_n1( + &mut self, + session_id: &str, + token_id: i32, + sampling: Option<&SamplingConfig>, + input: Option<&ActivationFrame>, + output_capacity: usize, + ) -> Result<(i32, Option, ActivationFrame)> { + let session = self.session(session_id)?; + let output = + session.decode_step_frame_sampled_mtp_n1(token_id, sampling, input, output_capacity)?; + self.add_session_tokens(session_id, 1); + Ok(output) + } + + pub fn decode_frame_batch_sampled( + &mut self, + requests: &[RuntimeDecodeFrameBatchRequest<'_>], + ) -> Result> { + if requests.is_empty() { + return Ok(Vec::new()); + } + Self::ensure_unique_frame_batch_sessions(requests)?; + for request in requests { + self.session(request.session_id)?; + } + + let mut lane_sessions = Vec::with_capacity(requests.len()); + for request in requests { + let lane_session = self.sessions.remove(request.session_id).ok_or_else(|| { + anyhow::anyhow!( + "session {} was not active after admission", + request.session_id + ) + })?; + lane_sessions.push((request.session_id.to_string(), lane_session)); + } + + let result = { + let mut decode_requests = lane_sessions + .iter_mut() + .zip(requests.iter()) + .map(|((_, lane_session), request)| DecodeFrameBatchRequest { + session: &mut lane_session.session, + token_id: request.token_id, + sampling: request.sampling, + input: request.input, + }) + .collect::>(); + StageSession::decode_step_frame_batch_sampled(&mut decode_requests) + }; + + for (session_id, lane_session) in lane_sessions { + self.sessions.insert(session_id, lane_session); + } + if result.is_ok() { + for request in requests { + self.add_session_tokens(request.session_id, 1); + } + } + result + } + pub fn verify_frame( &mut self, session_id: &str, token_ids: &[i32], input: Option<&ActivationFrame>, output_capacity: usize, + ) -> Result<(Vec, ActivationFrame)> { + self.verify_frame_sampled(session_id, token_ids, None, input, output_capacity) + } + + pub fn verify_frame_sampled( + &mut self, + session_id: &str, + token_ids: &[i32], + sampling: Option<&SamplingConfig>, + input: Option<&ActivationFrame>, + output_capacity: usize, ) -> Result<(Vec, ActivationFrame)> { let session = self.session(session_id)?; - let output = session.verify_tokens_frame(token_ids, input, output_capacity)?; + let output = + session.verify_tokens_frame_sampled(token_ids, sampling, input, output_capacity)?; self.add_session_tokens(session_id, token_ids.len() as u64); Ok(output) } + pub fn verify_frame_sampled_serial( + &mut self, + session_id: &str, + token_ids: &[i32], + sampling: Option<&SamplingConfig>, + input: Option<&ActivationFrame>, + output_capacity: usize, + ) -> Result<(Vec, ActivationFrame)> { + if token_ids.is_empty() { + bail!("serial verify_frame requires at least one token"); + } + let input_frames = split_activation_frame(input, token_ids.len())?; + let mut predicted_tokens = Vec::with_capacity(token_ids.len() + 2); + let mut output_frames = Vec::with_capacity(token_ids.len()); + let mut last_draft = None; + for (index, token_id) in token_ids.iter().copied().enumerate() { + let input_frame = input_frames.as_ref().map(|frames| &frames[index]); + let (predicted, native_mtp, output) = self.decode_frame_sampled_mtp_n1( + session_id, + token_id, + sampling, + input_frame, + output_capacity, + )?; + if predicted >= 0 { + predicted_tokens.push(predicted); + } + last_draft = native_mtp; + output_frames.push(output); + } + if let Some(draft) = last_draft { + predicted_tokens.push(draft.token_id); + predicted_tokens + .push(i32::try_from(draft.proposal_compute_us.max(0)).unwrap_or(i32::MAX)); + } + Ok((predicted_tokens, combine_activation_frames(&output_frames)?)) + } + pub fn checkpoint_session(&mut self, session_id: &str) -> Result<()> { let checkpoint = self.session(session_id)?.checkpoint()?; self.session_checkpoints @@ -303,6 +483,24 @@ impl RuntimeState { Ok(()) } + pub fn align_session_to_token_count_if_ahead( + &mut self, + session_id: &str, + token_count: u64, + ) -> Result> { + let Some(current) = self.session_token_counts.get(session_id).copied() else { + return Ok(None); + }; + if current <= token_count { + return Ok(None); + } + self.trim_session(session_id, token_count)?; + Ok(Some(RuntimeSessionAlignStats { + before_token_count: current, + after_token_count: token_count, + })) + } + fn session(&mut self, session_id: &str) -> Result<&mut StageSession> { if !self.sessions.contains_key(session_id) { let lane_session = self.take_idle_session().map(Ok).unwrap_or_else(|| { @@ -320,6 +518,31 @@ impl RuntimeState { .session) } + fn ensure_unique_batch_sessions(requests: &[RuntimeDecodeBatchRequest<'_>]) -> Result<()> { + let mut seen = BTreeSet::new(); + for request in requests { + if !seen.insert(request.session_id) { + bail!("duplicate session {} in decode batch", request.session_id); + } + } + Ok(()) + } + + fn ensure_unique_frame_batch_sessions( + requests: &[RuntimeDecodeFrameBatchRequest<'_>], + ) -> Result<()> { + let mut seen = BTreeSet::new(); + for request in requests { + if !seen.insert(request.session_id) { + bail!( + "duplicate session {} in decode frame batch", + request.session_id + ); + } + } + Ok(()) + } + fn active_session(&mut self, session_id: &str) -> Result<&mut StageSession> { self.sessions .get_mut(session_id) @@ -796,6 +1019,76 @@ impl RuntimeState { } } +fn split_activation_frame( + input: Option<&ActivationFrame>, + token_count: usize, +) -> Result>> { + let Some(input) = input else { + return Ok(None); + }; + if token_count == 0 { + bail!("cannot split activation frame for zero tokens"); + } + if input.desc.token_count as usize != token_count { + bail!( + "activation token count mismatch: frame={} tokens={}", + input.desc.token_count, + token_count + ); + } + if input.payload.len() % token_count != 0 { + bail!( + "activation payload is not divisible by token count: payload={} tokens={}", + input.payload.len(), + token_count + ); + } + let row_bytes = input.payload.len() / token_count; + let frames = input + .payload + .chunks(row_bytes) + .map(|row| { + let mut desc = input.desc; + desc.token_count = 1; + desc.sequence_count = 1; + desc.payload_bytes = row.len() as u64; + ActivationFrame { + desc, + payload: row.to_vec(), + } + }) + .collect(); + Ok(Some(frames)) +} + +fn combine_activation_frames(frames: &[ActivationFrame]) -> Result { + let Some(first) = frames.first() else { + bail!("cannot combine empty activation frames"); + }; + let mut desc = first.desc; + let mut payload = Vec::new(); + let mut token_count = 0u32; + for frame in frames { + if frame.desc.dtype != desc.dtype + || frame.desc.layout != desc.layout + || frame.desc.producer_stage_index != desc.producer_stage_index + || frame.desc.layer_start != desc.layer_start + || frame.desc.layer_end != desc.layer_end + || frame.desc.sequence_count != desc.sequence_count + || frame.desc.flags != desc.flags + { + bail!("cannot combine incompatible activation frames"); + } + token_count = token_count + .checked_add(frame.desc.token_count) + .context("combined activation token count overflow")?; + payload.extend_from_slice(&frame.payload); + } + desc.token_count = token_count; + desc.payload_bytes = payload.len() as u64; + Ok(ActivationFrame { desc, payload }) +} + /// Allocate the next lane slot. /// /// Prefers indices in `free_lane_indices` (lanes previously discarded diff --git a/docs/design/GLM47_SPD_EXECUTION_PLAN.md b/docs/design/GLM47_SPD_EXECUTION_PLAN.md new file mode 100644 index 0000000000..c03649ef56 --- /dev/null +++ b/docs/design/GLM47_SPD_EXECUTION_PLAN.md @@ -0,0 +1,135 @@ +# GLM 4.7 SPD Execution Plan + +This plan combines the local GLM 4.7 checkpoint, the GLM llama.cpp/Skippy work +on `feat/jianyang-glm-llama-patches`, the Skippy SPD proof handoff in PR #859, +and the reference implementation at `yuyijiong/speculative_pipeline_decoding`. + +The goal is to train a GLM 4.7 SPD sidecar model that can act as a wider +drafting oracle than native GLM MTP, then use the Speedy benchmark to compare +vanilla GLM decode against GLM decode with the verified SPD sidecar enabled. + +## Scope + +This effort is about trained SPD on top of the current native-MTP branch. Native +GLM MTP remains the verifier/correctness foundation and implementation base, but +its `N=1` proposal width is likely too narrow to amortize multi-stage Skippy +latency. SPD is the candidate `N>1` drafting oracle. It should not be treated as +a competing unverified benchmark lane. + +PR #860 is the compact GLM SPD donor branch for training, export, manifest, and +latency-model plumbing. PR #859 remains useful proof archaeology, but its broad +live-serving/protocol changes should only be revisited after offline GLM sidecar +quality justifies serving integration. The reference SPD repo remains the +training/evaluation source for the sidecar head. + +## Phase 1: Inspect GLM + +Inspect the local checkpoint before training or benchmarking: + +- architecture and `model_type` +- tokenizer identity and chat template presence +- `num_hidden_layers` +- `hidden_size` +- vocab size +- any GLM-specific auxiliary tensors preserved by llama.cpp patches + +The local `zai-org/GLM-4.7-Flash` snapshot reports: + +- architecture: `Glm4MoeLiteForCausalLM` +- model type: `glm4_moe_lite` +- target layers: `47` +- hidden size: `2048` +- vocab size: `154880` +- auxiliary layer-47 tensors: `eh_proj`, `enorm`, `hnorm` + +## Phase 2: Frontload Code Risk + +Before long GPU training jobs, make the GLM SPD path executable enough to expose +integration gaps: + +- add GLM model metadata inspection +- support explicit non-uniform `stage_layer_boundaries` +- derive GLM hidden-state tap rows from those boundaries +- generate a small GLM-tokenizer draft vocabulary for training smoke runs +- patch the reference trainer so explicit tap rows can bypass equal-stage + assumptions +- write manifest-compatible GLM SPD smoke artifacts +- validate the smoke manifest through `skippy-runtime` +- keep the smoke artifacts clearly separate from trained weights + +`evals/spd/glm47_frontload.py` owns this first executable surface. The default +GLM 4.7 Flash topology uses stage boundaries `15,31,47`, matching the 47-layer +target model without relying on equal layer division. + +## Phase 3: Speedy Baseline + +After the GLM code path exists, establish the vanilla GLM baseline: + +1. Run the Speedy benchmark against vanilla GLM decode from the local checkpoint + or Skippy package. +2. Freeze Speedy prompt set, generation settings, tokenizer, temperature, max + tokens, hardware, and runtime build. +3. Record Speedy throughput, latency distribution, generated token counts, and + output text. +4. If using Skippy, validate vanilla split correctness against non-split GLM. + +This baseline is the only performance comparison target for SPD. + +## Phase 4: Training Smoke + +Before a model-quality run, run a tiny real training smoke: + +- frozen local GLM 4.7 base model +- explicit stage boundaries `15,31,47` +- derived hidden tap rows `0,15,31,47;0,15,31;0,15` +- small generated GLM-tokenizer draft vocab +- a few rows from the training corpus +- `--skip-eval` unless the training checkpoint is produced successfully + +The success criterion is artifact flow, not acceptance quality: + +- `speculation_head_final.pt` +- `skippy-spd-head.json` +- optional `spd-head.safetensors` after export +- Rust SPD manifest validation + +The initial tiny GLM 4.7 smoke passed and the artifacts are stored in the +private Hugging Face model repo `meshllm/skippy-spd-glm47-train-smoke`. The +repo contains the manifest, serving safetensors export, and original reference +checkpoint for reproducing the Skippy SPD manifest and serving-artifact path. + +## Phase 5: Train And Evaluate SPD + +Build a GLM tokenizer-specific draft vocabulary. Do not reuse Qwen draft vocab. +Start with 32k tokens, then try 50k if vocab coverage limits acceptance. + +Train only the SPD sidecar head against the frozen GLM base model. Evaluate with +verification enabled and record accepted draft flags, acceptance rate, +equivalent accept length, theoretical gain, summaries, and raw traces. + +Compare these paths, with tok/s as the hard optimization target: + +- vanilla llama.cpp GLM target decode with zero MTP +- current Skippy GLM native-MTP verifier path +- Skippy GLM target decode with verified SPD sidecar, reported by draft width + `N=1,2,4,8` + +## Phase 6: Serving Decision + +Only wire the trained SPD head into live Skippy if the Speedy benchmark shows +it is materially faster than vanilla GLM while preserving target-equivalent +verified output. + +Serving integration then needs: + +- Rust safetensors loading for the SPD head +- GLM SPD forward pass for the recorded topology +- Skippy hidden-state taps or transport +- Python/Rust proposal parity on fixed taps +- verified proposal generation +- rollback/session trim for rejected proposals +- SPD metrics for draft, accept, reject, proposal time, verification time, and + end-to-end throughput + +If the trained sidecar is weak, keep the result as research evidence and do not +wire it into serving. diff --git a/docs/design/GLM_NATIVE_MTP_SKIPPY_ARCHITECTURE.md b/docs/design/GLM_NATIVE_MTP_SKIPPY_ARCHITECTURE.md new file mode 100644 index 0000000000..889359929c --- /dev/null +++ b/docs/design/GLM_NATIVE_MTP_SKIPPY_ARCHITECTURE.md @@ -0,0 +1,275 @@ +# GLM Native MTP in Skippy + +## Goal + +Build GLM native MTP in two ordered steps: + +1. Native MTP `n=1` correctness. +2. Batched verification for latency amortization. + +The existing batched verification ABI in the llama.cpp patch queue is historical. +Do not treat it as the architecture for this work. The new path should start +from native MTP proposal semantics and add batching only after one-token +correctness is proven. + +## Current GLM-4.7 Evidence + +HF metadata checked on 2026-06-15 shows `zai-org/GLM-4.7-Flash` is the native +MTP source checkpoint and `unsloth/GLM-4.7-Flash-GGUF` publishes the existing +Skippy-certified public GGUF: + +- base model: `zai-org/GLM-4.7-Flash`; +- model class: 30B-A3B MoE; +- checkpoint architecture: `Glm4MoeLiteForCausalLM`; +- checkpoint config: `num_hidden_layers = 47`, + `num_nextn_predict_layers = 1`, `hidden_size = 2048`; +- checkpoint MTP tensors live at layer `47`: + - `model.layers.47.eh_proj.weight`; + - `model.layers.47.enorm.weight`; + - `model.layers.47.hnorm.weight`; +- existing public Skippy artifact: + `unsloth/GLM-4.7-Flash-GGUF:Q4_K_M`; +- Skippy certified split plan: `layer_end=47`, `splits=15,31`, activation + width `2048`; +- Skippy wire dtype: `f16` by default, with q8 already validated for the + existing non-MTP parity path; +- upstream model card shows native framework MTP with + `--speculative-config.method mtp` and + `--speculative-config.num_speculative_tokens 1`. + +That makes GLM-4.7 Flash the right first target for native `n=1` MTP in Skippy: +it is small enough to exercise locally, already certified for stage splitting, +and has an upstream native one-token MTP path to compare against. + +The public `GLM-4.7-Flash-Q4_K_M.gguf` is not enough for this milestone. Its +GGUF metadata reports `general.architecture = "deepseek2"` and +`deepseek2.block_count = 47`, and a tensor-name scan found no `next`, `mtp`, +`eh_proj`, `enorm`, or `hnorm` tensors. The checkpoint has the native MTP +tensors, but the public GGUF appears to have dropped them. Phase 1 therefore +needs a custom GLM-4.7 GGUF quantization from checkpoint that preserves the +layer-47 MTP tensors, then a Skippy layer package built from that GGUF. +The target Hub repo for that custom artifact is +`meshllm/GLM-4.7-Flash-MTP-GGUF`. + +The model-backed Step 1 gate should use the correctness harness with a required +draft sideband: + +```bash +skippy-correctness chain \ + --model /path/to/GLM-4.7-Flash-Q4_K_M-mtp.gguf \ + --model-id meshllm/GLM-4.7-Flash-MTP-GGUF:Q4_K_M \ + --layer-end 48 \ + --splits 15,31 \ + --activation-wire-dtype f16 \ + --require-native-mtp-draft +``` + +Use the layer count from the MTP-preserving GGUF/package when setting +`--layer-end`. If the converter stores 47 authoritative target layers plus one +NextN layer, the Skippy proof command should use `--layer-end 48`; if the layer +package exposes only authoritative target layers and attaches MTP as sidecar +metadata, use the package topology instead. + +The important invariant is that the report must show +`matches = true`, `native_mtp.sideband_present = true`, and +`native_mtp.authoritative_matches_reply = true`. + +## Runtime Shape + +Native MTP is a proposer sidecar attached to the target runtime. It is not a +Skippy trunk stage and it is not an external draft model. + +```mermaid +flowchart LR + C["Coordinator / OpenAI frontend"] --> S0["Skippy stage 0"] + S0 --> S1["Skippy middle stages"] + S1 --> SF["Final Skippy target stage"] + SF --> LM["Target LM head"] + SF --> MTP["Native MTP sidecar"] + LM --> R["Target token"] + MTP --> D["Draft token"] +``` + +The final target stage is the natural owner for native MTP because it already +has the last target hidden state and logits. The MTP sidecar should run against +that target state and return one proposed future token. + +## Step 1: Native MTP n=1 Correctness + +The first milestone should not try to amortize network latency. It should prove +that Skippy can obtain one native MTP draft token, compare it against the next +authoritative target decode, and report accept/reject metrics without changing +greedy output. + +```mermaid +sequenceDiagram + participant C as Coordinator + participant S0 as Stage 0 + participant SF as Final target stage + MTP + + C->>S0: decode current token x + S0->>SF: target activation + SF->>SF: target LM head selects token A + SF->>SF: native MTP proposes token B + SF->>C: target=A, draft=B + C->>C: emit A, remember draft B + + C->>S0: decode authoritative token A + S0->>SF: target activation + SF->>SF: target LM head selects token B' + SF->>C: target=B' + C->>C: accept if B == B', else reject +``` + +Important properties: + +- Target KV state remains authoritative. +- No target rollback is needed in this milestone. +- The MTP sidecar may need reset or advance logic after every observed target + token so its proposal state follows the committed target prefix. +- The visible token stream is always produced by target tokens. +- Metrics can still count `drafted`, `accepted`, `rejected`, accept rate, MTP + proposal latency, and target verification latency. + +This milestone can be slower than baseline. Its job is correctness and +instrumentation. + +## Step 1 ABI + +The smallest useful stage ABI is a target decode that can optionally return a +native MTP draft: + +```text +skippy_decode_step_frame_sampled_mtp_n1( + session, + token_id, + sampling, + input_activation, + output_activation, + out_target_token, + out_mtp_draft_token, + out_mtp_status, + out_mtp_compute_us, +) +``` + +The actual naming can differ, but the contract should stay this small: + +- one target input token; +- one target sampled/greedy output token; +- zero or one native MTP draft token; +- no speculative target commit semantics; +- no multi-token verification batch; +- no external draft model path. + +If the model has no compatible native MTP head, the ABI should return +`out_mtp_status = unavailable` and the Rust loop should behave exactly like +baseline Skippy. + +## Step 2: Batched Verification + +Only after Step 1 is correct should Skippy add batched verification. Batched +verification sends the target token and the one MTP draft token through the +target stage chain together as a provisional suffix. + +```mermaid +sequenceDiagram + participant C as Coordinator + participant S0 as Stage 0 + participant S1 as Middle stages + participant SF as Final target stage + MTP + + C->>S0: verification batch [A, draft B] + S0->>S1: activation batch [A, draft B] + S1->>SF: activation batch [A, draft B] + SF->>SF: target predicts [B', C'] + SF->>C: predictions [B', C'], optional next MTP draft + C->>C: if B == B': commit A,B; else commit A only + C->>S0: commit/trim decision +``` + +The target stages must treat the draft token as provisional. If the target does +not agree with the draft, every stage must discard KV/state for the rejected +draft position. + +```mermaid +flowchart TD + P["Committed prefix"] --> V["Run verification batch [A, draft B]"] + V --> Q{"Target B' == draft B?"} + Q -->|yes| AC["Commit A and B"] + Q -->|no| RJ["Commit A, trim draft B"] + AC --> N["Continue from accepted prefix"] + RJ --> N +``` + +This is where latency can be amortized. With an accept rate near 67%, one stage +chain traversal often advances more than one visible token. + +## Batched Verification Contract + +The new batched contract should be transactional and native-MTP-aware: + +```text +VerificationBatch { + committed_prefix_tokens: u64, + tokens: [target_token, mtp_draft_token], + draft_start: 1, + draft_count: 1, + mode: native_mtp_n1, +} + +VerificationBatchReply { + target_predictions: [token_after_target, token_after_draft], + mtp_draft_after_last_committed: optional token, + timings, +} + +CommitDecision { + commit_token_count: 1 or 2, +} +``` + +The commit decision is separate from the batch execution because only the +coordinator can compare the returned target prediction with the draft token and +decide how much of the provisional suffix is valid. + +## Sampling Scope + +The first correctness gate should be greedy or otherwise deterministic. Batched +verification with sampling is a later problem because sampler history must be +advanced consistently for each provisionally accepted position. + +For the GLM coding-agent lab goal, the first usable target is: + +- deterministic decode; +- native MTP `n=1`; +- accept/reject metrics; +- no target rollback in Step 1; +- transactional target rollback only in Step 2. + +## Non-Goals + +- Reusing the discarded batched verification ABI as the design center. +- Treating MTP heads as ordinary Skippy trunk layers. +- External draft-model speculation. +- `n=2` or `n=3` before `n=1` is proven. +- Multi-host throughput claims before single-node correctness and metrics pass. + +## Acceptance Gates + +Step 1 is complete when: + +1. A GLM model with native MTP produces nonzero MTP draft attempts through + Skippy. +2. Greedy output is byte-identical to Skippy baseline. +3. Accept/reject metrics are emitted. +4. Disabling native MTP returns to baseline behavior. + +Step 2 is complete when: + +1. Batched verification commits accepted draft positions and trims rejected + provisional positions correctly. +2. Greedy output remains byte-identical to Skippy baseline. +3. Stage telemetry reports batch size, accepted draft count, rejected draft + count, commit count, trim count, and latency. +4. `skippy-bench` artifacts show throughput versus Skippy baseline. diff --git a/evals/spd/README.md b/evals/spd/README.md new file mode 100644 index 0000000000..2ad2e78cfe --- /dev/null +++ b/evals/spd/README.md @@ -0,0 +1,526 @@ +# GLM 4.7 SPD-on-MTP Eval Notes + +This directory contains the reproducible training, export, and latency-model +tools for the GLM 4.7 SPD-on-MTP experiment. The active path is to train a new +GLM 4.7 sidecar head, export it as a Skippy-readable artifact, and use verified +target decoding to decide whether SPD can provide useful `N > 1` draft tokens. + +SPD is treated here as a separate trained sidecar head. It proposes draft tokens +from selected target-model hidden states; the target model still verifies every +accepted token. The work in this directory proves the training/evaluation path +and records the artifact contract Skippy needs before serving the head from +Rust. + +The Qwen results below are background evidence from the donor proof. They are +useful because they show that SPD can be a strong wide drafting oracle, but they +are not the model-selection target for this branch. + +## What Works + +- GLM 4.7 checkpoint inspection, non-uniform stage-boundary metadata, and smoke + artifact generation are supported through `glm47_frontload.py`. +- The reference training wrapper accepts explicit GLM tap rows derived from + `--stage-layer-boundaries 15,31,47` and can build a GLM-tokenizer draft vocab + for smoke runs. +- The training wrapper can now emit a generic contiguous-layer topology plan + that records randomized logical hidden-state tap layouts without training a + fixed-stage head. +- `generic_layer_tap_sidecar.py` can train, evaluate, and export a + topology-independent layer-tap sidecar that uses logical hidden-state taps, + tap features, randomized contiguous layouts, and tap dropout metadata instead + of fixed stage projection tensors. +- A real SPD head can be trained locally for `Qwen/Qwen3-0.6B` with the paper's + reference implementation. +- A real pretrained SPD head for `Qwen/Qwen3.5-4B` reaches high acceptance on + local eval prompts. +- Real per-sample SPD eval traces can be fed into a Skippy split-stage latency + model to estimate how much pipeline bubble/activation-hop latency SPD can + hide. +- `skippy-runtime` can parse and validate the SPD head manifest/checkpoint + binding, including a Rust-readable safetensors serving checkpoint. It does + not execute the head yet. + +## What Does Not Work Yet + +- We have not trained a topology-independent GLM 4.7 production sidecar for + this branch yet. +- The donor SPD architecture still owns fixed `stage_projs.{stage}` projection + tensors. Use `generic_layer_tap_sidecar.py` for the topology-independent path + instead of extending the donor head further. +- We have not established generic GLM sidecar acceptance/EAL across + `N=1,2,4,8`. +- Skippy/Rust does not yet run the SPD head forward pass. +- Skippy does not yet expose the live hidden-state taps required by the head. +- No live Skippy generation request has used trained SPD proposals yet. +- The `.pt` checkpoint is a proof/training artifact. Export it to + `spd-head.safetensors` before Rust-side serving work. + +## Open Training Data + +The local Qwen3-0.6B proof uses: + +- dataset: `HuggingFaceH4/ultrachat_200k` +- split: `train_sft` +- rows: first `1024` rows for the recorded local proof + +The reference SPD repository lists the intended training corpus family as: + +- UltraChat-200k +- ShareGPT +- SmolTalk +- SmolTalk-Chinese + +MT-Bench, HumanEval, and GSM8K prompts are used here only for evaluation. + +## GLM 4.7 Frontload Smoke Path + +`glm47_frontload.py` inspects a local GLM 4.7 checkpoint and writes a tiny +manifest-compatible SPD smoke artifact before any long training run. This is +for frontloading integration risk: GLM model metadata, non-uniform stage +boundaries, hidden-state tap rows, and Rust manifest validation. The generated +weights are shape fixtures, not a trained SPD head. + +The default local checkpoint path points at the cached `zai-org/GLM-4.7-Flash` +snapshot when present. Override it with `--model-path` on another machine. + +```bash +python evals/spd/glm47_frontload.py \ + --model-path /path/to/GLM-4.7-Flash \ + --work-dir /tmp/skippy-spd-glm47-frontload \ + --num-stages 3 \ + --stage-layer-boundaries 15,31,47 \ + --num-spec-layers 1 \ + --draft-vocab-size 8 \ + --write-smoke-artifacts +``` + +The command writes: + +- `glm47-spd-frontload.json` — checkpoint inspection and derived topology +- `speculation_head_final.pt` — placeholder provenance file +- `spd-head.safetensors` — tiny Rust-readable serving shape fixture +- `skippy-spd-head.json` — manifest with `stage_layer_boundaries` + +Validate the smoke manifest without building native llama.cpp: + +```bash +SKIPPY_SPD_MANIFEST=/tmp/skippy-spd-glm47-frontload/glm47-spd-frontload/skippy-spd-head.json \ + cargo test -p skippy-runtime --features dynamic-native-runtime \ + validates_external_manifest_when_skippy_spd_manifest_is_set +``` + +The inspected local GLM 4.7 Flash checkpoint currently reports +`model_type = glm4_moe_lite`, architecture `Glm4MoeLiteForCausalLM`, +`num_hidden_layers = 47`, `hidden_size = 2048`, `vocab_size = 154880`, and +layer-47 auxiliary tensors `eh_proj`, `enorm`, and `hnorm`. Because 47 target +layers do not divide evenly into the old equal-stage assumptions, GLM SPD +manifests carry explicit `stage_layer_boundaries`. + +### GLM Training Smoke Command + +After the frontload smoke passes, run a tiny real training smoke with a +tokenizer-specific draft vocab. This is not a model-quality run; it verifies +that the reference trainer can load GLM, accept the non-uniform tap topology, +produce a real `speculation_head_final.pt`, and write a Skippy manifest. + +```bash +python evals/spd/hf_train_eval_qwen06.py \ + --work-dir /tmp/skippy-spd-glm47-train-smoke \ + --model-name /path/to/GLM-4.7-Flash \ + --dataset HuggingFaceH4/ultrachat_200k \ + --dataset-split train_sft \ + --train-rows 8 \ + --skip-eval \ + --num-stages 3 \ + --stage-layer-boundaries 15,31,47 \ + --num-spec-layers 1 \ + --max-length 128 \ + --batch-size 1 \ + --gradient-accumulation-steps 1 \ + --build-draft-vocab-size 1024 \ + --draft-vocab-json '' \ + --device cuda \ + --upload-repo none +``` + +`--stage-layer-boundaries` derives the reference trainer's +`--shallow_hidden_layer_indices` as `0,15,31,47;0,15,31;0,15`. You can override +that directly with `--shallow-hidden-layer-indices` when testing another tap +layout. `--build-draft-vocab-size` builds a GLM-tokenizer draft vocab from the +loaded training rows and passes the generated JSON into the reference trainer. + +The first GLM 4.7 smoke artifact is uploaded to the private Hugging Face model +repo `meshllm/skippy-spd-glm47-train-smoke`. It contains the Skippy manifest, +the Rust-readable `spd-head.safetensors` export, the original reference +`speculation_head_final.pt`, and a smoke-focused model card. This artifact is +for training/export/manifest validation only, not production-quality decoding. + +## Generic Topology Plan + +The generic GLM 4.7 sidecar target is not "train one head for +`15,31,47`." Skippy nodes may carry any contiguous layer ranges, so the head +must learn from logical hidden-state evidence: + +- hidden-state index `0` means token embeddings before layer 0 +- hidden-state index `k` means the output after target layer `k - 1` +- physical hosts only decide which logical taps are cheap to expose + +Use `--topology-policy generic-plan` to write randomized contiguous-layer +layouts and tap rows before implementing or launching generic training: + +```bash +python evals/spd/hf_train_eval_qwen06.py \ + --topology-policy generic-plan \ + --model-name /path/to/GLM-4.7-Flash \ + --topology-plan-samples 32 \ + --topology-min-stages 2 \ + --topology-max-stages 6 \ + --topology-tap-dropout 0.25 \ + --num-spec-layers 4 \ + --draft-top-k 4 \ + --upload-repo none +``` + +The command writes `topology/skippy-spd-topology-plan.json` under the run +artifact directory and exits before cloning or training. That exit is +intentional: the current reference implementation would otherwise produce a +fixed-stage head. The next model patch should consume these logical tap plans +with masks or tap dropout so one exported sidecar can be evaluated against many +candidate Skippy contiguous-layer topologies. + +## Generic Layer-Tap Sidecar + +`generic_layer_tap_sidecar.py` is the first non-donor sidecar path. It trains a +small token oracle over a set of logical hidden-state taps: + +- `hidden[layer_index = 0]`: token embeddings before target layer 0 +- `hidden[layer_index = k]`: output after target layer `k - 1` +- tap features: normalized layer depth plus an embedding-row flag +- tap mask/dropout: randomly withholds intermediate taps during training + +The exported manifest uses: + +- `source.format = generic-layer-tap-sidecar-v1` +- `topology.head_kind = generic-layer-tap-v1` +- serving tensors such as `tap_proj.*`, `depth_proj.*`, `tap_norm.*`, + `output_norm.*`, and `draft_heads.{n}.*` + +Run a local contract smoke without loading GLM: + +```bash +uv run evals/spd/generic_layer_tap_sidecar.py \ + --smoke-synthetic \ + --work-dir /tmp/skippy-spd-generic-layer-tap-smoke \ + --model-name GLM-4.7-Flash-shape-only \ + --topology-num-hidden-layers 47 \ + --topology-plan-samples 8 \ + --topology-min-stages 2 \ + --topology-max-stages 4 \ + --topology-tap-dropout 0.25 \ + --num-spec-layers 2 \ + --draft-top-k 1 \ + --draft-vocab-size 64 \ + --synthetic-hidden-size 32 \ + --synthetic-vocab-size 128 \ + --synthetic-train-examples 48 \ + --synthetic-eval-examples 24 \ + --batch-size 8 \ + --epochs 1 \ + --device cpu \ + --export-dtype float32 +``` + +Validate the exported generic manifest with Rust: + +```bash +SKIPPY_SPD_MANIFEST=/tmp/skippy-spd-generic-layer-tap-smoke/artifacts//train/skippy-spd-head.json \ + cargo test -p skippy-runtime --lib \ + validates_external_manifest_when_skippy_spd_manifest_is_set +``` + +Run a small real GLM 4.7 quality gate by replacing `--smoke-synthetic` with the +local checkpoint path and small train/eval row counts first: + +```bash +export HF_TOKEN="$(cat /Volumes/models/huggingface/token 2>/dev/null || true)" +export HF_HOME=/tmp/codex-hf-home +export HF_HUB_CACHE=/tmp/codex-hf-hub +export HF_DATASETS_CACHE=/tmp/codex-hf-datasets + +uv run evals/spd/generic_layer_tap_sidecar.py \ + --model-name /path/to/GLM-4.7-Flash \ + --work-dir /tmp/skippy-spd-glm47-generic-layer-tap-n1 \ + --train-rows 128 \ + --eval-rows 16 \ + --positions-per-row 4 \ + --max-length 512 \ + --num-spec-layers 1 \ + --draft-vocab-size 4096 \ + --topology-plan-samples 32 \ + --topology-min-stages 2 \ + --topology-max-stages 6 \ + --topology-tap-dropout 0.25 \ + --batch-size 8 \ + --epochs 1 \ + --device cuda \ + --export-dtype float16 +``` + +On `micstudio`, keep the model path on the shared NFS mount but use local +`/tmp` HF caches for dataset/Hub metadata. The NFS-backed shared cache can fail +dataset metadata locks with `OSError: [Errno 77] No locks available`. + +First recorded real GLM mechanics gate: + +| Field | Value | +| --- | --- | +| Artifact | `/Users/micn/spd-runs/glm47-generic-layer-tap-n1-smoke/artifacts/20260617-221949` | +| Model type | `glm4_moe_lite` | +| Hidden / layers / vocab | `2048` / `47` / `154880` | +| `num_spec_layers` | `1` | +| Topology samples | `8` layouts, stages `2..4`, tap dropout `0.25` | +| Draft vocab | `702` observed ids from 2 training rows | +| Eval examples | `1` | +| Draft vocab label coverage | `1.0` | +| Acceptance | `0.0` | +| Equivalent accept length | `0.0` | +| Proposal latency | `183.98 ms/example` | +| Eval wall time | `0.209 s` | +| Serving checkpoint | `spd-head.safetensors`, `F16`, `11.3 MB`, sha256 `3afea1ba44fc32ea9c77237952eb71010a0ba4d05d17cb6e756ffb799cfb3904` | +| Rust validation | `SKIPPY_SPD_MANIFEST=... cargo test -p skippy-runtime --lib validates_external_manifest_when_skippy_spd_manifest_is_set` passed | + +This is a tokenizer/model/export/manifest mechanics pass, not a quality result. +The next quality gate needs more rows, more positions, and then the `N=1,2,4` +sweep. + +To avoid paying full GLM target-forward cost for every sidecar architecture +iteration, write a reusable hidden-state example cache. Prefer generating the +cache with the widest draft window you plan to test, for example +`--num-spec-layers 4`; narrower `N=1` and `N=2` runs can load the same cache and +truncate labels during training/eval. + +```bash +uv run evals/spd/generic_layer_tap_sidecar.py \ + --model-name /path/to/GLM-4.7-Flash \ + --examples-cache-out /tmp/glm47-layer-tap-examples.pt \ + --num-spec-layers 4 \ + ...same extraction/training flags... +``` + +Then train another sidecar on the exact same examples without reloading GLM: + +```bash +uv run evals/spd/generic_layer_tap_sidecar.py \ + --examples-cache-in /tmp/glm47-layer-tap-examples.pt \ + --encoder mean_mlp \ + --mlp-ratio 2.0 \ + --num-spec-layers 2 \ + --eval-top-k 1,2,4,8,16,32 \ + --batch-size 16 \ + --epochs 3 \ + --device mps \ + --export-dtype float16 +``` + +Supported encoders: + +- `--encoder mean`: masked mean pooling over encoded taps +- `--encoder mean_mlp`: masked mean pooling plus a residual MLP over the pooled + representation +- `--encoder attention`: learned query plus multi-head attention over encoded + taps; this has not beaten mean pooling in the current GLM 4.7 gate + +Use `--eval-top-k` to report all-label diagnostic hit rates. These diagnostics +are separate from sequential acceptance/EAL and are intended to show whether +the correct target token is close to the top of the sidecar distribution. + +All encoder variants export `generic-layer-tap-v1` safetensors manifests that +validate through `skippy-runtime`; the manifest records the exact safetensors +tensor count for encoder-specific tensors. + +For the fixed GLM layer-tap control, bypass randomized layer-tap sampling and +pin the logical evidence schedule: + +```bash +uv run evals/spd/generic_layer_tap_sidecar.py \ + --model-name /path/to/GLM-4.7-Flash \ + --fixed-layer-taps 0,12,24,35,47 \ + --extract-batch-size 4 \ + --examples-cache-out /tmp/glm47-fixed-s4-examples-n4.pt \ + --num-spec-layers 4 \ + ...same extraction/training flags... +``` + +The fixed taps are logical hidden-state indices, not Skippy host or stage IDs. +Use this control to prove GLM-specific SPD learnability before returning to +randomized layer-tap training. Increase `--extract-batch-size` only while hidden +state extraction fits comfortably in memory. + +## Reproduce Qwen3-0.6B Training + +This is the smallest useful proof that the training path and artifact shape +work. It trains a real head from open data. + +```bash +python evals/spd/hf_train_eval_qwen06.py \ + --work-dir /tmp/skippy-spd-qwen06-proof \ + --model-name Qwen/Qwen3-0.6B \ + --dataset HuggingFaceH4/ultrachat_200k \ + --dataset-split train_sft \ + --train-rows 1024 \ + --eval-rows-per-set 8 \ + --num-stages 2 \ + --num-spec-layers 4 \ + --max-length 256 \ + --max-new-tokens 64 \ + --draft-top-k 4 \ + --device mps \ + --upload-repo none +``` + +Use `--device cuda` on a GPU host. The runner also supports HF Jobs, but that is +only a convenience wrapper; the proof is ordinary Python plus open data. + +Recorded local result: + +| Model | Head | Eval draft top-k | Generated tokens | Accepted flags | Acceptance | Equivalent accept length | Theoretical gain | +| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | +| `Qwen/Qwen3-0.6B` | locally trained, 4 spec layers | 4 | 1536 | 326 / 1536 | 0.5628 | 1.1257 | 12.67% | + +This proves the training/export path, but it is not the high-gain target. + +## Reproduce Qwen3.5-4B Pretrained Head Eval + +This is the strongest current model-quality signal. It uses an author-published +SPD head and evaluates it locally against the reference verifier. + +```bash +python evals/spd/hf_train_eval_qwen06.py \ + --work-dir /tmp/skippy-spd-qwen35-4b-pretrained-s4l4 \ + --model-name Qwen/Qwen3.5-4B \ + --spec-head-repo yuyijiong/speculative_pipeline_decoding \ + --spec-head-file Qwen3.5-4B_s4_l4.pt \ + --manifest-base-model-path Qwen/Qwen3.5-4B \ + --skip-train \ + --device mps \ + --eval-rows-per-set 8 \ + --max-new-tokens 64 \ + --draft-top-k 4 \ + --upload-repo none +``` + +Use `--device cuda` on a GPU host. The first run downloads the base model and +the SPD head. + +Recorded local result: + +| Model | Head | Eval draft top-k | Generated tokens | Accepted flags | Acceptance | Equivalent accept length | Theoretical gain | +| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | +| `Qwen/Qwen3.5-4B` | pretrained, 4 stages / 4 spec layers | 4 | 1536 | 1230 / 1536 | 0.6176 | 2.4704 | 163.39% | + +Per-dataset theoretical gains from the same run: + +| Dataset | Acceptance | Equivalent accept length | Theoretical gain | +| --- | ---: | ---: | ---: | +| MT-Bench | 0.4918 | 1.9673 | 98.42% | +| HumanEval | 0.8797 | 3.5189 | 254.18% | +| GSM8K | 0.5926 | 2.3704 | 137.58% | + +## Latency Simulation From Real Traces + +`simulate_latency.py` consumes the raw `eval/raw/*per_sample.jsonl` file emitted +by the reference evaluator. It does not invent acceptance; it uses the real +`new_tokens`, `decode_loop_steps`, and accepted-flag counters from the run. + +```bash +python evals/spd/simulate_latency.py \ + --raw /tmp/skippy-spd-qwen35-4b-pretrained-s4l4/artifacts//eval/raw/pipeline_eval__train__speculation_head_final__nt24__per_sample.jsonl \ + --stage-ms 4,4,4,4 \ + --hop-ms 0,1,5,10,25 +``` + +Recorded Qwen3.5-4B trace with a four-stage `4ms,4ms,4ms,4ms` model: + +| Hop ms | Serial split tok/s | SPD pipeline tok/s | SPD vs serial split | Paper-like gain | P50 serial ms | P50 SPD ms | +| ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 0 | 62.50 | 617.61 | 9.882x | 2.470x | 1024.00 | 106.50 | +| 1 | 52.63 | 494.09 | 9.388x | 2.470x | 1216.00 | 133.12 | +| 5 | 32.26 | 274.49 | 8.509x | 2.470x | 1984.00 | 239.62 | +| 10 | 21.74 | 176.46 | 8.117x | 2.470x | 2944.00 | 372.75 | +| 25 | 10.99 | 85.19 | 7.752x | 2.470x | 5824.00 | 772.12 | + +The `paper-like gain` column is based on the SPD trace alone. The `SPD vs serial +split` column models a Skippy-specific comparison where ordinary split serving +must traverse every stage/hop for each generated token before the next target +token is known. + +## Export the Serving Checkpoint + +After training or downloading a reference SPD head, export the PyTorch +checkpoint to a Rust-readable serving artifact: + +```bash +python evals/spd/export_spd_head.py \ + --checkpoint /tmp/skippy-spd-qwen35-4b-pretrained-s4l4/artifacts//train/speculation_head_final.pt \ + --manifest /tmp/skippy-spd-qwen35-4b-pretrained-s4l4/artifacts//train/skippy-spd-head.json \ + --base-model-path Qwen/Qwen3.5-4B +``` + +The exporter writes `spd-head.safetensors` next to the manifest and adds an +optional `serving_checkpoint` section to `skippy-spd-head.json`. The original +`.pt` checkpoint remains referenced for provenance. + +Validate an exported local head through Rust with: + +```bash +SKIPPY_SPD_MANIFEST=/tmp/skippy-spd-qwen35-4b-pretrained-s4l4/artifacts//train/skippy-spd-head.json \ + cargo test -p skippy-runtime validates_external_manifest_when_skippy_spd_manifest_is_set +``` + +## Artifact Contract + +The proof runner writes: + +- `train/speculation_head_final.pt` +- `train/spd-head.safetensors` after export +- `train/skippy-spd-head.json` +- `eval/raw/*.jsonl` +- `eval/summary/*.json` + +The manifest schema is `skippy-spd-head/v1`. It binds a head checkpoint to: + +- base model path/id +- checkpoint format/version +- checkpoint byte size and sha256 +- hidden size +- base vocab size +- draft vocab size and optional draft token ids +- number of target stages +- number of spec layers +- shallow hidden-layer tap indices +- optional safetensors serving checkpoint path, size, checksum, tensor count, + and dtype + +Rust validation lives in `crates/skippy-runtime/src/spd.rs`. + +## Next Engineering Steps + +1. Add a tensor loader for the SPD head weights and draft vocab mapping. +2. Implement the Qwen3.5-4B SPD forward pass in Rust for the recorded topology. +3. Capture Skippy hidden-state taps and compare Rust top-k proposals to the + Python reference on the same taps. +4. Wire live proposal generation into `skippy-server`. +5. Verify every accepted token through the normal target stages. +6. Use the Speedy benchmark for the final vanilla target versus verified + target+SPD sidecar comparison; keep latency simulation as supporting + analysis only. + +## Next Research Steps + +1. Train a head for a larger Qwen-family model to prove scaling beyond the + pretrained 4B artifact. +2. Keep the draft vocab capped at 32k or 50k first. +3. Record acceptance, equivalent accept length, and latency simulation from the + same eval prompts. +4. Only after that, evaluate custom large MoE targets. Very large MoE models + need activation-capture support and are not the right first scaling proof. diff --git a/evals/spd/export_spd_head.py b/evals/spd/export_spd_head.py new file mode 100644 index 0000000000..d6e0554363 --- /dev/null +++ b/evals/spd/export_spd_head.py @@ -0,0 +1,197 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "numpy", +# "safetensors>=0.5.0", +# "torch>=2.8.0", +# ] +# /// +"""Export a reference SPD PyTorch checkpoint into a Skippy serving artifact.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path +from typing import Any + + +SERVING_FORMAT = "safetensors-spd-head-v1" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Export an SPD .pt checkpoint to a Rust-readable safetensors artifact" + ) + parser.add_argument("--checkpoint", required=True, help="Input speculation_head_final.pt") + parser.add_argument("--manifest", required=True, help="Input skippy-spd-head.json") + parser.add_argument( + "--manifest-out", + default="", + help="Manifest to write. Defaults to updating --manifest in place.", + ) + parser.add_argument( + "--out-dir", + default="", + help="Output directory. Defaults to the manifest output directory.", + ) + parser.add_argument("--out-name", default="spd-head.safetensors") + parser.add_argument( + "--dtype", + choices=("keep", "bfloat16", "float16", "float32"), + default="keep", + help="Tensor dtype for the serving artifact.", + ) + parser.add_argument( + "--base-model-path", + default="", + help="Optional manifest base_model_path override, for portable public manifests.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + checkpoint_path = Path(args.checkpoint) + manifest_path = Path(args.manifest) + manifest_out = Path(args.manifest_out) if args.manifest_out else manifest_path + out_dir = Path(args.out_dir) if args.out_dir else manifest_out.parent + output_path = out_dir / args.out_name + + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + tensors, checkpoint_config = load_tensors(checkpoint_path, args.dtype) + out_dir.mkdir(parents=True, exist_ok=True) + + from safetensors.torch import save_file + + dtype_label = common_dtype_label(tensors) + metadata = { + "format": SERVING_FORMAT, + "source_checkpoint_sha256": manifest["checkpoint"]["sha256"], + "source_format": manifest["source"]["format"], + "tensor_count": str(len(tensors)), + "dtype": dtype_label, + } + base_model_path = args.base_model_path.strip() + if base_model_path: + manifest["source"]["base_model_path"] = base_model_path + metadata["base_model_path"] = manifest["source"]["base_model_path"] + + save_file(tensors, output_path, metadata=metadata) + + manifest["serving_checkpoint"] = { + "path": manifest_relative_path(output_path, manifest_out.parent), + "sha256": file_sha256(output_path), + "bytes": output_path.stat().st_size, + "format": SERVING_FORMAT, + "tensor_count": len(tensors), + "dtype": dtype_label, + } + if checkpoint_config.get("version") is not None: + manifest["source"]["checkpoint_version"] = int(checkpoint_config["version"]) + + manifest_out.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print( + json.dumps( + { + "serving_checkpoint": str(output_path), + "manifest": str(manifest_out), + "format": SERVING_FORMAT, + "tensor_count": len(tensors), + "dtype": dtype_label, + "bytes": output_path.stat().st_size, + "sha256": manifest["serving_checkpoint"]["sha256"], + }, + indent=2, + sort_keys=True, + ) + ) + + +def load_tensors(checkpoint_path: Path, dtype: str) -> tuple[dict[str, Any], dict[str, Any]]: + import torch + + try: + checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + except TypeError: + checkpoint = torch.load(checkpoint_path, map_location="cpu") + + if not isinstance(checkpoint, dict): + raise RuntimeError(f"{checkpoint_path} must contain a dict checkpoint") + + state_dict = checkpoint.get("state_dict") or checkpoint.get("model_state_dict") + if state_dict is None: + state_dict = checkpoint + if not isinstance(state_dict, dict): + raise RuntimeError(f"{checkpoint_path} does not contain a state dict") + + target_dtype = { + "keep": None, + "bfloat16": torch.bfloat16, + "float16": torch.float16, + "float32": torch.float32, + }[dtype] + + tensors: dict[str, Any] = {} + for name in sorted(state_dict): + value = state_dict[name] + if not torch.is_tensor(value): + continue + tensor = value.detach().cpu().contiguous() + if target_dtype is not None: + tensor = tensor.to(target_dtype) + tensors[name] = tensor + + if not tensors: + raise RuntimeError(f"{checkpoint_path} did not contain any tensors") + config = checkpoint.get("config") if isinstance(checkpoint.get("config"), dict) else {} + return tensors, config + + +def common_dtype_label(tensors: dict[str, Any]) -> str: + import torch + + labels = { + torch.bfloat16: "BF16", + torch.float16: "F16", + torch.float32: "F32", + torch.float64: "F64", + torch.int64: "I64", + torch.int32: "I32", + torch.int16: "I16", + torch.int8: "I8", + torch.uint8: "U8", + torch.bool: "BOOL", + } + seen = {labels.get(tensor.dtype, str(tensor.dtype)) for tensor in tensors.values()} + if len(seen) == 1: + return next(iter(seen)) + return "mixed" + + +def manifest_relative_path(path: Path, manifest_dir: Path) -> str: + resolved_path = path.resolve() + resolved_manifest_dir = manifest_dir.resolve() + try: + relative = resolved_path.relative_to(resolved_manifest_dir) + except ValueError as exc: + raise RuntimeError( + f"{path} must be inside the manifest directory {manifest_dir}" + ) from exc + if any(part in ("", ".", "..") for part in relative.parts): + raise RuntimeError(f"unsafe manifest-relative path: {relative}") + return relative.as_posix() + + +def file_sha256(path: Path) -> str: + hasher = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + hasher.update(chunk) + return hasher.hexdigest() + + +if __name__ == "__main__": + main() diff --git a/evals/spd/generic_layer_tap_sidecar.py b/evals/spd/generic_layer_tap_sidecar.py new file mode 100644 index 0000000000..44ebe7af6e --- /dev/null +++ b/evals/spd/generic_layer_tap_sidecar.py @@ -0,0 +1,1053 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "accelerate>=1.0.0", +# "datasets>=3.0.0", +# "numpy", +# "safetensors>=0.5.0", +# "torch>=2.8.0", +# "transformers>=5.6.0", +# ] +# /// +"""Train/export/evaluate a topology-independent SPD layer-tap sidecar. + +This is the generic GLM 4.7 path. Unlike the donor SPD head, this sidecar does +not own per-stage projection tensors. It consumes a set of logical hidden-state +taps plus tap features, then predicts the next N draft tokens. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import random +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +from torch import nn + +from topology_plan import build_topology_plan + + +SERVING_FORMAT = "safetensors-spd-head-v1" +SOURCE_FORMAT = "generic-layer-tap-sidecar-v1" +HEAD_KIND = "generic-layer-tap-v1" + + +@dataclass +class TapExample: + hidden: torch.Tensor + features: torch.Tensor + labels: list[int] + topology_key: str + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Train a generic SPD layer-tap sidecar") + parser.add_argument("--model-name", default="GLM-4.7-Flash-shape-only") + parser.add_argument("--dataset", default="HuggingFaceH4/ultrachat_200k") + parser.add_argument("--dataset-split", default="train_sft") + parser.add_argument("--work-dir", default="/tmp/skippy-spd-generic-layer-tap") + parser.add_argument("--topology-plan", default="") + parser.add_argument("--topology-plan-samples", type=int, default=32) + parser.add_argument("--topology-min-stages", type=int, default=2) + parser.add_argument("--topology-max-stages", type=int, default=6) + parser.add_argument("--topology-seed", type=int, default=47) + parser.add_argument("--topology-tap-dropout", type=float, default=0.25) + parser.add_argument("--topology-num-hidden-layers", type=int, default=47) + parser.add_argument( + "--fixed-layer-taps", + default="", + help=( + "Comma-separated logical hidden-state indices for a fixed layer-tap " + "control run, for example 0,12,24,35,47." + ), + ) + parser.add_argument("--num-spec-layers", type=int, default=1) + parser.add_argument("--draft-top-k", type=int, default=1) + parser.add_argument("--draft-vocab-size", type=int, default=4096) + parser.add_argument("--train-rows", type=int, default=128) + parser.add_argument("--eval-rows", type=int, default=16) + parser.add_argument("--positions-per-row", type=int, default=4) + parser.add_argument("--max-length", type=int, default=512) + parser.add_argument("--extract-batch-size", type=int, default=1) + parser.add_argument("--epochs", type=int, default=1) + parser.add_argument("--batch-size", type=int, default=8) + parser.add_argument("--learning-rate", type=float, default=3e-4) + parser.add_argument("--examples-cache-in", default="") + parser.add_argument("--examples-cache-out", default="") + parser.add_argument("--encoder", choices=("mean", "mean_mlp", "attention"), default="mean") + parser.add_argument("--attention-heads", type=int, default=4) + parser.add_argument("--mlp-ratio", type=float, default=2.0) + parser.add_argument("--eval-top-k", default="1,2,4,8,16,32") + parser.add_argument("--device", choices=("auto", "cuda", "mps", "cpu"), default="auto") + parser.add_argument("--dtype", choices=("float32", "float16", "bfloat16"), default="float32") + parser.add_argument("--attn-implementation", default="sdpa") + parser.add_argument( + "--model-device-map", + choices=("single", "auto", "cpu"), + default="single", + help="Device map for loading the target model during hidden-state extraction.", + ) + parser.add_argument("--local-files-only", action="store_true") + parser.add_argument("--export-dtype", choices=("float32", "float16", "bfloat16"), default="float16") + parser.add_argument("--smoke-synthetic", action="store_true") + parser.add_argument("--synthetic-hidden-size", type=int, default=64) + parser.add_argument("--synthetic-vocab-size", type=int, default=512) + parser.add_argument("--synthetic-train-examples", type=int, default=64) + parser.add_argument("--synthetic-eval-examples", type=int, default=32) + parser.add_argument("--manifest-base-model-path", default="") + return parser.parse_args() + + +class GenericLayerTapSidecar(nn.Module): + def __init__( + self, + hidden_size: int, + draft_vocab_size: int, + num_spec_layers: int, + *, + encoder_kind: str = "mean", + attention_heads: int = 4, + mlp_ratio: float = 2.0, + ) -> None: + super().__init__() + self.encoder_kind = encoder_kind + self.tap_proj = nn.Linear(hidden_size, hidden_size) + self.depth_proj = nn.Linear(2, hidden_size) + self.tap_norm = nn.LayerNorm(hidden_size) + if encoder_kind == "attention": + self.attn_query = nn.Parameter(torch.zeros(1, 1, hidden_size)) + self.tap_attention = nn.MultiheadAttention( + embed_dim=hidden_size, + num_heads=attention_heads, + batch_first=True, + ) + self.attn_norm = nn.LayerNorm(hidden_size) + elif encoder_kind == "mean_mlp": + mlp_hidden_size = max(1, int(hidden_size * mlp_ratio)) + self.pool_mlp = nn.Sequential( + nn.Linear(hidden_size, mlp_hidden_size), + nn.GELU(), + nn.Linear(mlp_hidden_size, hidden_size), + ) + elif encoder_kind != "mean": + raise ValueError(f"unsupported encoder_kind: {encoder_kind}") + self.output_norm = nn.LayerNorm(hidden_size) + self.draft_heads = nn.ModuleList( + [nn.Linear(hidden_size, draft_vocab_size) for _ in range(num_spec_layers)] + ) + + def forward( + self, + hidden: torch.Tensor, + features: torch.Tensor, + mask: torch.Tensor, + ) -> list[torch.Tensor]: + encoded = torch.tanh(self.tap_norm(self.tap_proj(hidden) + self.depth_proj(features))) + if self.encoder_kind == "attention": + query = self.attn_query.expand(encoded.shape[0], -1, -1) + attended, _ = self.tap_attention( + query, + encoded, + encoded, + key_padding_mask=~mask, + need_weights=False, + ) + pooled = self.output_norm(self.attn_norm(attended[:, 0, :])) + else: + masked = encoded * mask.unsqueeze(-1).to(encoded.dtype) + denom = mask.sum(dim=1, keepdim=True).clamp_min(1).to(encoded.dtype) + pooled = masked.sum(dim=1) / denom + if self.encoder_kind == "mean_mlp": + pooled = pooled + self.pool_mlp(pooled) + pooled = self.output_norm(pooled) + return [head(pooled) for head in self.draft_heads] + + +def main() -> None: + args = parse_args() + started = time.perf_counter() + artifact_dir = Path(args.work_dir).expanduser().resolve() / "artifacts" / timestamp() + train_dir = artifact_dir / "train" + eval_dir = artifact_dir / "eval" / "summary" + train_dir.mkdir(parents=True, exist_ok=True) + eval_dir.mkdir(parents=True, exist_ok=True) + + plan = load_or_build_topology_plan(args) + cache_in = Path(args.examples_cache_in).expanduser() if args.examples_cache_in else None + if cache_in is not None: + model_meta, draft_token_ids, train_examples, eval_examples, plan = load_examples_cache(cache_in) + elif args.smoke_synthetic: + model_meta, draft_token_ids, train_examples, eval_examples = synthetic_examples(args, plan) + else: + model_meta, draft_token_ids, train_examples, eval_examples = real_glm_examples(args, plan) + if args.examples_cache_out and cache_in is None: + write_examples_cache( + Path(args.examples_cache_out).expanduser(), + model_meta, + draft_token_ids, + train_examples, + eval_examples, + plan, + args, + ) + train_examples = align_example_width(train_examples, int(args.num_spec_layers), "train") + eval_examples = align_example_width(eval_examples, int(args.num_spec_layers), "eval") + + device = resolve_device(args.device) + sidecar = GenericLayerTapSidecar( + hidden_size=int(model_meta["hidden_size"]), + draft_vocab_size=len(draft_token_ids), + num_spec_layers=int(args.num_spec_layers), + encoder_kind=args.encoder, + attention_heads=int(args.attention_heads), + mlp_ratio=float(args.mlp_ratio), + ).to(device) + optimizer = torch.optim.AdamW(sidecar.parameters(), lr=float(args.learning_rate)) + label_map = {token_id: index for index, token_id in enumerate(draft_token_ids)} + + train_started = time.perf_counter() + train_loss = train_sidecar( + sidecar, + optimizer, + train_examples, + label_map, + batch_size=int(args.batch_size), + epochs=int(args.epochs), + device=device, + ) + train_elapsed = time.perf_counter() - train_started + + eval_summary = evaluate_sidecar( + sidecar, + eval_examples, + label_map, + batch_size=int(args.batch_size), + device=device, + top_k_values=parse_top_k_values(args.eval_top_k), + ) + eval_summary["train_loss"] = train_loss + eval_summary["train_wall_seconds"] = train_elapsed + eval_summary["total_wall_seconds"] = time.perf_counter() - started + eval_summary["encoder"] = { + "kind": args.encoder, + "attention_heads": int(args.attention_heads) if args.encoder == "attention" else None, + "mlp_ratio": float(args.mlp_ratio) if args.encoder == "mean_mlp" else None, + } + eval_summary["topology_policy"] = plan["policy"] + eval_summary["topology_eval"] = summarize_examples_by_topology(eval_examples) + eval_summary["model"] = model_meta + + checkpoint_path = train_dir / "generic-layer-tap-sidecar.pt" + checkpoint = { + "format": SOURCE_FORMAT, + "version": 1, + "config": { + "head_kind": HEAD_KIND, + "hidden_size": int(model_meta["hidden_size"]), + "vocab_size": int(model_meta["vocab_size"]), + "draft_vocab_size": len(draft_token_ids), + "num_spec_layers": int(args.num_spec_layers), + "encoder_kind": args.encoder, + "attention_heads": int(args.attention_heads), + "mlp_ratio": float(args.mlp_ratio) if args.encoder == "mean_mlp" else None, + "max_taps": max_taps(plan), + "tap_feature_size": 2, + "draft_token_ids": draft_token_ids, + "topology_plan": plan, + }, + "state_dict": {name: value.detach().cpu() for name, value in sidecar.state_dict().items()}, + "eval_summary": eval_summary, + } + torch.save(checkpoint, checkpoint_path) + + serving_path = train_dir / "spd-head.safetensors" + save_serving_safetensors(sidecar, serving_path, args.export_dtype) + manifest_path = train_dir / "skippy-spd-head.json" + write_manifest( + args=args, + manifest_path=manifest_path, + checkpoint_path=checkpoint_path, + serving_path=serving_path, + model_meta=model_meta, + draft_token_ids=draft_token_ids, + plan=plan, + ) + summary_path = eval_dir / "generic_layer_tap_eval_summary.json" + summary_path.write_text(json.dumps(eval_summary, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + print( + json.dumps( + { + "artifact_dir": str(artifact_dir), + "manifest": str(manifest_path), + "serving_checkpoint": str(serving_path), + "summary": str(summary_path), + "acceptance_rate": eval_summary["acceptance_rate"], + "equivalent_accept_length": eval_summary["equivalent_accept_length"], + "proposal_latency_ms": eval_summary["proposal_latency_ms"], + "eval_wall_seconds": eval_summary["eval_wall_seconds"], + "top_k": eval_summary["top_k_diagnostics"]["overall"], + }, + indent=2, + sort_keys=True, + ) + ) + + +def load_or_build_topology_plan(args: argparse.Namespace) -> dict[str, Any]: + if args.topology_plan: + return json.loads(Path(args.topology_plan).expanduser().read_text(encoding="utf-8")) + return build_topology_plan(args) + + +def write_examples_cache( + path: Path, + model_meta: dict[str, Any], + draft_token_ids: list[int], + train_examples: list[TapExample], + eval_examples: list[TapExample], + plan: dict[str, Any], + args: argparse.Namespace, +) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + payload = { + "schema": "skippy-spd-layer-tap-examples/v1", + "created_at": timestamp(), + "model": model_meta, + "draft_token_ids": draft_token_ids, + "topology_plan": plan, + "metadata": { + "num_spec_layers": int(args.num_spec_layers), + "train_rows": int(args.train_rows), + "eval_rows": int(args.eval_rows), + "positions_per_row": int(args.positions_per_row), + "max_length": int(args.max_length), + }, + "train_examples": [example_to_record(example) for example in train_examples], + "eval_examples": [example_to_record(example) for example in eval_examples], + } + torch.save(payload, path) + print(f"wrote layer-tap examples cache -> {path}", flush=True) + + +def load_examples_cache( + path: Path, +) -> tuple[dict[str, Any], list[int], list[TapExample], list[TapExample], dict[str, Any]]: + try: + payload = torch.load(path, map_location="cpu", weights_only=False) + except TypeError: + payload = torch.load(path, map_location="cpu") + if payload.get("schema") != "skippy-spd-layer-tap-examples/v1": + raise RuntimeError(f"unsupported examples cache schema in {path}") + train_examples = [record_to_example(record) for record in payload["train_examples"]] + eval_examples = [record_to_example(record) for record in payload["eval_examples"]] + print( + f"loaded layer-tap examples cache -> {path} " + f"(train={len(train_examples)}, eval={len(eval_examples)})", + flush=True, + ) + return ( + dict(payload["model"]), + [int(token_id) for token_id in payload["draft_token_ids"]], + train_examples, + eval_examples, + dict(payload["topology_plan"]), + ) + + +def example_to_record(example: TapExample) -> dict[str, Any]: + return { + "hidden": example.hidden.detach().cpu().to(torch.float16), + "features": example.features.detach().cpu().to(torch.float16), + "labels": [int(label) for label in example.labels], + "topology_key": example.topology_key, + } + + +def record_to_example(record: dict[str, Any]) -> TapExample: + return TapExample( + hidden=record["hidden"].float(), + features=record["features"].float(), + labels=[int(label) for label in record["labels"]], + topology_key=str(record["topology_key"]), + ) + + +def align_example_width( + examples: list[TapExample], + num_spec_layers: int, + label: str, +) -> list[TapExample]: + widths = [len(example.labels) for example in examples] + too_narrow = [width for width in widths if width < num_spec_layers] + if too_narrow: + raise RuntimeError( + f"{label} examples cannot satisfy --num-spec-layers={num_spec_layers}; " + f"found label widths including {too_narrow[:5]}" + ) + if any(width > num_spec_layers for width in widths): + print( + f"truncating {label} examples to --num-spec-layers={num_spec_layers}", + flush=True, + ) + return [ + TapExample( + hidden=example.hidden, + features=example.features, + labels=example.labels[:num_spec_layers], + topology_key=example.topology_key, + ) + for example in examples + ] + return examples + + +def synthetic_examples( + args: argparse.Namespace, + plan: dict[str, Any], +) -> tuple[dict[str, Any], list[int], list[TapExample], list[TapExample]]: + rng = random.Random(int(args.topology_seed)) + hidden_size = int(args.synthetic_hidden_size) + vocab_size = int(args.synthetic_vocab_size) + draft_token_ids = list(range(min(int(args.draft_vocab_size), vocab_size))) + model_meta = { + "model_name": args.model_name, + "model_type": "synthetic", + "hidden_size": hidden_size, + "vocab_size": vocab_size, + "num_hidden_layers": int(plan["model"]["num_hidden_layers"]), + } + train = [ + synthetic_example(hidden_size, draft_token_ids, int(args.num_spec_layers), plan, rng) + for _ in range(int(args.synthetic_train_examples)) + ] + evals = [ + synthetic_example(hidden_size, draft_token_ids, int(args.num_spec_layers), plan, rng) + for _ in range(int(args.synthetic_eval_examples)) + ] + return model_meta, draft_token_ids, train, evals + + +def synthetic_example( + hidden_size: int, + draft_token_ids: list[int], + num_spec_layers: int, + plan: dict[str, Any], + rng: random.Random, +) -> TapExample: + indices = sample_tap_indices(plan, rng) + hidden = torch.randn(len(indices), hidden_size) + features = tap_features(indices, int(plan["model"]["num_hidden_layers"])) + labels = [rng.choice(draft_token_ids) for _ in range(num_spec_layers)] + return TapExample( + hidden=hidden, + features=features, + labels=labels, + topology_key=topology_key(indices), + ) + + +def real_glm_examples( + args: argparse.Namespace, + plan: dict[str, Any], +) -> tuple[dict[str, Any], list[int], list[TapExample], list[TapExample]]: + from transformers import AutoModelForCausalLM, AutoTokenizer + + device = resolve_device(args.device) + dtype = torch_dtype(args.dtype) + tokenizer = AutoTokenizer.from_pretrained(args.model_name, trust_remote_code=True) + rows = load_training_rows(args.dataset, args.dataset_split, int(args.train_rows + args.eval_rows)) + texts = [row_text_for_vocab(tokenizer, row.get("messages") or []) for row in rows] + draft_token_ids = build_draft_vocab(tokenizer, texts[: int(args.train_rows)], int(args.draft_vocab_size)) + model = load_target_model(args, dtype, device) + model.eval() + model_meta = { + "model_name": args.model_name, + "model_type": getattr(model.config, "model_type", None), + "hidden_size": int(model.config.hidden_size), + "vocab_size": int(model.config.vocab_size), + "num_hidden_layers": int(model.config.num_hidden_layers), + } + train_texts = texts[: int(args.train_rows)] + eval_texts = texts[int(args.train_rows) : int(args.train_rows + args.eval_rows)] + rng = random.Random(int(args.topology_seed)) + model_device = next(model.parameters()).device + train = collect_examples_from_texts(args, tokenizer, model, train_texts, plan, rng, model_device) + evals = collect_examples_from_texts(args, tokenizer, model, eval_texts, plan, rng, model_device) + return model_meta, draft_token_ids, train, evals + + +def load_target_model(args: argparse.Namespace, dtype: torch.dtype, device: torch.device) -> Any: + from transformers import AutoModelForCausalLM + + kwargs: dict[str, Any] = { + "trust_remote_code": True, + "torch_dtype": dtype, + "attn_implementation": args.attn_implementation, + "low_cpu_mem_usage": True, + "local_files_only": bool(args.local_files_only), + } + if args.model_device_map == "auto": + kwargs["device_map"] = "auto" + elif args.model_device_map == "cpu": + kwargs["device_map"] = {"": "cpu"} + elif device.type != "cpu": + kwargs["device_map"] = {"": device.type} + model = AutoModelForCausalLM.from_pretrained(args.model_name, **kwargs) + if args.model_device_map == "single" and device.type == "cpu": + model = model.to(device) + return model + + +def collect_examples_from_texts( + args: argparse.Namespace, + tokenizer: Any, + model: Any, + texts: list[str], + plan: dict[str, Any], + rng: random.Random, + device: torch.device, +) -> list[TapExample]: + examples: list[TapExample] = [] + num_layers = int(plan["model"]["num_hidden_layers"]) + batch_size = max(1, int(args.extract_batch_size)) + if getattr(tokenizer, "pad_token_id", None) is None: + tokenizer.pad_token = tokenizer.eos_token or tokenizer.unk_token + for start in range(0, len(texts), batch_size): + batch_texts = texts[start : start + batch_size] + encoded = tokenizer( + batch_texts, + return_tensors="pt", + truncation=True, + max_length=int(args.max_length), + add_special_tokens=True, + padding=True, + ) + input_ids = encoded["input_ids"].to(device) + attention_mask = encoded["attention_mask"].to(device) + with torch.no_grad(): + output = model( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + use_cache=False, + ) + hidden_states = [state.detach().cpu().float() for state in output.hidden_states] + input_ids_cpu = input_ids.detach().cpu() + attention_mask_cpu = attention_mask.detach().cpu() + for row in range(input_ids_cpu.shape[0]): + examples.extend( + collect_examples_from_encoded_row( + args, + hidden_states, + input_ids_cpu, + attention_mask_cpu, + row, + plan, + rng, + num_layers, + ) + ) + print( + f"collected {len(examples)} examples from {min(start + batch_size, len(texts))}/{len(texts)} texts", + flush=True, + ) + if not examples: + raise RuntimeError("no real GLM examples collected") + return examples + + +def collect_examples_from_encoded_row( + args: argparse.Namespace, + hidden_states: list[torch.Tensor], + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + row: int, + plan: dict[str, Any], + rng: random.Random, + num_layers: int, +) -> list[TapExample]: + real_positions = attention_mask[row].nonzero(as_tuple=False).flatten().tolist() + token_ids = [int(input_ids[row, col]) for col in real_positions] + if len(token_ids) <= int(args.num_spec_layers) + 1: + return [] + max_pos = len(token_ids) - int(args.num_spec_layers) - 1 + positions = sorted(rng.sample(range(max_pos), min(int(args.positions_per_row), max_pos))) + examples: list[TapExample] = [] + for pos in positions: + indices = sample_tap_indices(plan, rng) + tensor_pos = int(real_positions[pos]) + hidden = torch.stack([hidden_states[index][row, tensor_pos] for index in indices], dim=0) + features = tap_features(indices, num_layers) + labels = [token_ids[pos + offset] for offset in range(1, int(args.num_spec_layers) + 1)] + examples.append( + TapExample( + hidden=hidden, + features=features, + labels=labels, + topology_key=topology_key(indices), + ) + ) + return examples + + +def sample_tap_indices(plan: dict[str, Any], rng: random.Random) -> list[int]: + layout = rng.choice(plan["layouts"]) + row = list(layout["shallow_hidden_layer_indices"][0]) + dropout = float(layout.get("tap_dropout", {}).get("probability", 0.0)) + required = set(int(index) for index in layout.get("tap_dropout", {}).get("required_indices", [])) + kept = [index for index in row if index in required or rng.random() >= dropout] + if not kept: + kept = [row[0], row[-1]] + return sorted(set(int(index) for index in kept)) + + +def topology_key(indices: list[int]) -> str: + return ",".join(str(index) for index in indices) + + +def tap_features(indices: list[int], num_layers: int) -> torch.Tensor: + rows = [] + for index in indices: + rows.append([float(index) / float(max(1, num_layers)), 1.0 if index == 0 else 0.0]) + return torch.tensor(rows, dtype=torch.float32) + + +def train_sidecar( + sidecar: GenericLayerTapSidecar, + optimizer: torch.optim.Optimizer, + examples: list[TapExample], + label_map: dict[int, int], + *, + batch_size: int, + epochs: int, + device: torch.device, +) -> float: + sidecar.train() + losses: list[float] = [] + rng = random.Random(17) + for _ in range(max(1, epochs)): + rng.shuffle(examples) + for start in range(0, len(examples), batch_size): + batch = examples[start : start + batch_size] + hidden, features, mask, labels = collate_batch(batch, label_map, device) + logits = sidecar(hidden, features, mask) + loss = batch_loss(logits, labels) + optimizer.zero_grad(set_to_none=True) + loss.backward() + optimizer.step() + losses.append(float(loss.detach().cpu())) + return sum(losses) / max(1, len(losses)) + + +def evaluate_sidecar( + sidecar: GenericLayerTapSidecar, + examples: list[TapExample], + label_map: dict[int, int], + *, + batch_size: int, + device: torch.device, + top_k_values: list[int], +) -> dict[str, Any]: + sidecar.eval() + accepted = 0 + proposal_slots = 0 + covered_labels = 0 + total_labels = 0 + proposal_time = 0.0 + top_k_stats = new_top_k_stats(top_k_values, len(examples[0].labels)) + by_topology = { + example.topology_key: {"examples": 0, "accepted": 0, "slots": 0} + for example in examples + } + eval_started = time.perf_counter() + with torch.no_grad(): + for start in range(0, len(examples), batch_size): + batch = examples[start : start + batch_size] + hidden, features, mask, labels = collate_batch(batch, label_map, device) + proposal_started = time.perf_counter() + logits = sidecar(hidden, features, mask) + if device.type == "mps": + torch.mps.synchronize() + elif device.type == "cuda": + torch.cuda.synchronize() + proposal_time += time.perf_counter() - proposal_started + update_top_k_stats(logits, labels, top_k_stats) + predictions = torch.stack([head.argmax(dim=-1) for head in logits], dim=1) + for row_idx in range(predictions.shape[0]): + topology = batch[row_idx].topology_key + by_topology[topology]["examples"] += 1 + for spec_idx in range(predictions.shape[1]): + label = int(labels[row_idx, spec_idx].detach().cpu()) + total_labels += 1 + if label < 0: + break + covered_labels += 1 + proposal_slots += 1 + by_topology[topology]["slots"] += 1 + if int(predictions[row_idx, spec_idx].detach().cpu()) != label: + break + accepted += 1 + by_topology[topology]["accepted"] += 1 + eval_wall = time.perf_counter() - eval_started + denominator = max(1, proposal_slots) + topology_rows = [] + for key, row in sorted(by_topology.items()): + slots = max(1, row["slots"]) + topology_rows.append( + { + "tap_indices": key, + "examples": row["examples"], + "accepted_draft_tokens": row["accepted"], + "proposal_slots": row["slots"], + "acceptance_rate": row["accepted"] / slots, + } + ) + return { + "examples": len(examples), + "accepted_draft_tokens": accepted, + "proposal_slots": proposal_slots, + "acceptance_rate": accepted / denominator, + "equivalent_accept_length": accepted / max(1, len(examples)), + "draft_vocab_label_coverage": covered_labels / max(1, total_labels), + "covered_labels": covered_labels, + "total_labels": total_labels, + "proposal_latency_ms": (proposal_time / max(1, len(examples))) * 1000.0, + "proposal_wall_seconds": proposal_time, + "eval_wall_seconds": eval_wall, + "top_k_diagnostics": summarize_top_k_stats(top_k_stats), + "by_topology": topology_rows, + } + + +def parse_top_k_values(raw: str) -> list[int]: + values = [] + for part in raw.split(","): + part = part.strip() + if not part: + continue + value = int(part) + if value <= 0: + raise ValueError("--eval-top-k values must be positive") + values.append(value) + return sorted(set(values)) or [1] + + +def new_top_k_stats(top_k_values: list[int], num_spec_layers: int) -> dict[str, Any]: + return { + "k_values": list(top_k_values), + "overall": {k: {"hits": 0, "labels": 0} for k in top_k_values}, + "by_spec_layer": [ + {k: {"hits": 0, "labels": 0} for k in top_k_values} + for _ in range(num_spec_layers) + ], + } + + +def update_top_k_stats( + logits: list[torch.Tensor], + labels: torch.Tensor, + stats: dict[str, Any], +) -> None: + for spec_idx, head_logits in enumerate(logits): + target = labels[:, spec_idx] + valid = target.ge(0) + if not valid.any(): + continue + valid_targets = target[valid].detach().cpu() + valid_logits = head_logits[valid].detach() + max_k = min(max(stats["k_values"]), valid_logits.shape[-1]) + top_indices = valid_logits.topk(k=max_k, dim=-1).indices.detach().cpu() + for k in stats["k_values"]: + clipped_k = min(k, top_indices.shape[-1]) + hits = int(top_indices[:, :clipped_k].eq(valid_targets.unsqueeze(1)).any(dim=1).sum()) + labels_count = int(valid_targets.numel()) + stats["overall"][k]["hits"] += hits + stats["overall"][k]["labels"] += labels_count + stats["by_spec_layer"][spec_idx][k]["hits"] += hits + stats["by_spec_layer"][spec_idx][k]["labels"] += labels_count + + +def summarize_top_k_stats(stats: dict[str, Any]) -> dict[str, Any]: + return { + "overall": top_k_rows(stats["overall"], stats["k_values"]), + "by_spec_layer": [ + { + "spec_layer": spec_idx, + "rows": top_k_rows(layer_stats, stats["k_values"]), + } + for spec_idx, layer_stats in enumerate(stats["by_spec_layer"]) + ], + } + + +def top_k_rows(rows: dict[int, dict[str, int]], k_values: list[int]) -> list[dict[str, Any]]: + result = [] + for k in k_values: + labels = rows[k]["labels"] + hits = rows[k]["hits"] + result.append( + { + "k": k, + "hits": hits, + "labels": labels, + "hit_rate": hits / max(1, labels), + } + ) + return result + + +def summarize_examples_by_topology(examples: list[TapExample]) -> list[dict[str, Any]]: + counts: dict[str, int] = {} + for example in examples: + counts[example.topology_key] = counts.get(example.topology_key, 0) + 1 + return [{"tap_indices": key, "examples": value} for key, value in sorted(counts.items())] + + +def collate_batch( + batch: list[TapExample], + label_map: dict[int, int], + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + max_taps = max(example.hidden.shape[0] for example in batch) + hidden_size = batch[0].hidden.shape[1] + num_spec_layers = len(batch[0].labels) + hidden = torch.zeros(len(batch), max_taps, hidden_size, dtype=torch.float32) + features = torch.zeros(len(batch), max_taps, 2, dtype=torch.float32) + mask = torch.zeros(len(batch), max_taps, dtype=torch.bool) + labels = torch.full((len(batch), num_spec_layers), -100, dtype=torch.long) + for row, example in enumerate(batch): + taps = example.hidden.shape[0] + hidden[row, :taps] = example.hidden + features[row, :taps] = example.features + mask[row, :taps] = True + for spec_idx, token_id in enumerate(example.labels): + labels[row, spec_idx] = label_map.get(int(token_id), -100) + return hidden.to(device), features.to(device), mask.to(device), labels.to(device) + + +def batch_loss(logits: list[torch.Tensor], labels: torch.Tensor) -> torch.Tensor: + losses = [] + for spec_idx, head_logits in enumerate(logits): + target = labels[:, spec_idx] + valid = target.ge(0) + if valid.any(): + losses.append(nn.functional.cross_entropy(head_logits[valid], target[valid])) + if not losses: + return logits[0].sum() * 0.0 + return torch.stack(losses).mean() + + +def save_serving_safetensors( + sidecar: GenericLayerTapSidecar, + output_path: Path, + dtype_name: str, +) -> None: + from safetensors.torch import save_file + + dtype = torch_dtype(dtype_name) + tensors = { + name: tensor.detach().cpu().to(dtype) + for name, tensor in sorted(sidecar.state_dict().items()) + } + save_file( + tensors, + output_path, + metadata={ + "format": SERVING_FORMAT, + "source_format": SOURCE_FORMAT, + "head_kind": HEAD_KIND, + "dtype": safetensors_dtype_label(dtype), + "tensor_count": str(len(tensors)), + }, + ) + + +def write_manifest( + *, + args: argparse.Namespace, + manifest_path: Path, + checkpoint_path: Path, + serving_path: Path, + model_meta: dict[str, Any], + draft_token_ids: list[int], + plan: dict[str, Any], +) -> None: + checkpoint_rel = checkpoint_path.name + serving_rel = serving_path.name + max_stages = int(plan["policy"]["max_stages"]) + manifest = { + "schema": "skippy-spd-head/v1", + "checkpoint": { + "path": checkpoint_rel, + "sha256": file_sha256(checkpoint_path), + "bytes": checkpoint_path.stat().st_size, + }, + "serving_checkpoint": { + "path": serving_rel, + "sha256": file_sha256(serving_path), + "bytes": serving_path.stat().st_size, + "format": SERVING_FORMAT, + "tensor_count": len(dict(torch.load(checkpoint_path, map_location="cpu")["state_dict"])), + "dtype": safetensors_dtype_label(torch_dtype(args.export_dtype)), + }, + "source": { + "format": SOURCE_FORMAT, + "reference_repo": None, + "base_model_path": args.manifest_base_model_path.strip() or args.model_name, + "model_type": model_meta.get("model_type"), + "checkpoint_version": 1, + }, + "topology": { + "hidden_size": int(model_meta["hidden_size"]), + "vocab_size": int(model_meta["vocab_size"]), + "draft_vocab_size": len(draft_token_ids), + "head_kind": HEAD_KIND, + "encoder_kind": args.encoder, + "attention_heads": int(args.attention_heads) if args.encoder == "attention" else None, + "mlp_ratio": float(args.mlp_ratio) if args.encoder == "mean_mlp" else None, + "num_stages": max_stages, + "stage_layer_boundaries": None, + "num_spec_layers": int(args.num_spec_layers), + "max_taps": max_taps(plan), + "tap_feature_size": 2, + "trained_with_use_deepest": False, + "shallow_hidden_layer_indices": representative_tap_rows(plan), + "spec_init_from_base_layers": None, + "draft_token_ids": draft_token_ids, + }, + } + manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def representative_tap_rows(plan: dict[str, Any]) -> list[list[int]]: + rows: list[list[int]] = [] + seen: set[tuple[int, ...]] = set() + for layout in plan["layouts"]: + row = tuple(int(index) for index in layout["shallow_hidden_layer_indices"][0]) + if row in seen: + continue + seen.add(row) + rows.append(list(row)) + if len(rows) >= 16: + break + return rows + + +def max_taps(plan: dict[str, Any]) -> int: + return max( + len(layout["shallow_hidden_layer_indices"][0]) + for layout in plan["layouts"] + ) + + +def load_training_rows(dataset_name: str, split: str, limit: int) -> list[dict[str, Any]]: + from datasets import load_dataset + + ds = load_dataset(dataset_name, split=f"{split}[:{max(1, limit)}]") + rows = [] + for row in ds: + messages = row.get("messages") or row.get("conversations") + if messages: + rows.append({"messages": normalize_messages(messages)}) + return rows + + +def normalize_messages(messages: list[dict[str, Any]]) -> list[dict[str, str]]: + normalized = [] + for message in messages: + role = message.get("role") or message.get("from") + content = message.get("content") or message.get("value") + if role == "human": + role = "user" + elif role == "gpt": + role = "assistant" + if role is not None and content is not None: + normalized.append({"role": str(role), "content": str(content)}) + return normalized + + +def row_text_for_vocab(tokenizer: Any, messages: list[dict[str, Any]]) -> str: + try: + rendered = tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=False, + enable_thinking=False, + ) + if isinstance(rendered, str): + return rendered + except Exception: + pass + return "\n".join(message.get("content", "") for message in messages) + + +def build_draft_vocab(tokenizer: Any, texts: list[str], draft_vocab_size: int) -> list[int]: + from collections import Counter + + counts: Counter[int] = Counter() + for text in texts: + counts.update(int(token_id) for token_id in tokenizer.encode(text, add_special_tokens=False)) + for token_id in ( + getattr(tokenizer, "eos_token_id", None), + getattr(tokenizer, "pad_token_id", None), + getattr(tokenizer, "bos_token_id", None), + ): + if token_id is not None: + counts[int(token_id)] += 1 + if not counts: + raise RuntimeError("could not build draft vocab") + ids = [ + token_id + for token_id, _ in sorted(counts.items(), key=lambda item: (-item[1], item[0]))[ + : max(1, draft_vocab_size) + ] + ] + return sorted(set(ids)) + + +def resolve_device(value: str) -> torch.device: + if value == "cuda" or (value == "auto" and torch.cuda.is_available()): + return torch.device("cuda") + if value == "mps" or (value == "auto" and torch.backends.mps.is_available()): + return torch.device("mps") + return torch.device("cpu") + + +def torch_dtype(value: str) -> torch.dtype: + return { + "float32": torch.float32, + "float16": torch.float16, + "bfloat16": torch.bfloat16, + }[value] + + +def safetensors_dtype_label(dtype: torch.dtype) -> str: + return { + torch.float32: "F32", + torch.float16: "F16", + torch.bfloat16: "BF16", + }[dtype] + + +def file_sha256(path: Path) -> str: + hasher = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + hasher.update(chunk) + return hasher.hexdigest() + + +def timestamp() -> str: + return time.strftime("%Y%m%d-%H%M%S", time.gmtime()) + + +if __name__ == "__main__": + main() diff --git a/evals/spd/glm47_frontload.py b/evals/spd/glm47_frontload.py new file mode 100755 index 0000000000..d3819ba735 --- /dev/null +++ b/evals/spd/glm47_frontload.py @@ -0,0 +1,357 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.10" +# dependencies = [] +# /// +"""Frontload GLM 4.7 SPD integration checks. + +This utility is intentionally lightweight: it inspects a local GLM checkpoint, +derives SPD topology metadata for non-uniform Skippy stage boundaries, and can +write tiny manifest-compatible smoke artifacts. The smoke artifacts validate the +Skippy SPD manifest and serving-checkpoint shape contract; they are not trained +weights. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import struct +import subprocess +import sys +import time +from pathlib import Path +from typing import Any + + +REFERENCE_REPO = "https://github.com/yuyijiong/speculative_pipeline_decoding.git" +DEFAULT_GLM47_FLASH_SNAPSHOT = ( + Path.home() + / ".cache/huggingface/hub/models--zai-org--GLM-4.7-Flash/" + / "snapshots/7dd20894a642a0aa287e9827cb1a1f7f91386b67" +) +SMOKE_CHECKPOINT_FORMAT = "torch-speculation-head-v10" +SERVING_FORMAT = "safetensors-spd-head-v1" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Inspect and frontload GLM 4.7 SPD metadata") + parser.add_argument( + "--model-path", + default=str(DEFAULT_GLM47_FLASH_SNAPSHOT), + help="Local GLM checkpoint directory containing config.json.", + ) + parser.add_argument("--work-dir", default="/tmp/skippy-spd-glm47-frontload") + parser.add_argument("--reference-repo", default=REFERENCE_REPO) + parser.add_argument( + "--patch-reference", + action="store_true", + help="Clone and patch the SPD reference repo's model-type allowlist for GLM.", + ) + parser.add_argument( + "--write-smoke-artifacts", + action="store_true", + help="Write tiny manifest-compatible SPD smoke artifacts.", + ) + parser.add_argument("--num-stages", type=int, default=3) + parser.add_argument( + "--stage-layer-boundaries", + default="", + help="Comma-separated target layer end indices, e.g. 15,31,47.", + ) + parser.add_argument("--num-spec-layers", type=int, default=1) + parser.add_argument("--draft-vocab-size", type=int, default=8) + parser.add_argument("--out-name", default="glm47-spd-frontload") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + model_path = Path(args.model_path).expanduser().resolve() + work_dir = Path(args.work_dir).expanduser().resolve() + out_dir = work_dir / args.out_name + out_dir.mkdir(parents=True, exist_ok=True) + + inspection = inspect_checkpoint(model_path) + boundaries = resolve_stage_boundaries( + args.stage_layer_boundaries, + num_stages=args.num_stages, + num_layers=inspection["num_hidden_layers"], + ) + hidden_indices = derive_hidden_tap_indices(boundaries) + topology = { + "hidden_size": inspection["hidden_size"], + "vocab_size": inspection["vocab_size"], + "draft_vocab_size": args.draft_vocab_size, + "num_stages": len(boundaries), + "stage_layer_boundaries": boundaries, + "num_spec_layers": args.num_spec_layers, + "trained_with_use_deepest": False, + "shallow_hidden_layer_indices": hidden_indices, + "spec_init_from_base_layers": None, + "draft_token_ids": list(range(args.draft_vocab_size)), + } + + report = { + "model_path": str(model_path), + "inspected_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "checkpoint": inspection, + "topology": topology, + "notes": [ + "Smoke artifacts are shape-contract fixtures, not trained SPD weights.", + "Stage boundaries are target layer end indices.", + "Hidden-state indices follow Hugging Face convention: 0 is embeddings, k is output after layer k-1.", + ], + } + report_path = out_dir / "glm47-spd-frontload.json" + write_json(report_path, report) + + if args.patch_reference: + reference_dir = work_dir / "speculative_pipeline_decoding" + clone_reference(args.reference_repo, reference_dir) + patch_reference_for_glm(reference_dir) + + if args.write_smoke_artifacts: + write_smoke_artifacts( + out_dir=out_dir, + model_path=model_path, + inspection=inspection, + topology=topology, + reference_repo=args.reference_repo, + ) + + print(json.dumps(report, indent=2, sort_keys=True)) + + +def inspect_checkpoint(model_path: Path) -> dict[str, Any]: + config_path = model_path / "config.json" + if not config_path.is_file(): + raise FileNotFoundError(f"GLM config.json not found: {config_path}") + config = json.loads(config_path.read_text(encoding="utf-8")) + weight_map = read_weight_map(model_path) + auxiliary_tensors = sorted( + name + for name in weight_map + if any(part in name.lower() for part in ("eh_proj", "enorm", "hnorm", "nextn", "mtp")) + ) + return { + "architectures": config.get("architectures", []), + "model_type": config.get("model_type"), + "hidden_size": positive_int(config, "hidden_size"), + "vocab_size": positive_int(config, "vocab_size"), + "num_hidden_layers": positive_int(config, "num_hidden_layers"), + "num_nextn_predict_layers": int(config.get("num_nextn_predict_layers") or 0), + "num_attention_heads": config.get("num_attention_heads"), + "num_key_value_heads": config.get("num_key_value_heads"), + "rope_theta": config.get("rope_theta"), + "tokenizer": inspect_tokenizer(model_path), + "weight_shards": len({shard for shard in weight_map.values()}), + "tensor_count": len(weight_map), + "auxiliary_tensors": auxiliary_tensors, + } + + +def read_weight_map(model_path: Path) -> dict[str, str]: + index_path = model_path / "model.safetensors.index.json" + if not index_path.is_file(): + return {} + index = json.loads(index_path.read_text(encoding="utf-8")) + weight_map = index.get("weight_map") or {} + if not isinstance(weight_map, dict): + raise RuntimeError(f"invalid weight_map in {index_path}") + return {str(k): str(v) for k, v in weight_map.items()} + + +def inspect_tokenizer(model_path: Path) -> dict[str, Any]: + tokenizer_config = model_path / "tokenizer_config.json" + if not tokenizer_config.is_file(): + return {} + config = json.loads(tokenizer_config.read_text(encoding="utf-8")) + return { + "tokenizer_class": config.get("tokenizer_class"), + "model_max_length": config.get("model_max_length"), + "has_chat_template": bool(config.get("chat_template") or (model_path / "chat_template.jinja").is_file()), + } + + +def positive_int(config: dict[str, Any], key: str) -> int: + value = int(config.get(key) or 0) + if value <= 0: + raise RuntimeError(f"config field {key!r} must be a positive integer") + return value + + +def resolve_stage_boundaries(value: str, *, num_stages: int, num_layers: int) -> list[int]: + if value.strip(): + boundaries = [int(part.strip()) for part in value.split(",") if part.strip()] + elif num_layers == 47 and num_stages == 3: + boundaries = [15, 31, 47] + else: + boundaries = [ + round(num_layers * (stage + 1) / num_stages) for stage in range(num_stages) + ] + if not boundaries: + raise RuntimeError("stage_layer_boundaries must not be empty") + if boundaries[-1] != num_layers: + raise RuntimeError( + f"last stage boundary must equal num_hidden_layers={num_layers}, got {boundaries[-1]}" + ) + if any(left >= right for left, right in zip(boundaries, boundaries[1:])): + raise RuntimeError(f"stage boundaries must be strictly increasing: {boundaries}") + return boundaries + + +def derive_hidden_tap_indices(boundaries: list[int]) -> list[list[int]]: + rows: list[list[int]] = [] + for depth in range(len(boundaries), 0, -1): + rows.append([0, *boundaries[:depth]]) + return rows + + +def clone_reference(repo_url: str, dest: Path) -> None: + if dest.exists(): + print(f"reference repo already exists: {dest}", file=sys.stderr) + return + run(["git", "clone", "--depth", "1", repo_url, str(dest)]) + + +def patch_reference_for_glm(reference_dir: Path) -> None: + pipeline_model = reference_dir / "pipeline_model.py" + replace_once( + pipeline_model, + 'supported = {"qwen3", "qwen3_moe", "qwen3_5", "qwen3_5_text", "qwen3_5_moe", "qwen3_5_moe_text", "llama"}', + 'supported = {"qwen3", "qwen3_moe", "qwen3_5", "qwen3_5_text", "qwen3_5_moe", "qwen3_5_moe_text", "llama", "glm4_moe_lite"}', + ) + + +def replace_once(path: Path, old: str, new: str) -> None: + text = path.read_text(encoding="utf-8") + if old not in text: + if new in text: + return + raise RuntimeError(f"expected text not found in {path}: {old[:80]!r}") + path.write_text(text.replace(old, new, 1), encoding="utf-8") + + +def write_smoke_artifacts( + *, + out_dir: Path, + model_path: Path, + inspection: dict[str, Any], + topology: dict[str, Any], + reference_repo: str, +) -> None: + checkpoint_path = out_dir / "speculation_head_final.pt" + checkpoint_payload = { + "note": "GLM SPD frontload placeholder; not a torch training checkpoint.", + "config": checkpoint_config(model_path, inspection, topology), + } + checkpoint_path.write_text(json.dumps(checkpoint_payload, indent=2, sort_keys=True) + "\n") + serving_path = out_dir / "spd-head.safetensors" + write_smoke_safetensors(serving_path, topology) + manifest = { + "schema": "skippy-spd-head/v1", + "checkpoint": { + "path": checkpoint_path.name, + "sha256": file_sha256(checkpoint_path), + "bytes": checkpoint_path.stat().st_size, + }, + "serving_checkpoint": { + "path": serving_path.name, + "sha256": file_sha256(serving_path), + "bytes": serving_path.stat().st_size, + "format": SERVING_FORMAT, + "tensor_count": smoke_tensor_count(topology), + "dtype": "F32", + }, + "source": { + "format": SMOKE_CHECKPOINT_FORMAT, + "reference_repo": reference_repo, + "base_model_path": str(model_path), + "model_type": inspection["model_type"], + "checkpoint_version": 10, + }, + "topology": topology, + } + write_json(out_dir / "skippy-spd-head.json", manifest) + + +def checkpoint_config( + model_path: Path, + inspection: dict[str, Any], + topology: dict[str, Any], +) -> dict[str, Any]: + return { + "version": 10, + "base_model_path": str(model_path), + "model_type": inspection["model_type"], + **topology, + } + + +def write_smoke_safetensors(path: Path, topology: dict[str, Any]) -> None: + hidden = int(topology["hidden_size"]) + tensors: list[tuple[str, str, list[int]]] = [] + for stage, indices in enumerate(topology["shallow_hidden_layer_indices"]): + tensors.append((f"stage_projs.{stage}.weight", "F32", [hidden, hidden * len(indices)])) + tensors.append(("g0_proj.weight", "F32", [hidden, hidden])) + tensors.append(("lm_head.weight", "F32", [int(topology["draft_vocab_size"]), hidden])) + for layer in range(int(topology["num_spec_layers"])): + tensors.append((f"spec_layers.{layer}.input_layernorm.weight", "F32", [hidden])) + tensors.append((f"spec_layers.{layer}.post_attention_layernorm.weight", "F32", [hidden])) + + header_entries: dict[str, Any] = { + "__metadata__": { + "format": SERVING_FORMAT, + "purpose": "glm47-spd-frontload-smoke", + } + } + data_len = 0 + for name, dtype, shape in tensors: + byte_len = tensor_byte_len(dtype, shape) + header_entries[name] = { + "dtype": dtype, + "shape": shape, + "data_offsets": [data_len, data_len + byte_len], + } + data_len += byte_len + header = json.dumps(header_entries, sort_keys=True, separators=(",", ":")).encode() + with path.open("wb") as handle: + handle.write(struct.pack(" int: + return len(topology["shallow_hidden_layer_indices"]) + 2 + 2 * int(topology["num_spec_layers"]) + + +def tensor_byte_len(dtype: str, shape: list[int]) -> int: + sizes = {"F32": 4} + elements = 1 + for dimension in shape: + elements *= int(dimension) + return elements * sizes[dtype] + + +def write_json(path: Path, value: dict[str, Any]) -> None: + path.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(f"wrote {path}", file=sys.stderr) + + +def file_sha256(path: Path) -> str: + hasher = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + hasher.update(chunk) + return hasher.hexdigest() + + +def run(cmd: list[str]) -> None: + print("+", " ".join(cmd), file=sys.stderr) + subprocess.run(cmd, check=True) + + +if __name__ == "__main__": + main() diff --git a/evals/spd/hf_train_eval_qwen06.py b/evals/spd/hf_train_eval_qwen06.py new file mode 100755 index 0000000000..78031333a3 --- /dev/null +++ b/evals/spd/hf_train_eval_qwen06.py @@ -0,0 +1,1136 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "accelerate>=1.0.0", +# "datasets>=3.0.0", +# "huggingface_hub>=0.30.0", +# "numpy", +# "pyarrow", +# "setproctitle", +# "torch>=2.8.0", +# "tqdm", +# "transformers>=5.6.0", +# ] +# /// +"""Train and evaluate an SPD speculation head on Hugging Face Jobs or locally. + +This is intentionally a proof runner, not serving code. It produces a real +`speculation_head_final.pt` from the reference SPD implementation, evaluates it, +and uploads the checkpoint + eval summaries to a private HF model repo by +default. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import shutil +import subprocess +import sys +import time +from pathlib import Path +from typing import Any + +from topology_plan import write_topology_plan + + +REFERENCE_REPO = "https://github.com/yuyijiong/speculative_pipeline_decoding.git" +DEFAULT_MODEL = "Qwen/Qwen3-0.6B" +DEFAULT_DATASET = "HuggingFaceH4/ultrachat_200k" +DEFAULT_DATASET_SPLIT = "train_sft" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Run a real SPD head proof or smoke job") + parser.add_argument("--work-dir", default="/tmp/skippy-spd-qwen06-proof") + parser.add_argument("--reference-repo", default=REFERENCE_REPO) + parser.add_argument("--model-name", default=DEFAULT_MODEL) + parser.add_argument("--dataset", default=DEFAULT_DATASET) + parser.add_argument("--dataset-split", default=DEFAULT_DATASET_SPLIT) + parser.add_argument("--train-rows", type=int, default=1024) + parser.add_argument("--eval-rows-per-set", type=int, default=8) + parser.add_argument("--num-stages", type=int, default=2) + parser.add_argument( + "--topology-policy", + choices=("fixed", "generic-plan"), + default="fixed", + help=( + "Topology handling. 'fixed' preserves the reference trainer contract. " + "'generic-plan' writes randomized contiguous-layer tap plans and exits; " + "it is the scaffold for topology-independent sidecar training." + ), + ) + parser.add_argument( + "--topology-plan-out", + default="", + help=( + "Output JSON path for --topology-policy generic-plan. Defaults under " + "the current artifact directory." + ), + ) + parser.add_argument( + "--topology-plan-samples", + type=int, + default=32, + help="Number of randomized contiguous split layouts to write in generic-plan mode.", + ) + parser.add_argument( + "--topology-min-stages", + type=int, + default=2, + help="Minimum stage count to sample in generic-plan mode.", + ) + parser.add_argument( + "--topology-max-stages", + type=int, + default=6, + help="Maximum stage count to sample in generic-plan mode.", + ) + parser.add_argument( + "--topology-seed", + type=int, + default=47, + help="Random seed for generic-plan topology sampling.", + ) + parser.add_argument( + "--topology-tap-dropout", + type=float, + default=0.25, + help="Recorded tap-dropout probability for future generic topology training.", + ) + parser.add_argument( + "--topology-num-hidden-layers", + type=int, + default=0, + help="Override target layer count for generic-plan mode when config loading is unavailable.", + ) + parser.add_argument( + "--stage-layer-boundaries", + default="", + help=( + "Comma-separated target layer end indices for non-uniform topologies, " + "for example 15,31,47 for GLM 4.7 Flash." + ), + ) + parser.add_argument( + "--shallow-hidden-layer-indices", + default="", + help=( + "Explicit semicolon-separated HF hidden-state tap rows for [g_n..g_1]. " + "Overrides --stage-layer-boundaries when set." + ), + ) + parser.add_argument("--num-spec-layers", type=int, default=1) + parser.add_argument("--epochs", type=int, default=1) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--gradient-accumulation-steps", type=int, default=8) + parser.add_argument( + "--learning-rate", + type=float, + default=1e-5, + help="Learning rate for the speculation head trainer.", + ) + parser.add_argument( + "--log-interval", + type=int, + default=20, + help="Training log interval passed to the reference trainer.", + ) + parser.add_argument( + "--warmup-steps", + type=int, + default=100, + help="Learning-rate warmup steps passed to the reference trainer.", + ) + parser.add_argument( + "--save-steps", + type=int, + default=5000, + help="Checkpoint save interval passed to the reference trainer.", + ) + parser.add_argument("--max-length", type=int, default=512) + parser.add_argument("--max-new-tokens", type=int, default=64) + parser.add_argument("--draft-top-k", type=int, default=1) + parser.add_argument("--attn-implementation", default="sdpa") + parser.add_argument( + "--device", + choices=("auto", "cuda", "mps", "cpu"), + default="auto", + help="Device for local/reference execution. HF GPU jobs can leave this as auto.", + ) + parser.add_argument( + "--draft-vocab-json", + default="draft_vocab/ultrachat_qwen3_0.6b_top_32k.json", + help=( + "Draft vocab JSON path. Relative paths are resolved inside the reference repo; " + "absolute paths are passed through. Empty disables reduced draft vocab." + ), + ) + parser.add_argument( + "--build-draft-vocab-size", + type=int, + default=0, + help="Build a tokenizer-specific draft vocab from the loaded train rows before training.", + ) + parser.add_argument( + "--draft-vocab-out", + default="", + help="Output JSON for --build-draft-vocab-size. Defaults under the artifact data dir.", + ) + parser.add_argument( + "--upload-repo", + default="auto", + help="HF model repo for artifacts. Use 'auto' for /skippy-spd-qwen06-proof.", + ) + parser.add_argument("--public", action="store_true", help="Create upload repo as public") + parser.add_argument( + "--spec-head-path", + default="", + help="Existing speculation_head checkpoint to evaluate instead of training.", + ) + parser.add_argument( + "--spec-head-repo", + default="", + help="HF model repo containing an existing speculation_head checkpoint.", + ) + parser.add_argument( + "--spec-head-file", + default="", + help="Filename inside --spec-head-repo for an existing speculation_head checkpoint.", + ) + parser.add_argument( + "--manifest-base-model-path", + default="", + help="Override the base_model_path written to the Skippy SPD manifest.", + ) + parser.add_argument("--skip-train", action="store_true") + parser.add_argument("--skip-eval", action="store_true") + return parser.parse_args() + + +def run(cmd: list[str], *, cwd: Path | None = None, env: dict[str, str] | None = None) -> None: + print("+", " ".join(cmd), flush=True) + subprocess.run(cmd, cwd=str(cwd) if cwd else None, env=env, check=True) + + +def clone_reference(repo_url: str, dest: Path) -> None: + if dest.exists(): + print(f"reference repo already exists: {dest}", flush=True) + return + run(["git", "clone", "--depth", "1", repo_url, str(dest)]) + + +def patch_reference_for_proof(reference_dir: Path) -> None: + write_qwen3_nonthink_template(reference_dir / "qwen3-nonthink-template") + write_glm4_moe_lite_template(reference_dir / "glm4-moe-lite-template") + replace_once( + reference_dir / "train.py", + ' report_to="wandb",\n', + " report_to=[],\n", + ) + patch_reference_for_transformers(reference_dir) + patch_reference_for_glm_training_smoke(reference_dir) + + +def write_qwen3_nonthink_template(path: Path) -> None: + path.write_text( + """{%- for message in messages %} +{%- if message['role'] == 'system' %} +<|im_start|>system +{{ message['content'] }}<|im_end|> +{%- elif message['role'] == 'user' %} +<|im_start|>user +{{ message['content'] }}<|im_end|> +{%- elif message['role'] == 'assistant' %} +{% generation %}<|im_start|>assistant +{{ message['content'] }}<|im_end|>{% endgeneration %} +{%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} +<|im_start|>assistant +{%- endif %} +""", + encoding="utf-8", + ) + + +def write_glm4_moe_lite_template(path: Path) -> None: + path.write_text( + """[gMASK] +{%- macro visible_text(content) -%} + {%- if content is string -%} + {{- content }} + {%- elif content is iterable and content is not mapping -%} + {%- for item in content -%} + {%- if item is mapping and item.type == 'text' -%} + {{- item.text }} + {%- elif item is string -%} + {{- item }} + {%- endif -%} + {%- endfor -%} + {%- else -%} + {{- content }} + {%- endif -%} +{%- endmacro -%} +{% for m in messages %} +{%- if m.role == 'system' -%} +<|system|>{{ visible_text(m.content) }} +{%- elif m.role == 'user' -%} +<|user|>{{ visible_text(m.content) }} +{%- elif m.role == 'assistant' -%} +{% generation %}<|assistant|>{{ visible_text(m.content) }}{% endgeneration %} +{%- endif -%} +{%- endfor -%} +{%- if add_generation_prompt -%} +<|assistant|>{{- '' if (enable_thinking is defined and not enable_thinking) else '' -}} +{%- endif -%} +""", + encoding="utf-8", + ) + + +def patch_reference_for_transformers(reference_dir: Path) -> None: + patch_reference_linear_cache_import(reference_dir / "pipeline_linear_cache.py") + replace_once( + reference_dir / "pipeline_model.py", + ' "cache_position": cache_position,\n', + "", + ) + + +def patch_reference_for_glm_training_smoke(reference_dir: Path) -> None: + patch_pipeline_model_for_glm(reference_dir / "pipeline_model.py") + patch_train_for_glm_template(reference_dir / "train.py") + patch_train_for_label_filtering(reference_dir / "train.py") + patch_train_for_training_controls(reference_dir / "train.py") + + +def patch_train_for_glm_template(path: Path) -> None: + replace_once( + path, + ''' if model_type is not None and not _model_type_looks_like_qwen(model_type): + log.info("Skip Qwen file chat template (model_type=%s).", model_type) + return + if template_dir is None: + template_dir = "." +''', + ''' if template_dir is None: + template_dir = "." + if model_type is not None and str(model_type).lower() == "glm4_moe_lite": + path = os.path.join(template_dir, "glm4-moe-lite-template") + if not os.path.isfile(path): + log.warning("GLM chat template not found: %s (use tokenizer default).", path) + return + with open(path, "r", encoding="utf-8") as f: + tokenizer.chat_template = f.read() + log.info("Loaded GLM chat template from %s (model_type=%s).", path, model_type) + return + if model_type is not None and not _model_type_looks_like_qwen(model_type): + log.info("Skip Qwen file chat template (model_type=%s).", model_type) + return +''', + ) + + +def patch_train_for_label_filtering(path: Path) -> None: + replace_once( + path, + 'ENCODE_PIPELINE_CACHE_VERSION = "spd-encode-1"\n', + 'ENCODE_PIPELINE_CACHE_VERSION = "spd-encode-2-next-labels"\n', + ) + replace_once( + path, + '''def _encoded_example_length_ok( + ex: Dict[str, Any], + *, + min_length: int, + max_length: int, + max_length_overflow: str, +) -> bool: + ids = ex.get("input_ids") + if ids is None: + return False + n = len(ids) + if n < min_length: + return False + if max_length_overflow == "discard" and n > max_length: + return False + return True +''', + '''def _encoded_example_length_ok( + ex: Dict[str, Any], + *, + min_length: int, + max_length: int, + max_length_overflow: str, +) -> bool: + ids = ex.get("input_ids") + if ids is None: + return False + n = len(ids) + if n < min_length: + return False + if max_length_overflow == "discard" and n > max_length: + return False + labels = ex.get("labels") + if labels is None: + return False + labels = labels[:n] + if not any(int(label) != -100 for label in labels): + return False + if n <= 1 or not any(int(label) != -100 for label in labels[1:n]): + return False + return True +''', + ) + + +def patch_train_for_training_controls(path: Path) -> None: + replace_once( + path, + ' p.add_argument("--log_interval", type=int, default=20)\n', + ''' p.add_argument("--log_interval", type=int, default=20) + p.add_argument("--warmup_steps", type=int, default=100) + p.add_argument("--save_steps", type=int, default=5000) +''', + ) + replace_once( + path, + " warmup_steps=100,\n", + " warmup_steps=args.warmup_steps,\n", + ) + replace_once( + path, + " save_steps=5000,\n", + " save_steps=args.save_steps,\n", + ) + + +def patch_pipeline_model_for_glm(path: Path) -> None: + replace_once( + path, + 'supported = {"qwen3", "qwen3_moe", "qwen3_5", "qwen3_5_text", "qwen3_5_moe", "qwen3_5_moe_text", "llama"}', + 'supported = {"qwen3", "qwen3_moe", "qwen3_5", "qwen3_5_text", "qwen3_5_moe", "qwen3_5_moe_text", "llama", "glm4_moe_lite"}', + ) + replace_once( + path, + ''' if self.num_layers % self.num_stages != 0: + raise ValueError( + f"num_layers ({self.num_layers}) must be divisible by num_stages ({self.num_stages})" + ) + self.layers_per_stage = self.num_layers // self.num_stages +''', + ''' if self.num_layers % self.num_stages != 0: + if shallow_hidden_layer_indices is None: + raise ValueError( + f"num_layers ({self.num_layers}) must be divisible by num_stages ({self.num_stages}) " + "unless shallow_hidden_layer_indices supplies an explicit non-uniform topology" + ) + self.layers_per_stage = max(1, (self.num_layers + self.num_stages - 1) // self.num_stages) + else: + self.layers_per_stage = self.num_layers // self.num_stages +''', + ) + replace_once( + path, + " self.shallow_hidden_layer_indices = self._normalize_stage_feature_indices(shallow_hidden_layer_indices)\n", + """ self.shallow_hidden_layer_indices = self._normalize_stage_feature_indices(shallow_hidden_layer_indices) + self.stage_layer_ranges = self._derive_stage_layer_ranges() +""", + ) + replace_once( + path, + ''' def _snap_indices_needed(self) -> Set[int]: + want: Set[int] = set() + for row in self.shallow_hidden_layer_indices: + for idx in row: + want.add(int(idx)) + return want +''', + ''' def _snap_indices_needed(self) -> Set[int]: + want: Set[int] = set() + for row in self.shallow_hidden_layer_indices: + for idx in row: + want.add(int(idx)) + return want + + def _derive_stage_layer_ranges(self) -> List[Tuple[int, int]]: + rows = self.shallow_hidden_layer_indices + deepest_row = tuple(sorted({int(x) for x in rows[0] if int(x) > 0})) if rows else () + if ( + len(deepest_row) >= self.num_stages + and deepest_row[-1] == self.num_layers + and all(a < b for a, b in zip(deepest_row, deepest_row[1:])) + ): + ranges: List[Tuple[int, int]] = [] + start = 0 + for end in deepest_row[: self.num_stages]: + ranges.append((start, end)) + start = end + return ranges + + ranges = [] + for stage_idx in range(self.num_stages): + start = min(stage_idx * self.layers_per_stage, self.num_layers) + end = min((stage_idx + 1) * self.layers_per_stage, self.num_layers) + ranges.append((start, end)) + return ranges +''', + ) + replace_once( + path, + ''' start_layer = stage_idx * lps + end_layer = (stage_idx + 1) * lps +''', + ''' start_layer, end_layer = self.stage_layer_ranges[stage_idx] +''', + ) + patch_pipeline_model_for_draft_vocab_training(path) + + +def patch_pipeline_model_for_draft_vocab_training(path: Path) -> None: + replace_once( + path, + ''' teacher_argmax_full = teacher_logits.argmax(dim=-1) + if self._use_draft_vocab: + teacher_argmax_draft = self._token_id_to_draft_idx.to(teacher_argmax_full.device)[teacher_argmax_full] + teacher_in_draft = teacher_argmax_draft >= 0 + valid_mask_2d = valid_mask_2d & teacher_in_draft + target_for_acc_2d = teacher_argmax_draft.to(spec_logits.device) + else: + target_for_acc_2d = teacher_argmax_full.to(spec_logits.device) +''', + ''' if self._use_draft_vocab: + # The KL target has already been sliced to the draft vocabulary above. + # Filtering by the full-vocab argmax can drop every assistant position + # when the reduced draft vocab is small, yielding an exact zero loss. + target_for_acc_2d = teacher_target.argmax(dim=-1).to(spec_logits.device) + else: + teacher_argmax_full = teacher_logits.argmax(dim=-1) + target_for_acc_2d = teacher_argmax_full.to(spec_logits.device) +''', + ) + replace_once( + path, + ''' start_layer = stage_idx * lps + end_layer = (stage_idx + 1) * lps +''', + ''' start_layer, end_layer = self.stage_layer_ranges[stage_idx] +''', + ) + + +def patch_reference_linear_cache_import(path: Path) -> None: + replace_once( + path, + '''from transformers.cache_utils import ( + LinearAttentionAndFullAttentionLayer, + LinearAttentionCacheLayerMixin, + LinearAttentionLayer, +) +''', + '''from transformers.cache_utils import LinearAttentionCacheLayerMixin, LinearAttentionLayer + +try: + from transformers.cache_utils import LinearAttentionAndFullAttentionLayer +except ImportError: + class LinearAttentionAndFullAttentionLayer(LinearAttentionLayer): + """Compatibility placeholder for Transformers 4.x without hybrid linear caches.""" + + pass + +''', + ) + + +def patch_reference_for_device(reference_dir: Path, device: str) -> None: + if device == "auto": + return + patch_train_for_device(reference_dir / "train.py") + patch_eval_for_device(reference_dir / "eval.py") + patch_pipeline_inference_for_device(reference_dir / "pipeline_inference.py") + + +def replace_once(path: Path, old: str, new: str) -> None: + text = path.read_text(encoding="utf-8") + if old not in text: + if new in text: + return + raise RuntimeError(f"expected text not found in {path}: {old[:80]!r}") + path.write_text(text.replace(old, new, 1), encoding="utf-8") + + +def patch_train_for_device(path: Path) -> None: + replace_once( + path, + ' device_map: Any = {"": local_rank} if torch.cuda.is_available() else "cpu"\n', + ''' spd_device = os.environ.get("SPD_DEVICE", "auto").lower() + if spd_device == "mps": + device_map: Any = {"": "mps"} + elif spd_device == "cpu": + device_map = "cpu" + elif spd_device == "cuda": + device_map = {"": local_rank} + else: + device_map = {"": local_rank} if torch.cuda.is_available() else "cpu" + +''', + ) + replace_once( + path, + " torch_dtype=torch.bfloat16,\n", + ''' torch_dtype=( + torch.float32 if os.environ.get("SPD_DEVICE", "auto").lower() == "mps" else torch.bfloat16 + ), +''', + ) + replace_once( + path, + " bf16=True,\n fp16=False,\n", + ''' bf16=torch.cuda.is_available(), + fp16=False, +''', + ) + + +def patch_eval_for_device(path: Path) -> None: + for _ in range(2): + replace_once( + path, + ''' dtype=torch.bfloat16, + device_map={"": 0} if torch.cuda.is_available() else None, + attn_implementation="flash_attention_2", +''', + ''' dtype=torch.float16 if os.environ.get("SPD_DEVICE", "auto").lower() == "mps" else torch.bfloat16, + device_map={"": "mps"} if os.environ.get("SPD_DEVICE", "auto").lower() == "mps" else ({"": 0} if torch.cuda.is_available() else None), + attn_implementation=os.environ.get("SPD_ATTN_IMPLEMENTATION", "sdpa"), +''', + ) + replace_once( + path, + ' map_loc = "cuda"\n', + ' map_loc = "mps" if os.environ.get("SPD_DEVICE", "auto").lower() == "mps" else ("cuda" if torch.cuda.is_available() else "cpu")\n', + ) + + +def patch_pipeline_inference_for_device(path: Path) -> None: + replace_once( + path, + ''' dtype=dtype, + device_map={"":0}, + trust_remote_code=True, +''', + ''' dtype=torch.float16 if os.environ.get("SPD_DEVICE", "auto").lower() == "mps" else dtype, + device_map={"": "mps"} if os.environ.get("SPD_DEVICE", "auto").lower() == "mps" else {"": 0}, + trust_remote_code=True, +''', + ) + replace_once( + path, + ' map_loc = "cuda" if torch.cuda.is_available() else "cpu"\n', + ' map_loc = "mps" if os.environ.get("SPD_DEVICE", "auto").lower() == "mps" else ("cuda" if torch.cuda.is_available() else "cpu")\n', + ) + + +def load_training_rows(dataset_name: str, split: str, limit: int) -> list[dict[str, Any]]: + from datasets import load_dataset + + wanted = max(1, int(limit)) + split_expr = f"{split}[:{wanted}]" + print(f"loading training data: {dataset_name} {split_expr}", flush=True) + try: + ds = load_dataset(dataset_name, split=split_expr) + except Exception: + fallback = f"train[:{wanted}]" + print(f"failed to load split {split_expr!r}; trying {fallback!r}", flush=True) + ds = load_dataset(dataset_name, split=fallback) + + rows: list[dict[str, Any]] = [] + for row in ds: + messages = row.get("messages") or row.get("conversations") + if not messages: + continue + normalized = [] + for message in messages: + role = message.get("role") or message.get("from") + content = message.get("content") or message.get("value") + if role is None or content is None: + normalized = [] + break + if role == "human": + role = "user" + if role == "gpt": + role = "assistant" + normalized.append({"role": str(role), "content": str(content)}) + if normalized: + rows.append({"messages": normalized}) + if not rows: + raise RuntimeError(f"no usable messages/conversations rows found in {dataset_name}") + return rows[:wanted] + + +def write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as handle: + for row in rows: + handle.write(json.dumps(row, ensure_ascii=False) + "\n") + print(f"wrote {len(rows)} rows -> {path}", flush=True) + + +def create_mini_eval_data(reference_dir: Path, out_dir: Path, rows_per_set: int) -> None: + source_root = reference_dir / "eval_data" + for name in ("mt_bench", "humaneval", "gsm8k"): + source = source_root / name / "question.jsonl" + dest_dir = out_dir / name + dest_dir.mkdir(parents=True, exist_ok=True) + dest = dest_dir / "question.jsonl" + copied = 0 + with source.open("r", encoding="utf-8") as src, dest.open("w", encoding="utf-8") as dst: + for line in src: + if copied >= rows_per_set: + break + if line.strip(): + dst.write(line) + copied += 1 + print(f"mini eval {name}: {copied} rows -> {dest}", flush=True) + + +def parse_stage_layer_boundaries(value: str) -> list[int] | None: + value = (value or "").strip() + if not value: + return None + boundaries = [int(part.strip()) for part in value.split(",") if part.strip()] + if not boundaries: + raise RuntimeError("--stage-layer-boundaries must not be empty when set") + if any(left >= right for left, right in zip(boundaries, boundaries[1:])): + raise RuntimeError(f"--stage-layer-boundaries must be strictly increasing: {boundaries}") + return boundaries + + +def derive_hidden_tap_indices(boundaries: list[int]) -> list[list[int]]: + rows: list[list[int]] = [] + for depth in range(len(boundaries), 0, -1): + rows.append([0, *boundaries[:depth]]) + return rows + + +def hidden_tap_rows_arg(args: argparse.Namespace) -> str: + explicit = args.shallow_hidden_layer_indices.strip() + if explicit: + return explicit + boundaries = parse_stage_layer_boundaries(args.stage_layer_boundaries) + if boundaries is None: + return "" + if len(boundaries) != args.num_stages: + raise RuntimeError( + f"--stage-layer-boundaries has {len(boundaries)} entries but --num-stages is {args.num_stages}" + ) + return ";".join(",".join(str(index) for index in row) for row in derive_hidden_tap_indices(boundaries)) + + +def resolve_draft_vocab_path(reference_dir: Path, value: str) -> Path: + path = Path(value) + if path.is_absolute(): + return path + return reference_dir / path + + +def build_draft_vocab_json( + *, + args: argparse.Namespace, + rows: list[dict[str, Any]], + output_path: Path, + vocab_size: int, +) -> Path: + from collections import Counter + + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(args.model_name, trust_remote_code=True) + counts: Counter[int] = Counter() + for row in rows: + text = row_text_for_vocab(tokenizer, row.get("messages") or []) + if not text: + continue + counts.update(int(token_id) for token_id in tokenizer.encode(text, add_special_tokens=False)) + for token_id in ( + getattr(tokenizer, "eos_token_id", None), + getattr(tokenizer, "pad_token_id", None), + getattr(tokenizer, "bos_token_id", None), + ): + if token_id is not None: + counts[int(token_id)] += 1 + if not counts: + raise RuntimeError("could not build draft vocab: no tokens counted") + token_ids = [ + token_id + for token_id, _ in sorted(counts.items(), key=lambda item: (-item[1], item[0]))[ + : max(1, int(vocab_size)) + ] + ] + token_ids = sorted(set(token_ids)) + output_path.parent.mkdir(parents=True, exist_ok=True) + output = { + "draft_vocab_size": len(token_ids), + "token_ids": token_ids, + "metadata": { + "base_model_path": args.model_name, + "source": "hf_train_eval_qwen06.py --build-draft-vocab-size", + "train_rows": len(rows), + "requested_vocab_size": int(vocab_size), + }, + } + output_path.write_text(json.dumps(output, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(f"wrote draft vocab ({len(token_ids)} ids) -> {output_path}", flush=True) + return output_path + + +def row_text_for_vocab(tokenizer: Any, messages: list[dict[str, Any]]) -> str: + if not messages: + return "" + try: + rendered = tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=False, + enable_thinking=False, + ) + if isinstance(rendered, str): + return rendered + except Exception: + pass + return "\n".join(str(message.get("content", "")) for message in messages) + + +def train_head(args: argparse.Namespace, reference_dir: Path, train_jsonl: Path, train_dir: Path) -> Path: + cmd = [ + sys.executable, + "train.py", + "--model_name", + args.model_name, + "--data_path", + str(train_jsonl), + "--num_stages", + str(args.num_stages), + "--num_spec_layers", + str(args.num_spec_layers), + "--epochs", + str(args.epochs), + "--batch_size", + str(args.batch_size), + "--gradient_accumulation_steps", + str(args.gradient_accumulation_steps), + "--lr", + str(args.learning_rate), + "--max_length", + str(args.max_length), + "--max_length_overflow", + "truncate", + "--min_length", + "10", + "--num_proc", + "2", + "--attn_implementation", + args.attn_implementation, + "--output_dir", + str(train_dir), + "--log_interval", + str(args.log_interval), + "--warmup_steps", + str(args.warmup_steps), + "--save_steps", + str(args.save_steps), + ] + hidden_rows = hidden_tap_rows_arg(args) + if hidden_rows: + cmd.extend(["--shallow_hidden_layer_indices", hidden_rows]) + if args.draft_vocab_json: + cmd.extend(["--draft_vocab_json", str(resolve_draft_vocab_path(reference_dir, args.draft_vocab_json))]) + started = time.perf_counter() + run(cmd, cwd=reference_dir, env=reference_env(args)) + elapsed = time.perf_counter() - started + ckpt = train_dir / "speculation_head_final.pt" + if not ckpt.is_file(): + raise FileNotFoundError(f"training did not produce {ckpt}") + print(f"training complete in {elapsed / 60.0:.2f} min: {ckpt}", flush=True) + write_skippy_spd_manifest(args, ckpt, train_dir / "skippy-spd-head.json") + return ckpt + + +def prepare_existing_spec_head(args: argparse.Namespace, train_dir: Path) -> Path | None: + local_path = args.spec_head_path.strip() + repo = args.spec_head_repo.strip() + filename = args.spec_head_file.strip() + if not local_path and not repo and not filename: + return None + if local_path and (repo or filename): + raise RuntimeError("--spec-head-path cannot be combined with --spec-head-repo/file") + if bool(repo) != bool(filename): + raise RuntimeError("--spec-head-repo and --spec-head-file must be set together") + + if repo: + from huggingface_hub import hf_hub_download + + source = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model")) + else: + source = Path(local_path).expanduser().resolve() + if not source.is_file(): + raise FileNotFoundError(f"speculation head checkpoint not found: {source}") + + train_dir.mkdir(parents=True, exist_ok=True) + dest = train_dir / "speculation_head_final.pt" + if source.resolve() != dest.resolve(): + if dest.exists(): + dest.unlink() + try: + os.link(source, dest) + except OSError: + shutil.copy2(source, dest) + write_skippy_spd_manifest(args, dest, train_dir / "skippy-spd-head.json") + print(f"prepared existing speculation head -> {dest}", flush=True) + return dest + + +def file_sha256(path: Path) -> str: + hasher = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + hasher.update(chunk) + return hasher.hexdigest() + + +def resolve_manifest_model_type(config: dict[str, Any], model_name: str) -> str | None: + model_type = config.get("model_type") + if isinstance(model_type, str) and model_type: + return model_type + + try: + from transformers import AutoConfig + + base_config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) + except Exception as exc: # noqa: BLE001 + print(f"warning: could not resolve base model_type for manifest: {exc}", flush=True) + return None + + model_type = getattr(base_config, "model_type", None) + return model_type if isinstance(model_type, str) and model_type else None + + +def write_skippy_spd_manifest(args: argparse.Namespace, ckpt: Path, manifest_path: Path) -> None: + import torch + + try: + checkpoint = torch.load(ckpt, map_location="cpu", weights_only=False) + except TypeError: + checkpoint = torch.load(ckpt, map_location="cpu") + config = checkpoint.get("config") if isinstance(checkpoint, dict) else None + if not isinstance(config, dict): + raise RuntimeError(f"{ckpt} does not contain a config dict") + + draft_token_ids = config.get("draft_token_ids") + stage_layer_boundaries = config.get("stage_layer_boundaries") + if stage_layer_boundaries is None: + stage_layer_boundaries = parse_stage_layer_boundaries(args.stage_layer_boundaries) + manifest_base_model_path = args.manifest_base_model_path.strip() + if not manifest_base_model_path: + manifest_base_model_path = config.get("base_model_path") or args.model_name + + manifest = { + "schema": "skippy-spd-head/v1", + "checkpoint": { + "path": ckpt.name, + "sha256": file_sha256(ckpt), + "bytes": ckpt.stat().st_size, + }, + "source": { + "format": "torch-speculation-head-v10", + "reference_repo": args.reference_repo, + "base_model_path": manifest_base_model_path, + "model_type": resolve_manifest_model_type(config, args.model_name), + "checkpoint_version": int(config.get("version", 0)), + }, + "topology": { + "hidden_size": int(config["hidden_size"]), + "vocab_size": int(config["vocab_size"]), + "draft_vocab_size": int(config.get("draft_vocab_size", config["vocab_size"])), + "num_stages": int(config["num_stages"]), + "stage_layer_boundaries": stage_layer_boundaries, + "num_spec_layers": int(config["num_spec_layers"]), + "trained_with_use_deepest": bool(config.get("trained_with_use_deepest", False)), + "shallow_hidden_layer_indices": config["shallow_hidden_layer_indices"], + "spec_init_from_base_layers": config.get("spec_init_from_base_layers"), + "draft_token_ids": draft_token_ids, + }, + } + manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(f"wrote Skippy SPD manifest -> {manifest_path}", flush=True) + + +def evaluate_head( + args: argparse.Namespace, + reference_dir: Path, + ckpt: Path, + eval_data: Path, + eval_dir: Path, +) -> None: + cmd = [ + sys.executable, + "eval.py", + "--spec_head_ckpt", + str(ckpt), + "--base_model_path", + args.model_name, + "--data_dir", + str(eval_data), + "--output_dir", + str(eval_dir), + "--gpus", + "0", + "--max_new_tokens", + str(args.max_new_tokens), + "--temperature", + "0.0", + "--draft_top_k", + str(args.draft_top_k), + "--no-baseline", + ] + started = time.perf_counter() + run(cmd, cwd=reference_dir, env=reference_env(args)) + elapsed = time.perf_counter() - started + print(f"eval complete in {elapsed / 60.0:.2f} min", flush=True) + + +def print_eval_summary(eval_dir: Path) -> None: + summary_dir = eval_dir / "summary" + if not summary_dir.is_dir(): + print(f"no eval summary dir found: {summary_dir}", flush=True) + return + for path in sorted(summary_dir.glob("*.json")): + with path.open("r", encoding="utf-8") as handle: + obj = json.load(handle) + print(f"summary: {path}", flush=True) + overall = obj.get("overall") or {} + if not overall and obj.get("results"): + overall = obj["results"][0].get("overall", {}) + interesting = { + key: overall.get(key) + for key in ( + "acceptance_rate", + "equivalent_accept_length", + "theoretical_speedup", + "new_tokens", + "decode_loop_steps", + ) + if key in overall + } + print(json.dumps(interesting or overall, indent=2, sort_keys=True), flush=True) + + +def resolve_upload_repo(value: str) -> str | None: + value = (value or "").strip() + if not value or value.lower() in {"none", "off", "false", "disabled"}: + return None + if value != "auto": + return value + from huggingface_hub import HfApi + + who = HfApi().whoami() + name = who.get("name") + if not name: + raise RuntimeError("could not resolve HF username for --upload-repo auto") + return f"{name}/skippy-spd-qwen06-proof" + + +def upload_artifacts(upload_repo: str | None, artifact_dir: Path, *, public: bool) -> None: + if upload_repo is None: + print(f"no upload repo configured; artifacts remain at {artifact_dir}", flush=True) + return + from huggingface_hub import HfApi + + api = HfApi() + api.create_repo(upload_repo, repo_type="model", private=not public, exist_ok=True) + api.upload_folder( + repo_id=upload_repo, + repo_type="model", + folder_path=str(artifact_dir), + path_in_repo=f"runs/{artifact_dir.name}", + ) + print(f"uploaded artifacts to hf://models/{upload_repo}/runs/{artifact_dir.name}", flush=True) + + +def reference_env(args: argparse.Namespace) -> dict[str, str]: + env = os.environ.copy() + if args.device != "auto": + env["SPD_DEVICE"] = args.device + env["SPD_ATTN_IMPLEMENTATION"] = args.attn_implementation + return env + + +def main() -> None: + args = parse_args() + work_dir = Path(args.work_dir).resolve() + reference_dir = work_dir / "speculative_pipeline_decoding" + artifact_dir = work_dir / "artifacts" / time.strftime("%Y%m%d-%H%M%S") + data_dir = work_dir / "data" + train_dir = artifact_dir / "train" + eval_dir = artifact_dir / "eval" + mini_eval_dir = data_dir / "mini_eval_data" + train_jsonl = data_dir / "train_conversations.jsonl" + + work_dir.mkdir(parents=True, exist_ok=True) + artifact_dir.mkdir(parents=True, exist_ok=True) + if args.topology_policy == "generic-plan": + write_topology_plan(args, artifact_dir) + print( + "generic-plan mode stops before training because the donor SPD head " + "still owns fixed per-stage projection tensors.", + flush=True, + ) + return + + clone_reference(args.reference_repo, reference_dir) + patch_reference_for_proof(reference_dir) + patch_reference_for_device(reference_dir, args.device) + + existing_ckpt = prepare_existing_spec_head(args, train_dir) + if existing_ckpt is not None: + ckpt = existing_ckpt + elif not args.skip_train: + rows = load_training_rows(args.dataset, args.dataset_split, args.train_rows) + write_jsonl(train_jsonl, rows) + if args.build_draft_vocab_size > 0: + draft_vocab_out = ( + Path(args.draft_vocab_out).expanduser().resolve() + if args.draft_vocab_out + else data_dir / f"draft_vocab_top_{args.build_draft_vocab_size}.json" + ) + args.draft_vocab_json = str( + build_draft_vocab_json( + args=args, + rows=rows, + output_path=draft_vocab_out, + vocab_size=args.build_draft_vocab_size, + ) + ) + ckpt = train_head(args, reference_dir, train_jsonl, train_dir) + else: + ckpt = train_dir / "speculation_head_final.pt" + if not ckpt.is_file(): + raise FileNotFoundError(f"--skip-train requires existing checkpoint at {ckpt}") + + if not args.skip_eval: + if mini_eval_dir.exists(): + shutil.rmtree(mini_eval_dir) + create_mini_eval_data(reference_dir, mini_eval_dir, args.eval_rows_per_set) + evaluate_head(args, reference_dir, ckpt, mini_eval_dir, eval_dir) + print_eval_summary(eval_dir) + + upload_repo = resolve_upload_repo(args.upload_repo) + upload_artifacts(upload_repo, artifact_dir, public=args.public) + + +if __name__ == "__main__": + main() diff --git a/evals/spd/simulate_latency.py b/evals/spd/simulate_latency.py new file mode 100644 index 0000000000..60106bf4d9 --- /dev/null +++ b/evals/spd/simulate_latency.py @@ -0,0 +1,293 @@ +#!/usr/bin/env python3 +"""Simulate split-stage latency from real SPD eval traces. + +This is not a serving benchmark. It consumes the per-sample JSONL emitted by the +reference SPD eval and applies an explicit latency model to its real +``new_tokens`` and ``decode_loop_steps`` counts. +""" + +from __future__ import annotations + +import argparse +import csv +import json +import statistics +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any + + +@dataclass(frozen=True) +class TraceTotals: + samples: int + stages: int + tokens: int + decode_steps: int + accepted_flags: int + acceptance_flags: int + + @property + def pipeline_cycles(self) -> float: + return self.decode_steps / max(self.stages, 1) + + @property + def aggregate_acceptance_rate(self) -> float: + if self.decode_steps == 0: + return 0.0 + return self.tokens / self.decode_steps + + @property + def equivalent_accept_length(self) -> float: + return self.stages * self.aggregate_acceptance_rate + + @property + def paper_theoretical_gain_pct(self) -> float: + if self.pipeline_cycles <= 0.0: + return 0.0 + return ((self.tokens / self.pipeline_cycles) - 1.0) * 100.0 + + +@dataclass(frozen=True) +class LatencyScenario: + stage_ms: tuple[float, ...] + hop_ms: float + + @property + def stages(self) -> int: + return len(self.stage_ms) + + @property + def serial_step_ms(self) -> float: + return sum(self.stage_ms) + self.hop_ms * max(self.stages - 1, 0) + + @property + def pipeline_slot_ms(self) -> float: + if not self.stage_ms: + return 0.0 + slots = [] + for index, stage_ms in enumerate(self.stage_ms): + outgoing_hop = self.hop_ms if index + 1 < self.stages else 0.0 + slots.append(stage_ms + outgoing_hop) + return max(slots) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--raw", required=True, type=Path, help="SPD eval raw per-sample JSONL") + parser.add_argument( + "--stage-ms", + default="4,4", + help="Comma-separated per-stage compute latency in ms, e.g. 4,4 or 3,5,6", + ) + parser.add_argument( + "--hop-ms", + default="0,1,5,10,25", + help="Comma-separated inter-stage activation hop latency scenarios in ms", + ) + parser.add_argument( + "--json", + action="store_true", + help="Emit JSON instead of the default human table", + ) + return parser.parse_args() + + +def parse_float_list(value: str, label: str) -> tuple[float, ...]: + try: + parsed = tuple(float(part.strip()) for part in value.split(",") if part.strip()) + except ValueError as error: + raise SystemExit(f"invalid {label}: {value!r}") from error + if not parsed: + raise SystemExit(f"{label} must contain at least one number") + if any(item < 0.0 for item in parsed): + raise SystemExit(f"{label} values must be non-negative") + return parsed + + +def load_rows(path: Path) -> list[dict[str, Any]]: + rows = [] + with path.open("r", encoding="utf-8") as handle: + for line_number, line in enumerate(handle, start=1): + line = line.strip() + if not line: + continue + try: + rows.append(json.loads(line)) + except json.JSONDecodeError as error: + raise SystemExit(f"{path}:{line_number}: invalid JSON") from error + if not rows: + raise SystemExit(f"{path} contains no rows") + return rows + + +def trace_totals(rows: list[dict[str, Any]]) -> TraceTotals: + stages = {int(row.get("num_stages") or 0) for row in rows} + stages.discard(0) + if len(stages) != 1: + raise SystemExit(f"expected one num_stages value, got {sorted(stages)}") + return TraceTotals( + samples=len(rows), + stages=stages.pop(), + tokens=sum(int(row.get("new_tokens") or 0) for row in rows), + decode_steps=sum(int(row.get("decode_loop_steps") or 0) for row in rows), + accepted_flags=sum(int(row.get("n_accepted") or 0) for row in rows), + acceptance_flags=sum(int(row.get("n_acceptance_flags") or 0) for row in rows), + ) + + +def percentile(values: list[float], pct: float) -> float: + if not values: + return 0.0 + values = sorted(values) + index = (len(values) - 1) * pct + lower = int(index) + upper = min(lower + 1, len(values) - 1) + fraction = index - lower + return values[lower] * (1.0 - fraction) + values[upper] * fraction + + +def simulate(rows: list[dict[str, Any]], totals: TraceTotals, scenario: LatencyScenario) -> dict[str, Any]: + if scenario.stages != totals.stages: + raise SystemExit( + f"--stage-ms has {scenario.stages} stages but trace has {totals.stages} stages" + ) + + serial_ms_by_sample = [ + int(row.get("new_tokens") or 0) * scenario.serial_step_ms for row in rows + ] + spd_ms_by_sample = [ + (int(row.get("decode_loop_steps") or 0) / totals.stages) * scenario.pipeline_slot_ms + for row in rows + ] + serial_total_ms = sum(serial_ms_by_sample) + spd_total_ms = sum(spd_ms_by_sample) + paper_like_no_spd_ms = totals.tokens * scenario.pipeline_slot_ms + + return { + "stage_ms": list(scenario.stage_ms), + "hop_ms": scenario.hop_ms, + "serial_step_ms": scenario.serial_step_ms, + "pipeline_slot_ms": scenario.pipeline_slot_ms, + "paper_like_no_spd_ms": paper_like_no_spd_ms, + "serial_split_no_spd_ms": serial_total_ms, + "spd_pipeline_ms": spd_total_ms, + "spd_vs_paper_like_no_spd": safe_ratio(paper_like_no_spd_ms, spd_total_ms), + "spd_vs_serial_split_no_spd": safe_ratio(serial_total_ms, spd_total_ms), + "serial_split_tok_s": safe_tok_s(totals.tokens, serial_total_ms), + "spd_pipeline_tok_s": safe_tok_s(totals.tokens, spd_total_ms), + "serial_request_p50_ms": percentile(serial_ms_by_sample, 0.5), + "serial_request_p95_ms": percentile(serial_ms_by_sample, 0.95), + "spd_request_p50_ms": percentile(spd_ms_by_sample, 0.5), + "spd_request_p95_ms": percentile(spd_ms_by_sample, 0.95), + "request_latency_p50_ratio": safe_ratio( + percentile(serial_ms_by_sample, 0.5), + percentile(spd_ms_by_sample, 0.5), + ), + "request_latency_p95_ratio": safe_ratio( + percentile(serial_ms_by_sample, 0.95), + percentile(spd_ms_by_sample, 0.95), + ), + } + + +def safe_ratio(numerator: float, denominator: float) -> float: + if denominator <= 0.0: + return 0.0 + return numerator / denominator + + +def safe_tok_s(tokens: int, total_ms: float) -> float: + if total_ms <= 0.0: + return 0.0 + return tokens / (total_ms / 1000.0) + + +def emit_table(totals: TraceTotals, results: list[dict[str, Any]]) -> None: + print("Trace") + print(f" samples: {totals.samples}") + print(f" stages: {totals.stages}") + print(f" generated tokens: {totals.tokens}") + print(f" decode loop steps: {totals.decode_steps}") + print(f" accepted draft flags: {totals.accepted_flags}/{totals.acceptance_flags}") + print(f" aggregate acceptance: {totals.aggregate_acceptance_rate:.4f}") + print(f" equivalent accept length: {totals.equivalent_accept_length:.4f}") + print(f" paper theoretical gain: {totals.paper_theoretical_gain_pct:.2f}%") + print() + print("Latency scenarios") + writer = csv.writer(sys.stdout) + writer.writerow( + [ + "hop_ms", + "slot_ms", + "serial_tok_s", + "spd_tok_s", + "spd_vs_serial", + "paper_like_gain", + "p50_serial_ms", + "p50_spd_ms", + "p95_serial_ms", + "p95_spd_ms", + ] + ) + for result in results: + writer.writerow( + [ + f"{result['hop_ms']:.3f}", + f"{result['pipeline_slot_ms']:.3f}", + f"{result['serial_split_tok_s']:.2f}", + f"{result['spd_pipeline_tok_s']:.2f}", + f"{result['spd_vs_serial_split_no_spd']:.3f}", + f"{result['spd_vs_paper_like_no_spd']:.3f}", + f"{result['serial_request_p50_ms']:.2f}", + f"{result['spd_request_p50_ms']:.2f}", + f"{result['serial_request_p95_ms']:.2f}", + f"{result['spd_request_p95_ms']:.2f}", + ] + ) + + +def main() -> None: + args = parse_args() + rows = load_rows(args.raw) + totals = trace_totals(rows) + stage_ms = parse_float_list(args.stage_ms, "--stage-ms") + hop_values = parse_float_list(args.hop_ms, "--hop-ms") + scenarios = [LatencyScenario(stage_ms=stage_ms, hop_ms=hop_ms) for hop_ms in hop_values] + results = [simulate(rows, totals, scenario) for scenario in scenarios] + + payload = { + "raw": str(args.raw), + "totals": { + "samples": totals.samples, + "stages": totals.stages, + "tokens": totals.tokens, + "decode_steps": totals.decode_steps, + "accepted_flags": totals.accepted_flags, + "acceptance_flags": totals.acceptance_flags, + "aggregate_acceptance_rate": totals.aggregate_acceptance_rate, + "equivalent_accept_length": totals.equivalent_accept_length, + "paper_theoretical_gain_pct": totals.paper_theoretical_gain_pct, + }, + "assumptions": { + "serial_split_no_spd": ( + "one generated token traverses every stage and inter-stage hop before the " + "next target token is known" + ), + "spd_pipeline": ( + "real reference decode_loop_steps are converted to pipeline cycles by " + "dividing by num_stages; each cycle costs the slowest stage slot" + ), + "stage_slot": "stage compute plus outgoing hop, except the final stage has no outgoing hop", + }, + "results": results, + } + + if args.json: + print(json.dumps(payload, indent=2, sort_keys=True)) + else: + emit_table(totals, results) + + +if __name__ == "__main__": + main() diff --git a/evals/spd/topology_plan.py b/evals/spd/topology_plan.py new file mode 100644 index 0000000000..41543ecd53 --- /dev/null +++ b/evals/spd/topology_plan.py @@ -0,0 +1,303 @@ +"""Generic SPD topology plan helpers. + +The current donor SPD head is fixed-stage. These helpers describe logical +hidden-state evidence separately from the donor architecture so the generic +sidecar path can evolve without pretending fixed `stage_projs` are topology +independent. +""" + +from __future__ import annotations + +import json +import random +import time +from argparse import Namespace +from pathlib import Path +from typing import Any + + +TOPOLOGY_PLAN_SCHEMA = "skippy-spd-topology-plan/v1" + + +def write_topology_plan(args: Namespace, artifact_dir: Path) -> Path: + plan = build_topology_plan(args) + if args.topology_plan_out: + path = Path(args.topology_plan_out).expanduser().resolve() + else: + path = artifact_dir / "topology" / "skippy-spd-topology-plan.json" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(plan, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(f"wrote generic SPD topology plan -> {path}", flush=True) + return path + + +def build_topology_plan(args: Namespace) -> dict[str, Any]: + model = resolve_model_config_summary(args) + num_layers = int(model["num_hidden_layers"]) + fixed_layer_taps = parse_fixed_layer_taps(getattr(args, "fixed_layer_taps", ""), num_layers) + if fixed_layer_taps: + return fixed_layer_tap_plan(args, model, fixed_layer_taps) + min_stages = validate_positive_int("topology_min_stages", args.topology_min_stages) + max_stages = validate_positive_int("topology_max_stages", args.topology_max_stages) + if min_stages > max_stages: + raise RuntimeError( + f"--topology-min-stages {min_stages} must be <= --topology-max-stages {max_stages}" + ) + if max_stages > num_layers: + raise RuntimeError( + f"--topology-max-stages {max_stages} cannot exceed num_hidden_layers {num_layers}" + ) + samples = validate_positive_int("topology_plan_samples", args.topology_plan_samples) + if not 0.0 <= args.topology_tap_dropout < 1.0: + raise RuntimeError("--topology-tap-dropout must be in [0.0, 1.0)") + + layouts = sample_topology_layouts( + num_layers=num_layers, + min_stages=min_stages, + max_stages=max_stages, + samples=samples, + seed=int(args.topology_seed), + tap_dropout=float(args.topology_tap_dropout), + ) + + return { + "schema": TOPOLOGY_PLAN_SCHEMA, + "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "model": model, + "hidden_state_convention": { + "index_0": "token embedding before layer 0", + "index_k": "hidden state after target layer k-1", + "final_index": num_layers, + }, + "policy": { + "name": "generic-contiguous-layer-plan", + "seed": int(args.topology_seed), + "samples": samples, + "min_stages": min_stages, + "max_stages": max_stages, + "tap_dropout": float(args.topology_tap_dropout), + "num_spec_layers": int(args.num_spec_layers), + "draft_top_k": int(args.draft_top_k), + }, + "layouts": layouts, + "notes": [ + "This plan is topology data for the generic sidecar path; it is not a trained head.", + "The current reference SPD architecture still has per-stage projection weights.", + "Generic training must consume logical layer-indexed taps plus masks rather than fixed stage IDs.", + ], + } + + +def parse_fixed_layer_taps(raw: str, num_layers: int) -> list[int]: + if not raw.strip(): + return [] + taps = [int(part.strip()) for part in raw.split(",") if part.strip()] + validate_sorted_unique_indices("fixed_layer_taps", taps) + if taps[0] != 0: + raise RuntimeError("--fixed-layer-taps must include 0 as the first tap") + if taps[-1] != num_layers: + raise RuntimeError( + f"--fixed-layer-taps must end at final hidden index {num_layers}, got {taps[-1]}" + ) + return taps + + +def fixed_layer_tap_plan( + args: Namespace, + model: dict[str, Any], + taps: list[int], +) -> dict[str, Any]: + num_layers = int(model["num_hidden_layers"]) + return { + "schema": TOPOLOGY_PLAN_SCHEMA, + "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "model": model, + "hidden_state_convention": { + "index_0": "token embedding before layer 0", + "index_k": "hidden state after target layer k-1", + "final_index": num_layers, + }, + "policy": { + "name": "fixed-layer-tap-plan", + "seed": int(args.topology_seed), + "samples": 1, + "min_stages": len(taps) - 1, + "max_stages": len(taps) - 1, + "tap_dropout": 0.0, + "num_spec_layers": int(args.num_spec_layers), + "draft_top_k": int(args.draft_top_k), + "fixed_layer_taps": list(taps), + }, + "layouts": [fixed_layer_tap_layout_record(taps)], + "notes": [ + "This plan fixes logical layer evidence for a GLM SPD control run.", + "Fixed layer taps are not physical Skippy stage or machine IDs.", + "Use this plan to prove model-specific SPD learnability before randomized layer-tap training.", + ], + } + + +def fixed_layer_tap_layout_record(taps: list[int]) -> dict[str, Any]: + return { + "source": "fixed-layer-taps", + "num_stages": len(taps) - 1, + "stage_layer_boundaries": taps[1:], + "shallow_hidden_layer_indices": [list(taps)], + "logical_hidden_taps": [logical_hidden_taps(taps)], + "tap_dropout": { + "probability": 0.0, + "required_indices": list(taps), + "dropout_applies_to": "none", + }, + } + + +def resolve_model_config_summary(args: Namespace) -> dict[str, Any]: + model_name = args.model_name + config_path = Path(model_name).expanduser() / "config.json" + if config_path.is_file(): + config = json.loads(config_path.read_text(encoding="utf-8")) + else: + try: + from transformers import AutoConfig + + hf_config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) + except Exception as exc: # noqa: BLE001 + if args.topology_num_hidden_layers <= 0: + raise RuntimeError( + "could not load model config for topology planning; pass " + "--topology-num-hidden-layers to generate a model-shape-only plan" + ) from exc + config = {} + else: + config = hf_config.to_dict() + + num_layers = int(config.get("num_hidden_layers") or args.topology_num_hidden_layers or 0) + validate_positive_int("num_hidden_layers", num_layers) + return { + "model_name": model_name, + "model_type": config.get("model_type"), + "architectures": config.get("architectures", []), + "num_hidden_layers": num_layers, + "hidden_size": config.get("hidden_size"), + "vocab_size": config.get("vocab_size"), + "num_nextn_predict_layers": config.get("num_nextn_predict_layers"), + } + + +def sample_topology_layouts( + *, + num_layers: int, + min_stages: int, + max_stages: int, + samples: int, + seed: int, + tap_dropout: float, +) -> list[dict[str, Any]]: + rng = random.Random(seed) + layouts: list[dict[str, Any]] = [] + seen: set[tuple[int, ...]] = set() + + for num_stages in range(min_stages, max_stages + 1): + boundaries = tuple(balanced_stage_boundaries(num_layers, num_stages)) + seen.add(boundaries) + layouts.append(topology_layout_record("balanced", boundaries, tap_dropout)) + + attempts = 0 + while len(layouts) < samples and attempts < samples * 64: + attempts += 1 + num_stages = rng.randint(min_stages, max_stages) + boundaries = tuple(random_stage_boundaries(num_layers, num_stages, rng)) + if boundaries in seen: + continue + seen.add(boundaries) + layouts.append(topology_layout_record("random", boundaries, tap_dropout)) + + if len(layouts) < samples: + raise RuntimeError( + f"could only generate {len(layouts)} unique topology layouts after {attempts} attempts" + ) + return layouts[:samples] + + +def topology_layout_record( + source: str, + boundaries: tuple[int, ...], + tap_dropout: float, +) -> dict[str, Any]: + rows = derive_hidden_tap_indices(list(boundaries)) + return { + "source": source, + "num_stages": len(boundaries), + "stage_layer_boundaries": list(boundaries), + "shallow_hidden_layer_indices": rows, + "logical_hidden_taps": [logical_hidden_taps(row) for row in rows], + "tap_dropout": { + "probability": float(tap_dropout), + "required_indices": [0, boundaries[-1]], + "dropout_applies_to": "intermediate logical layer taps", + }, + } + + +def balanced_stage_boundaries(num_layers: int, num_stages: int) -> list[int]: + validate_positive_int("num_layers", num_layers) + validate_positive_int("num_stages", num_stages) + if num_stages > num_layers: + raise RuntimeError(f"num_stages {num_stages} cannot exceed num_layers {num_layers}") + boundaries = [round(num_layers * (stage + 1) / num_stages) for stage in range(num_stages)] + if boundaries[-1] != num_layers: + boundaries[-1] = num_layers + if any(left >= right for left, right in zip(boundaries, boundaries[1:])): + raise RuntimeError(f"balanced topology is not strictly increasing: {boundaries}") + return boundaries + + +def random_stage_boundaries(num_layers: int, num_stages: int, rng: random.Random) -> list[int]: + validate_positive_int("num_layers", num_layers) + validate_positive_int("num_stages", num_stages) + if num_stages > num_layers: + raise RuntimeError(f"num_stages {num_stages} cannot exceed num_layers {num_layers}") + if num_stages == 1: + return [num_layers] + cuts = sorted(rng.sample(range(1, num_layers), num_stages - 1)) + return [*cuts, num_layers] + + +def derive_hidden_tap_indices(boundaries: list[int]) -> list[list[int]]: + rows: list[list[int]] = [] + for depth in range(len(boundaries), 0, -1): + rows.append([0, *boundaries[:depth]]) + return rows + + +def logical_hidden_taps(indices: list[int]) -> list[dict[str, Any]]: + if not indices: + raise RuntimeError("hidden tap row must not be empty") + max_index = max(indices) + taps: list[dict[str, Any]] = [] + for index in indices: + normalized_depth = 0.0 if max_index == 0 else round(float(index) / float(max_index), 6) + taps.append( + { + "layer_index": int(index), + "kind": "embedding" if index == 0 else "layer_output", + "normalized_depth": normalized_depth, + } + ) + return taps + + +def validate_positive_int(name: str, value: int) -> int: + if value <= 0: + raise RuntimeError(f"{name} must be greater than zero, got {value}") + return value + + +def validate_sorted_unique_indices(name: str, values: list[int]) -> None: + if not values: + raise RuntimeError(f"{name} must not be empty") + if any(value < 0 for value in values): + raise RuntimeError(f"{name} must not contain negative indices: {values}") + if values != sorted(set(values)): + raise RuntimeError(f"{name} must be sorted and unique: {values}") diff --git a/scripts/affected-crates.sh b/scripts/affected-crates.sh index 15e75ba055..e9c25d9cf0 100755 --- a/scripts/affected-crates.sh +++ b/scripts/affected-crates.sh @@ -56,8 +56,10 @@ WORKSPACE_MEMBERS=( "skippy-server" "metrics-server" "skippy-model-package" + "skippy-quantize" "model-package" "skippy-correctness" + "llama-quant-ffi" "llama-spec-bench" "skippy-bench" "skippy-prompt" diff --git a/scripts/plan-clippy-batches.sh b/scripts/plan-clippy-batches.sh index 724545dcc1..b5155c1084 100644 --- a/scripts/plan-clippy-batches.sh +++ b/scripts/plan-clippy-batches.sh @@ -58,8 +58,10 @@ WORKSPACE_MEMBERS=( "skippy-server" "metrics-server" "skippy-model-package" + "skippy-quantize" "model-package" "skippy-correctness" + "llama-quant-ffi" "llama-spec-bench" "skippy-bench" "skippy-prompt" @@ -143,6 +145,7 @@ weights = { "model-package": 5, "skippy-bench": 4, "skippy-model-package": 4, + "skippy-quantize": 4, "openai-frontend": 4, "model-artifact": 4, "model-hf": 4, diff --git a/scripts/skippy-bf16-to-quant-layer-package.sh b/scripts/skippy-bf16-to-quant-layer-package.sh new file mode 100755 index 0000000000..0c6183cf0d --- /dev/null +++ b/scripts/skippy-bf16-to-quant-layer-package.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash +set -euo pipefail + +skippy_quantize_bin="${SKIPPY_QUANTIZE_BIN:-target/release/skippy-quantize}" +if [[ ! -x "$skippy_quantize_bin" ]]; then + echo "missing executable: $skippy_quantize_bin" >&2 + echo "build it with: just skippy-quantize-standalone-release-build" >&2 + exit 1 +fi + +exec "$skippy_quantize_bin" quantize-layer-package "$@" diff --git a/scripts/skippy-ci-smoke.sh b/scripts/skippy-ci-smoke.sh index b528ba9487..fcfdec0f29 100755 --- a/scripts/skippy-ci-smoke.sh +++ b/scripts/skippy-ci-smoke.sh @@ -473,14 +473,12 @@ assert_json "$REPORT_DIR/recurrent-kv-recurrent.json" \ '.matches == true and .cache_hit_matches == true and .suffix_prefill_matches == null and .state_payload_kind == "kv-recurrent" and .cache_hit_repeats == 2 and .state_bytes > 0 and .payload_digest.recurrent_bytes > 0 and .payload_digest.kv_bytes > 0' PROMPT_PORT="$(pick_port)" -PROMPT_RETURN_PORT="$(pick_port)" PROMPT_CONFIG="$WORK_DIR/prompt-stage.json" PROMPT_LOG="$WORK_DIR/prompt-stage.log" PROMPT_IN="$WORK_DIR/prompt-input.txt" PROMPT_OUT="$WORK_DIR/prompt-output.log" PROMPT_BIND="127.0.0.1:${PROMPT_PORT}" -PROMPT_RETURN_BIND="127.0.0.1:${PROMPT_RETURN_PORT}" -write_stage_config "$PROMPT_CONFIG" "$DENSE_MODEL_ID" "$DENSE_MODEL_PATH" "$DENSE_LAYER_END" "$PROMPT_CTX_SIZE" "$PROMPT_BIND" "resident-kv" "$PROMPT_N_BATCH" "$PROMPT_N_UBATCH" "tcp://${PROMPT_RETURN_BIND}" +write_stage_config "$PROMPT_CONFIG" "$DENSE_MODEL_ID" "$DENSE_MODEL_PATH" "$DENSE_LAYER_END" "$PROMPT_CTX_SIZE" "$PROMPT_BIND" "resident-kv" "$PROMPT_N_BATCH" "$PROMPT_N_UBATCH" "driver" make_long_prompt_file "$PROMPT_IN" OPENAI_PORT="$(pick_port)" @@ -680,7 +678,6 @@ LLAMA_STAGE_BUILD_DIR="$LLAMA_BUILD_DIR" \ --tokenizer-layer-start 0 \ --tokenizer-layer-end "$DENSE_LAYER_END" \ --first-stage-addr "$PROMPT_BIND" \ - --direct-return-bind-addr "$PROMPT_RETURN_BIND" \ --ctx-size "$PROMPT_CTX_SIZE" \ --activation-width 2048 \ --activation-wire-dtype f16 \ diff --git a/third_party/llama.cpp/patches/0083-Allow-DeepSeek2-NextN-tensor-mapping.patch b/third_party/llama.cpp/patches/0083-Allow-DeepSeek2-NextN-tensor-mapping.patch deleted file mode 100644 index bf0ad2daec..0000000000 --- a/third_party/llama.cpp/patches/0083-Allow-DeepSeek2-NextN-tensor-mapping.patch +++ /dev/null @@ -1,29 +0,0 @@ -From 29409539414f3400462fbc155128bd22f7eb9d2f Mon Sep 17 00:00:00 2001 -From: James Dumay -Date: Thu, 11 Jun 2026 08:26:33 +1000 -Subject: [PATCH 83/89] Allow DeepSeek2 NextN tensor mapping - ---- - gguf-py/gguf/constants.py | 6 ++++++ - 1 file changed, 6 insertions(+) - -diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py -index 4b6dfea6..48147001 100644 ---- a/gguf-py/gguf/constants.py -+++ b/gguf-py/gguf/constants.py -@@ -3031,6 +3031,12 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { - MODEL_TENSOR.FFN_DOWN_SHEXP, - MODEL_TENSOR.FFN_UP_SHEXP, - MODEL_TENSOR.FFN_EXP_PROBS_B, -+ MODEL_TENSOR.NEXTN_EH_PROJ, -+ MODEL_TENSOR.NEXTN_EMBED_TOKENS, -+ MODEL_TENSOR.NEXTN_ENORM, -+ MODEL_TENSOR.NEXTN_HNORM, -+ MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, -+ MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, - ], - MODEL_ARCH.DEEPSEEK2OCR: [ - MODEL_TENSOR.TOKEN_EMBD, --- -2.54.0 (Apple Git-156) - diff --git a/third_party/llama.cpp/patches/0084-Retry-remote-safetensor-range-reads.patch b/third_party/llama.cpp/patches/0084-Retry-remote-safetensor-range-reads.patch deleted file mode 100644 index 864bbb6f63..0000000000 --- a/third_party/llama.cpp/patches/0084-Retry-remote-safetensor-range-reads.patch +++ /dev/null @@ -1,59 +0,0 @@ -From e74383485f22d7458e6e97848269daa8a93d5274 Mon Sep 17 00:00:00 2001 -From: James Dumay -Date: Thu, 11 Jun 2026 09:20:58 +1000 -Subject: [PATCH 84/89] Retry remote safetensor range reads - ---- - gguf-py/gguf/utility.py | 28 ++++++++++++++++++++++++---- - 1 file changed, 24 insertions(+), 4 deletions(-) - -diff --git a/gguf-py/gguf/utility.py b/gguf-py/gguf/utility.py -index 154351d8..ec44e55f 100644 ---- a/gguf-py/gguf/utility.py -+++ b/gguf-py/gguf/utility.py -@@ -224,6 +224,8 @@ class SafetensorRemote: - If size is not specified, it will read the entire file. - """ - import requests -+ import sys -+ import time - from urllib.parse import urlparse - - parsed_url = urlparse(url) -@@ -233,11 +235,29 @@ class SafetensorRemote: - headers = cls._get_request_headers() - if size > -1: - headers["Range"] = f"bytes={start}-{start + size}" -- response = requests.get(url, allow_redirects=True, headers=headers) -- response.raise_for_status() -+ max_attempts = 8 -+ for attempt in range(max_attempts): -+ try: -+ response = requests.get(url, allow_redirects=True, headers=headers, timeout=120) -+ response.raise_for_status() -+ content = response.content[slice(size if size > -1 else None)] -+ if size < 0 or len(content) == size: -+ return content -+ raise requests.ConnectionError( -+ f"short read for range {start}+{size}: got {len(content)} bytes", -+ ) -+ except requests.RequestException as e: -+ if attempt + 1 >= max_attempts: -+ raise -+ sleep_s = min(2 ** attempt, 30) -+ print( -+ f"warning: retrying remote safetensor range read after error: {e} " -+ f"(attempt {attempt + 1}/{max_attempts}, sleep {sleep_s}s)", -+ file=sys.stderr, -+ ) -+ time.sleep(sleep_s) - -- # Get raw byte data -- return response.content[slice(size if size > -1 else None)] -+ raise RuntimeError("unreachable") - - @classmethod - def check_file_exist(cls, url: str) -> bool: --- -2.54.0 (Apple Git-156) - diff --git a/third_party/llama.cpp/patches/0085-Allow-quantizing-appended-MTP-layers.patch b/third_party/llama.cpp/patches/0085-Allow-quantizing-appended-MTP-layers.patch deleted file mode 100644 index 1d63b2a539..0000000000 --- a/third_party/llama.cpp/patches/0085-Allow-quantizing-appended-MTP-layers.patch +++ /dev/null @@ -1,25 +0,0 @@ -From b1c6596c2f333447967f157260074d43948f2261 Mon Sep 17 00:00:00 2001 -From: James Dumay -Date: Thu, 11 Jun 2026 10:46:54 +1000 -Subject: [PATCH 85/89] Allow quantizing appended MTP layers - ---- - src/llama-quant.cpp | 2 +- - 1 file changed, 1 insertion(+), 1 deletion(-) - -diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp -index d4e7738a..0fe62d6b 100644 ---- a/src/llama-quant.cpp -+++ b/src/llama-quant.cpp -@@ -847,7 +847,7 @@ static void init_quantize_state_counters(quantize_state_impl & qs, std::vector -Date: Thu, 11 Jun 2026 11:34:36 +1000 -Subject: [PATCH 86/89] Fix tokenizer regex during GGUF conversion - ---- - gguf-py/gguf/vocab.py | 22 +++++++++++++++++----- - 1 file changed, 17 insertions(+), 5 deletions(-) - -diff --git a/gguf-py/gguf/vocab.py b/gguf-py/gguf/vocab.py -index d93b94f2..52618906 100644 ---- a/gguf-py/gguf/vocab.py -+++ b/gguf-py/gguf/vocab.py -@@ -580,11 +580,23 @@ class LlamaHfVocab(Vocab): - - # Allow the tokenizer to default to slow or fast versions. - # Explicitly set tokenizer to use local paths. -- self.tokenizer = AutoTokenizer.from_pretrained( -- base_path, -- cache_dir=base_path, -- local_files_only=True, -- ) -+ tokenizer_kwargs = { -+ "cache_dir": base_path, -+ "local_files_only": True, -+ } -+ try: -+ self.tokenizer = AutoTokenizer.from_pretrained( -+ base_path, -+ fix_mistral_regex=True, -+ **tokenizer_kwargs, -+ ) -+ except TypeError as e: -+ if "fix_mistral_regex" not in str(e): -+ raise -+ self.tokenizer = AutoTokenizer.from_pretrained( -+ base_path, -+ **tokenizer_kwargs, -+ ) - assert self.tokenizer.is_fast # assume tokenizer.json is used # ty: ignore[unresolved-attribute] - - # Initialize lists and dictionaries for added tokens --- -2.54.0 (Apple Git-156) - diff --git a/third_party/llama.cpp/patches/0087-Fix-GLM-tokenizer-regex-during-conversion.patch b/third_party/llama.cpp/patches/0087-Fix-GLM-tokenizer-regex-during-conversion.patch deleted file mode 100644 index 20e7a5e549..0000000000 --- a/third_party/llama.cpp/patches/0087-Fix-GLM-tokenizer-regex-during-conversion.patch +++ /dev/null @@ -1,67 +0,0 @@ -From 6f8f9772cb7402e758bd10474d124ecad504e963 Mon Sep 17 00:00:00 2001 -From: James Dumay -Date: Thu, 11 Jun 2026 11:39:33 +1000 -Subject: [PATCH 87/89] Fix GLM tokenizer regex during conversion - ---- - conversion/base.py | 23 +++++++++++++++++------ - 1 file changed, 17 insertions(+), 6 deletions(-) - -diff --git a/conversion/base.py b/conversion/base.py -index 9d81c19b..ffe59da9 100644 ---- a/conversion/base.py -+++ b/conversion/base.py -@@ -1129,6 +1129,20 @@ class TextModel(ModelBase): - if "rope_type" not in self.rope_parameters and (rope_type := self.rope_parameters.get("type")) is not None: - self.rope_parameters["rope_type"] = rope_type - -+ def load_hf_tokenizer(self, **kwargs): -+ from transformers import AutoTokenizer -+ -+ try: -+ return AutoTokenizer.from_pretrained( -+ self.dir_model, -+ fix_mistral_regex=True, -+ **kwargs, -+ ) -+ except TypeError as e: -+ if "fix_mistral_regex" not in str(e): -+ raise -+ return AutoTokenizer.from_pretrained(self.dir_model, **kwargs) -+ - @classmethod - def __init_subclass__(cls): - # can't use an abstract property, because overriding it without type errors -@@ -1333,8 +1347,7 @@ class TextModel(ModelBase): - tokens: list[str] = [] - toktypes: list[int] = [] - -- from transformers import AutoTokenizer -- tokenizer = AutoTokenizer.from_pretrained(self.dir_model) -+ tokenizer = self.load_hf_tokenizer() - vocab_size = self.hparams.get("vocab_size", len(tokenizer.vocab)) # ty: ignore[unresolved-attribute] - assert max(tokenizer.vocab.values()) < vocab_size # ty: ignore[unresolved-attribute] - -@@ -2060,8 +2073,7 @@ class TextModel(ModelBase): - self.gguf_writer.add_pooling_type(pooling_type) - - def _set_vocab_glmedge(self): -- from transformers import AutoTokenizer -- tokenizer = AutoTokenizer.from_pretrained(self.dir_model) -+ tokenizer = self.load_hf_tokenizer() - special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True) - tokens, toktypes, tokpre = self.get_vocab_base() - self.gguf_writer.add_tokenizer_model("gpt2") -@@ -2075,8 +2087,7 @@ class TextModel(ModelBase): - special_vocab.add_to_gguf(self.gguf_writer) - - def _set_vocab_glm(self): -- from transformers import AutoTokenizer -- tokenizer = AutoTokenizer.from_pretrained(self.dir_model) -+ tokenizer = self.load_hf_tokenizer() - special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True) - tokens, toktypes, tokpre = self.get_vocab_base() - self.gguf_writer.add_tokenizer_model("gpt2") --- -2.54.0 (Apple Git-156) - diff --git a/third_party/llama.cpp/patches/0088-Recognize-fixed-GLM-tokenizer-signature.patch b/third_party/llama.cpp/patches/0088-Recognize-fixed-GLM-tokenizer-signature.patch deleted file mode 100644 index e420b0a86d..0000000000 --- a/third_party/llama.cpp/patches/0088-Recognize-fixed-GLM-tokenizer-signature.patch +++ /dev/null @@ -1,26 +0,0 @@ -From e0b78d67fa560e66ad16c6b3d75e73b80390f405 Mon Sep 17 00:00:00 2001 -From: James Dumay -Date: Thu, 11 Jun 2026 11:45:43 +1000 -Subject: [PATCH 88/89] Recognize fixed GLM tokenizer signature - ---- - conversion/base.py | 3 +++ - 1 file changed, 3 insertions(+) - -diff --git a/conversion/base.py b/conversion/base.py -index ffe59da9..6ece31fc 100644 ---- a/conversion/base.py -+++ b/conversion/base.py -@@ -1424,6 +1424,9 @@ class TextModel(ModelBase): - if chkhsh == "cdf5f35325780597efd76153d4d1c16778f766173908894c04afc20108536267": - # ref: https://huggingface.co/zai-org/GLM-4.7-Flash - res = "glm4" -+ if chkhsh == "bd30abdd4d2e79a3b76199e7380500230a211249db4abad44a34ab28117fd4c6": -+ # ref: https://huggingface.co/zai-org/GLM-4.7-Flash with fix_mistral_regex=True -+ res = "glm4" - if chkhsh == "1431a23e583c97432bc230bff598d103ddb5a1f89960c8f1d1051aaa944d0b35": - # ref: https://huggingface.co/sapienzanlp/Minerva-7B-base-v1.0 - res = "minerva-7b" --- -2.54.0 (Apple Git-156) - diff --git a/third_party/llama.cpp/patches/0089-Add-GLM-native-MTP-support.patch b/third_party/llama.cpp/patches/0089-Add-GLM-native-MTP-graph-support.patch similarity index 69% rename from third_party/llama.cpp/patches/0089-Add-GLM-native-MTP-support.patch rename to third_party/llama.cpp/patches/0089-Add-GLM-native-MTP-graph-support.patch index 679a07be11..0d8ab01587 100644 --- a/third_party/llama.cpp/patches/0089-Add-GLM-native-MTP-support.patch +++ b/third_party/llama.cpp/patches/0089-Add-GLM-native-MTP-graph-support.patch @@ -1,174 +1,17 @@ -From 928b982ab33185d9b99266aaf50fee7d6855f15e Mon Sep 17 00:00:00 2001 +From 2e7aafec4d0683739d093177653f610451d83621 Mon Sep 17 00:00:00 2001 From: Mesh-LLM CI -Date: Sat, 13 Jun 2026 12:28:34 +1000 -Subject: [PATCH 89/89] Add GLM native MTP support +Date: Wed, 17 Jun 2026 07:36:49 +1000 +Subject: [PATCH] Add GLM native MTP graph support --- - conversion/glm.py | 11 ++ - gguf-py/gguf/tensor_mapping.py | 1 + - scripts/check-glm51-mtp-metadata.py | 110 +++++++++++ - src/llama-context.cpp | 15 +- - src/llama-model.cpp | 8 + - src/models/deepseek2.cpp | 280 +++++++++++++++++++++++++++- - src/models/glm-dsa.cpp | 12 +- - src/models/models.h | 5 + - 8 files changed, 432 insertions(+), 10 deletions(-) - create mode 100755 scripts/check-glm51-mtp-metadata.py + src/llama-context.cpp | 15 ++- + src/llama-model.cpp | 8 ++ + src/models/deepseek2.cpp | 284 ++++++++++++++++++++++++++++++++++++++- + src/models/glm-dsa.cpp | 12 +- + src/models/models.h | 5 + + src/skippy.cpp | 3 + + 6 files changed, 317 insertions(+), 10 deletions(-) -diff --git a/conversion/glm.py b/conversion/glm.py -index 64193772..87d443fe 100644 ---- a/conversion/glm.py -+++ b/conversion/glm.py -@@ -204,10 +204,21 @@ class Glm4MoeModel(TextModel): - @ModelBase.register("Glm4MoeLiteForCausalLM") - class Glm4MoeLiteModel(DeepseekV2Model): - model_arch = gguf.MODEL_ARCH.DEEPSEEK2 -+ skip_mtp = False -+ -+ def __init__(self, *args, **kwargs): -+ super().__init__(*args, **kwargs) -+ self.block_count = self.hparams["num_hidden_layers"] + self.hparams.get("num_nextn_predict_layers", 0) -+ self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) - - def set_vocab(self): - return self._set_vocab_glm() - -+ def set_gguf_parameters(self): -+ super().set_gguf_parameters() -+ if (num_nextn_predict_layers := self.hparams.get("num_nextn_predict_layers")) is not None: -+ self.gguf_writer.add_nextn_predict_layers(num_nextn_predict_layers) -+ - - @ModelBase.register("GlmMoeDsaForCausalLM") - class GlmMoeDsaModel(DeepseekV2Model): -diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py -index 5f1e2881..0b1614a9 100644 ---- a/gguf-py/gguf/tensor_mapping.py -+++ b/gguf-py/gguf/tensor_mapping.py -@@ -470,6 +470,7 @@ class TensorNameMap: - - MODEL_TENSOR.FFN_EXP_PROBS_B: ( - "model.layers.{bid}.mlp.gate.e_score_correction", # deepseek-v3 dots1 -+ "model.layers.{bid}.mlp.gate.e_score_correction_bias", # glm-dsa - "model.layers.{bid}.mlp.moe_statics.e_score_correction", # ernie4.5-moe - "model.layers.{bid}.mlp.gate.expert_bias", # bailingmoe2 - "model.layers.{bid}.mlp.expert_bias", # afmoe -diff --git a/scripts/check-glm51-mtp-metadata.py b/scripts/check-glm51-mtp-metadata.py -new file mode 100755 -index 00000000..208b46b9 ---- /dev/null -+++ b/scripts/check-glm51-mtp-metadata.py -@@ -0,0 +1,110 @@ -+#!/usr/bin/env python3 -+from __future__ import annotations -+ -+import argparse -+import json -+import sys -+from pathlib import Path -+ -+ -+def load_json(path: Path) -> dict: -+ with path.open("r", encoding="utf-8") as f: -+ return json.load(f) -+ -+ -+def import_gguf(repo: Path): -+ sys.path.insert(0, str(repo / "gguf-py")) -+ import gguf # type: ignore -+ -+ return gguf -+ -+ -+def map_name(tensor_map, name: str) -> str | None: -+ return tensor_map.get_name(name, try_suffixes=(".weight", ".bias")) -+ -+ -+def require(condition: bool, message: str, failures: list[str]) -> None: -+ if not condition: -+ failures.append(message) -+ -+ -+def main() -> int: -+ parser = argparse.ArgumentParser(description="Validate GLM-5.1 native MTP metadata and GGUF tensor mapping.") -+ parser.add_argument("--repo", type=Path, default=Path(__file__).resolve().parents[1], help="llama.cpp checkout") -+ parser.add_argument("--meta-dir", type=Path, default=Path("/tmp/glm51-meta"), help="directory with config and safetensor index") -+ args = parser.parse_args() -+ -+ config = load_json(args.meta_dir / "config.json") -+ index = load_json(args.meta_dir / "model.safetensors.index.json") -+ weight_names = set(index["weight_map"].keys()) -+ -+ gguf = import_gguf(args.repo) -+ block_count = int(config["num_hidden_layers"]) + int(config.get("num_nextn_predict_layers", 0)) -+ tensor_map = gguf.get_tensor_name_map(gguf.MODEL_ARCH.GLM_DSA, block_count) -+ -+ failures: list[str] = [] -+ -+ require(config.get("architectures") == ["GlmMoeDsaForCausalLM"], "unexpected architectures", failures) -+ require(config.get("model_type") == "glm_moe_dsa", "unexpected model_type", failures) -+ require(config.get("num_hidden_layers") == 78, "unexpected num_hidden_layers", failures) -+ require(config.get("num_nextn_predict_layers") == 1, "unexpected num_nextn_predict_layers", failures) -+ -+ mtp_layer = int(config["num_hidden_layers"]) -+ require(f"model.layers.{mtp_layer}.eh_proj.weight" in weight_names, "missing MTP eh_proj", failures) -+ require(f"model.layers.{mtp_layer}.enorm.weight" in weight_names, "missing MTP enorm", failures) -+ require(f"model.layers.{mtp_layer}.hnorm.weight" in weight_names, "missing MTP hnorm", failures) -+ require(f"model.layers.{mtp_layer}.shared_head.norm.weight" in weight_names, "missing MTP shared head norm", failures) -+ require(f"model.layers.{mtp_layer}.embed_tokens.weight" not in weight_names, "unexpected MTP-specific embed_tokens", failures) -+ require(f"model.layers.{mtp_layer}.shared_head.head.weight" not in weight_names, "unexpected MTP-specific shared head", failures) -+ -+ expected_mappings = { -+ f"model.layers.{mtp_layer}.eh_proj.weight": f"blk.{mtp_layer}.nextn.eh_proj.weight", -+ f"model.layers.{mtp_layer}.enorm.weight": f"blk.{mtp_layer}.nextn.enorm.weight", -+ f"model.layers.{mtp_layer}.hnorm.weight": f"blk.{mtp_layer}.nextn.hnorm.weight", -+ f"model.layers.{mtp_layer}.shared_head.norm.weight": f"blk.{mtp_layer}.nextn.shared_head_norm.weight", -+ f"model.layers.{mtp_layer}.input_layernorm.weight": f"blk.{mtp_layer}.attn_norm.weight", -+ f"model.layers.{mtp_layer}.self_attn.indexer.k_norm.weight": f"blk.{mtp_layer}.indexer.k_norm.weight", -+ f"model.layers.{mtp_layer}.self_attn.indexer.k_norm.bias": f"blk.{mtp_layer}.indexer.k_norm.bias", -+ f"model.layers.{mtp_layer}.self_attn.indexer.weights_proj.weight": f"blk.{mtp_layer}.indexer.proj.weight", -+ f"model.layers.{mtp_layer}.self_attn.indexer.wk.weight": f"blk.{mtp_layer}.indexer.attn_k.weight", -+ f"model.layers.{mtp_layer}.self_attn.indexer.wq_b.weight": f"blk.{mtp_layer}.indexer.attn_q_b.weight", -+ f"model.layers.{mtp_layer}.mlp.gate.weight": f"blk.{mtp_layer}.ffn_gate_inp.weight", -+ f"model.layers.{mtp_layer}.mlp.gate.e_score_correction_bias": f"blk.{mtp_layer}.exp_probs_b", -+ f"model.layers.{mtp_layer}.mlp.shared_experts.down_proj.weight": f"blk.{mtp_layer}.ffn_down_shexp.weight", -+ f"model.layers.{mtp_layer}.mlp.shared_experts.gate_proj.weight": f"blk.{mtp_layer}.ffn_gate_shexp.weight", -+ f"model.layers.{mtp_layer}.mlp.shared_experts.up_proj.weight": f"blk.{mtp_layer}.ffn_up_shexp.weight", -+ } -+ -+ for src, expected in expected_mappings.items(): -+ require(src in weight_names, f"missing source tensor {src}", failures) -+ actual = map_name(tensor_map, src) -+ require(actual == expected, f"mapping mismatch {src}: expected {expected}, got {actual}", failures) -+ -+ expert_src = f"model.layers.{mtp_layer}.mlp.experts.down_proj.weight" -+ expert_expected = f"blk.{mtp_layer}.ffn_down_exps.weight" -+ expert_actual = map_name(tensor_map, expert_src) -+ require(expert_actual == expert_expected, f"merged expert mapping mismatch: expected {expert_expected}, got {expert_actual}", failures) -+ -+ print(f"repo={args.repo}") -+ print(f"meta_dir={args.meta_dir}") -+ print(f"architecture={config['architectures'][0]}") -+ print(f"model_type={config['model_type']}") -+ print(f"trunk_layers={config['num_hidden_layers']}") -+ print(f"nextn_layers={config.get('num_nextn_predict_layers', 0)}") -+ print(f"total_blocks={block_count}") -+ print(f"mtp_layer={mtp_layer}") -+ print(f"tensor_count={len(weight_names)}") -+ print(f"source_total_size={index.get('metadata', {}).get('total_size')}") -+ -+ if failures: -+ print("status=FAIL") -+ for failure in failures: -+ print(f"failure={failure}") -+ return 1 -+ -+ print("status=PASS") -+ return 0 -+ -+ -+if __name__ == "__main__": -+ raise SystemExit(main()) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 87f75b94..5cf50e6c 100644 --- a/src/llama-context.cpp @@ -217,7 +60,7 @@ index a1dc66c6..869f80b6 100644 GGML_ASSERT(hparams.is_swa_any()); diff --git a/src/models/deepseek2.cpp b/src/models/deepseek2.cpp -index d1b261b4..b71fc076 100644 +index d1b261b4..d0ddb3c5 100644 --- a/src/models/deepseek2.cpp +++ b/src/models/deepseek2.cpp @@ -20,6 +20,9 @@ void llama_model_deepseek2::load_arch_hparams(llama_model_loader & ml) { @@ -272,10 +115,21 @@ index d1b261b4..b71fc076 100644 cur = ggml_get_rows(ctx0, cur, inp_out_ids); inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); } -@@ -429,6 +445,13 @@ llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_p - } +@@ -430,6 +446,10 @@ llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_p cur = inpL; + if (stage_filtered && !stage_filter.include_output) { ++ if (!cparams.embeddings_nextn_masked && inp_out_ids) { ++ cur = ggml_get_rows(ctx0, cur, inp_out_ids); ++ } ++ + cb(cur, "stage_boundary", il_end - 1); + res->t_embd = cur; + ggml_build_forward_expand(gf, cur); +@@ -438,6 +458,13 @@ llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_p + + cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); + + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + @@ -283,10 +137,10 @@ index d1b261b4..b71fc076 100644 + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + - if (stage_filtered && !stage_filter.include_output) { - cb(cur, "stage_boundary", il_end - 1); - res->t_embd = cur; -@@ -449,3 +472,256 @@ llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_p + cb(cur, "result_norm", -1); + res->t_embd = cur; + +@@ -449,3 +476,256 @@ llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_p ggml_build_forward_expand(gf, cur); } @@ -602,6 +456,19 @@ index ee3aff07..75c0ffcf 100644 std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; +diff --git a/src/skippy.cpp b/src/skippy.cpp +index 4dcd0d96..20cff23f 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -880,6 +880,9 @@ struct skippy_filter_scope { + filter.enabled = true; + filter.layer_start = config->layer_start; + filter.layer_end = config->layer_end; ++ if (config->include_output && filter.layer_end < std::numeric_limits::max()) { ++ filter.layer_end += 1; ++ } + filter.include_embeddings = config->include_embeddings; + filter.include_output = config->include_output; + llama_model_loader_set_stage_filter(filter); -- 2.54.0 (Apple Git-156) - diff --git a/third_party/llama.cpp/patches/0090-Reduce-llama-quantize-BF16-memory-residency.patch b/third_party/llama.cpp/patches/0090-Reduce-llama-quantize-BF16-memory-residency.patch deleted file mode 100644 index 28134af205..0000000000 --- a/third_party/llama.cpp/patches/0090-Reduce-llama-quantize-BF16-memory-residency.patch +++ /dev/null @@ -1,201 +0,0 @@ -From 5916dc8a91bc3d60d6822866eee5ded34d95f858 Mon Sep 17 00:00:00 2001 -From: James Dumay -Date: Mon, 15 Jun 2026 09:09:02 +1000 -Subject: [PATCH 01/10] Reduce llama-quantize BF16 memory residency - ---- - src/llama-quant.cpp | 117 +++++++++++++++++++++++++++++--------------- - 1 file changed, 78 insertions(+), 39 deletions(-) - -diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp -index 912ef9ae7..d77866392 100644 ---- a/src/llama-quant.cpp -+++ b/src/llama-quant.cpp -@@ -209,9 +209,9 @@ struct tensor_metadata { - // dequantization - // - --static void llama_tensor_dequantize_impl( -+static void llama_tensor_dequantize_chunk_impl( - ggml_tensor * tensor, std::vector> & output, std::vector & workers, -- const size_t nelements, const int nthread -+ const int64_t start, const size_t nelements, const int nthread - ) { - if (output.size() < nelements) { - output.resize(nelements); -@@ -228,19 +228,6 @@ static void llama_tensor_dequantize_impl( - throw std::runtime_error(format("cannot dequantize/convert tensor type %s", ggml_type_name(tensor->type))); - } - -- if (nthread < 2) { -- if (tensor->type == GGML_TYPE_F16) { -- ggml_fp16_to_fp32_row((ggml_fp16_t *)tensor->data, f32_output, nelements); -- } else if (tensor->type == GGML_TYPE_BF16) { -- ggml_bf16_to_fp32_row((ggml_bf16_t *)tensor->data, f32_output, nelements); -- } else if (ggml_is_quantized(tensor->type)) { -- qtype->to_float(tensor->data, f32_output, nelements); -- } else { -- GGML_ABORT("fatal error"); // unreachable -- } -- return; -- } -- - size_t block_size; - if (tensor->type == GGML_TYPE_F16 || - tensor->type == GGML_TYPE_BF16) { -@@ -249,9 +236,27 @@ static void llama_tensor_dequantize_impl( - block_size = (size_t)ggml_blck_size(tensor->type); - } - -+ GGML_ASSERT(start >= 0); -+ GGML_ASSERT((size_t) start % block_size == 0); -+ GGML_ASSERT(nelements % block_size == 0); -+ - size_t block_size_bytes = ggml_type_size(tensor->type); -+ const size_t input_offset = ((size_t) start / block_size) * block_size_bytes; -+ uint8_t * input_data = (uint8_t *) tensor->data + input_offset; -+ -+ if (nthread < 2) { -+ if (tensor->type == GGML_TYPE_F16) { -+ ggml_fp16_to_fp32_row((ggml_fp16_t *) input_data, f32_output, nelements); -+ } else if (tensor->type == GGML_TYPE_BF16) { -+ ggml_bf16_to_fp32_row((ggml_bf16_t *) input_data, f32_output, nelements); -+ } else if (ggml_is_quantized(tensor->type)) { -+ qtype->to_float(input_data, f32_output, nelements); -+ } else { -+ GGML_ABORT("fatal error"); // unreachable -+ } -+ return; -+ } - -- GGML_ASSERT(nelements % block_size == 0); - size_t nblocks = nelements / block_size; - size_t blocks_per_thread = nblocks / nthread; - size_t spare_blocks = nblocks - (blocks_per_thread * nthread); // if blocks aren't divisible by thread count -@@ -264,7 +269,7 @@ static void llama_tensor_dequantize_impl( - size_t thr_elems = thr_blocks * block_size; // number of elements for this thread - size_t thr_block_bytes = thr_blocks * block_size_bytes; // number of input bytes for this thread - -- auto compute = [qtype] (ggml_type typ, uint8_t * inbuf, float * outbuf, int nels) { -+ auto compute = [qtype] (ggml_type typ, uint8_t * inbuf, float * outbuf, int64_t nels) { - if (typ == GGML_TYPE_F16) { - ggml_fp16_to_fp32_row((ggml_fp16_t *)inbuf, outbuf, nels); - } else if (typ == GGML_TYPE_BF16) { -@@ -273,7 +278,7 @@ static void llama_tensor_dequantize_impl( - qtype->to_float(inbuf, outbuf, nels); - } - }; -- workers.emplace_back(compute, tensor->type, (uint8_t *) tensor->data + in_buff_offs, f32_output + out_buff_offs, thr_elems); -+ workers.emplace_back(compute, tensor->type, input_data + in_buff_offs, f32_output + out_buff_offs, thr_elems); - in_buff_offs += thr_block_bytes; - out_buff_offs += thr_elems; - } -@@ -281,6 +286,13 @@ static void llama_tensor_dequantize_impl( - workers.clear(); - } - -+static void llama_tensor_dequantize_impl( -+ ggml_tensor * tensor, std::vector> & output, std::vector & workers, -+ const size_t nelements, const int nthread -+) { -+ llama_tensor_dequantize_chunk_impl(tensor, output, workers, 0, nelements, nthread); -+} -+ - // - // do we allow this tensor to be quantized? - // -@@ -1208,27 +1220,18 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: - throw std::runtime_error(format("Missing importance matrix for tensor %s in a very low-bit quantization", tensor->name)); - } - -- float * f32_data; -- -- if (tensor->type == GGML_TYPE_F32) { -- f32_data = (float *) tensor->data; -- } else if (ggml_is_quantized(tensor->type) && !params->allow_requantize) { -- throw std::runtime_error(format("requantizing from type %s is disabled", ggml_type_name(tensor->type))); -- } else { -- llama_tensor_dequantize_impl(tensor, f32_conv_buf, workers, nelements, nthread); -- f32_data = (float *) f32_conv_buf.data(); -- } -- - LLAMA_LOG_INFO("converting to %s .. ", ggml_type_name(new_type)); - fflush(stdout); - -- if (work.size() < (size_t)nelements * 4) { -- work.resize(nelements * 4); // upper bound on size -- } -- new_data = work.data(); -- - const int64_t n_per_row = tensor->ne[0]; - const int64_t nrows = tensor->ne[1]; -+ const size_t row_size = ggml_row_size(new_type, n_per_row); -+ const size_t work_size = row_size * nrows * tensor->ne[2]; -+ -+ if (work.size() < work_size) { -+ work.resize(work_size); -+ } -+ new_data = work.data(); - - static const int64_t min_chunk_size = 32 * 512; - const int64_t chunk_size = (n_per_row >= min_chunk_size ? n_per_row : n_per_row * ((min_chunk_size + n_per_row - 1)/n_per_row)); -@@ -1236,15 +1239,51 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: - const int64_t nelements_matrix = tensor->ne[0] * tensor->ne[1]; - const int64_t nchunk = (nelements_matrix + chunk_size - 1)/chunk_size; - const int64_t nthread_use = nthread > 1 ? std::max((int64_t)1, std::min((int64_t)nthread, nchunk)) : 1; -+ const bool stream_source = tensor->type == GGML_TYPE_F16 || tensor->type == GGML_TYPE_BF16; - - // quantize each expert separately since they have different importance matrices - new_size = 0; -- for (int64_t i03 = 0; i03 < tensor->ne[2]; ++i03) { -- const float * f32_data_03 = f32_data + i03 * nelements_matrix; -- void * new_data_03 = (char *)new_data + ggml_row_size(new_type, n_per_row) * i03 * nrows; -- const float * imatrix_03 = imatrix ? imatrix + i03 * n_per_row : nullptr; -+ if (stream_source) { -+ static const int64_t stream_chunk_size = 32 * 1024 * 1024; -+ const int64_t rows_per_stream_chunk = std::max(1, std::max(chunk_size, stream_chunk_size) / n_per_row); - -- new_size += llama_tensor_quantize_impl(new_type, f32_data_03, new_data_03, chunk_size, nrows, n_per_row, imatrix_03, workers, nthread_use); -+ for (int64_t i03 = 0; i03 < tensor->ne[2]; ++i03) { -+ void * new_data_03 = (char *)new_data + row_size * i03 * nrows; -+ const float * imatrix_03 = imatrix ? imatrix + i03 * n_per_row : nullptr; -+ -+ for (int64_t first_row = 0; first_row < nrows; first_row += rows_per_stream_chunk) { -+ const int64_t this_nrow = std::min(nrows - first_row, rows_per_stream_chunk); -+ const int64_t start = i03 * nelements_matrix + first_row * n_per_row; -+ const size_t chunk_nelements = this_nrow * n_per_row; -+ -+ llama_tensor_dequantize_chunk_impl(tensor, f32_conv_buf, workers, start, chunk_nelements, nthread); -+ -+ void * new_data_chunk = (char *) new_data_03 + row_size * first_row; -+ const int64_t chunk_nchunk = (chunk_nelements + chunk_size - 1) / chunk_size; -+ const int64_t chunk_nthread = nthread > 1 ? std::max((int64_t)1, std::min((int64_t)nthread, chunk_nchunk)) : 1; -+ -+ new_size += llama_tensor_quantize_impl(new_type, (float *) f32_conv_buf.data(), new_data_chunk, chunk_size, this_nrow, n_per_row, imatrix_03, workers, chunk_nthread); -+ } -+ } -+ } else { -+ float * f32_data; -+ -+ if (tensor->type == GGML_TYPE_F32) { -+ f32_data = (float *) tensor->data; -+ } else if (ggml_is_quantized(tensor->type) && !params->allow_requantize) { -+ throw std::runtime_error(format("requantizing from type %s is disabled", ggml_type_name(tensor->type))); -+ } else { -+ llama_tensor_dequantize_impl(tensor, f32_conv_buf, workers, nelements, nthread); -+ f32_data = (float *) f32_conv_buf.data(); -+ } -+ -+ for (int64_t i03 = 0; i03 < tensor->ne[2]; ++i03) { -+ const float * f32_data_03 = f32_data + i03 * nelements_matrix; -+ void * new_data_03 = (char *)new_data + row_size * i03 * nrows; -+ const float * imatrix_03 = imatrix ? imatrix + i03 * n_per_row : nullptr; -+ -+ new_size += llama_tensor_quantize_impl(new_type, f32_data_03, new_data_03, chunk_size, nrows, n_per_row, imatrix_03, workers, nthread_use); -+ } - } - LLAMA_LOG_INFO("size = %8.2f MiB -> %8.2f MiB\n", tensor_size/1024.0/1024.0, new_size/1024.0/1024.0); - } --- -2.54.0 (Apple Git-156) - diff --git a/third_party/llama.cpp/patches/0091-Release-mmap-tensor-pages-after-quantization.patch b/third_party/llama.cpp/patches/0091-Release-mmap-tensor-pages-after-quantization.patch deleted file mode 100644 index 24b328fcd1..0000000000 --- a/third_party/llama.cpp/patches/0091-Release-mmap-tensor-pages-after-quantization.patch +++ /dev/null @@ -1,59 +0,0 @@ -From 460c8000aeae571b256901a53c9b53ac6d473faf Mon Sep 17 00:00:00 2001 -From: James Dumay -Date: Mon, 15 Jun 2026 09:25:06 +1000 -Subject: [PATCH 02/10] Release mmap tensor pages after quantization - ---- - src/llama-model-loader.cpp | 10 ++++++++++ - src/llama-model-loader.h | 1 + - src/llama-quant.cpp | 1 + - 3 files changed, 12 insertions(+) - -diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp -index 0d1cf3cc3..feb9d31c8 100644 ---- a/src/llama-model-loader.cpp -+++ b/src/llama-model-loader.cpp -@@ -1404,6 +1404,16 @@ void llama_model_loader::load_data_for(struct ggml_tensor * cur) const { - } - } - -+void llama_model_loader::unmap_data_for(struct ggml_tensor * cur) const { -+ if (!use_mmap) { -+ return; -+ } -+ -+ const auto & w = require_weight(ggml_get_name(cur)); -+ auto & mapping = mappings.at(w.idx); -+ mapping->unmap_fragment(w.offs, w.offs + ggml_nbytes(cur)); -+} -+ - bool llama_model_loader::load_all_data( - struct ggml_context * ctx, - llama_buf_map & bufs, -diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h -index c476026d3..39e20ff9a 100644 ---- a/src/llama-model-loader.h -+++ b/src/llama-model-loader.h -@@ -192,6 +192,7 @@ struct llama_model_loader { - - // for backwards compatibility, does not support ggml-backend - void load_data_for(struct ggml_tensor * cur) const; -+ void unmap_data_for(struct ggml_tensor * cur) const; - - // Returns false if cancelled by progress_callback - bool load_all_data( -diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp -index d77866392..3899ca896 100644 ---- a/src/llama-quant.cpp -+++ b/src/llama-quant.cpp -@@ -1298,6 +1298,7 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: - // write tensor data + padding - fout.write((const char *) new_data, new_size); - zeros(fout, GGML_PAD(new_size, align) - new_size); -+ ml.unmap_data_for(tensor); - } // no --dry-run - } // main loop - --- -2.54.0 (Apple Git-156) - diff --git a/third_party/llama.cpp/patches/0092-Extract-tensor-float-conversion-helper.patch b/third_party/llama.cpp/patches/0092-Extract-tensor-float-conversion-helper.patch deleted file mode 100644 index e478290aa6..0000000000 --- a/third_party/llama.cpp/patches/0092-Extract-tensor-float-conversion-helper.patch +++ /dev/null @@ -1,69 +0,0 @@ -From a391159ab6bef9eb0f14dbf027f9d761a18da445 Mon Sep 17 00:00:00 2001 -From: James Dumay -Date: Mon, 15 Jun 2026 10:21:56 +1000 -Subject: [PATCH 03/10] Extract tensor float conversion helper - ---- - src/llama-quant.cpp | 32 +++++++++++++++----------------- - 1 file changed, 15 insertions(+), 17 deletions(-) - -diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp -index 3899ca896..5b14869b6 100644 ---- a/src/llama-quant.cpp -+++ b/src/llama-quant.cpp -@@ -209,6 +209,18 @@ struct tensor_metadata { - // dequantization - // - -+static void llama_tensor_to_float_chunk( -+ ggml_type type, const ggml_type_traits * qtype, const void * input, float * output, int64_t nelements -+) { -+ if (type == GGML_TYPE_F16) { -+ ggml_fp16_to_fp32_row((const ggml_fp16_t *) input, output, nelements); -+ } else if (type == GGML_TYPE_BF16) { -+ ggml_bf16_to_fp32_row((const ggml_bf16_t *) input, output, nelements); -+ } else { -+ qtype->to_float(input, output, nelements); -+ } -+} -+ - static void llama_tensor_dequantize_chunk_impl( - ggml_tensor * tensor, std::vector> & output, std::vector & workers, - const int64_t start, const size_t nelements, const int nthread -@@ -245,15 +257,7 @@ static void llama_tensor_dequantize_chunk_impl( - uint8_t * input_data = (uint8_t *) tensor->data + input_offset; - - if (nthread < 2) { -- if (tensor->type == GGML_TYPE_F16) { -- ggml_fp16_to_fp32_row((ggml_fp16_t *) input_data, f32_output, nelements); -- } else if (tensor->type == GGML_TYPE_BF16) { -- ggml_bf16_to_fp32_row((ggml_bf16_t *) input_data, f32_output, nelements); -- } else if (ggml_is_quantized(tensor->type)) { -- qtype->to_float(input_data, f32_output, nelements); -- } else { -- GGML_ABORT("fatal error"); // unreachable -- } -+ llama_tensor_to_float_chunk(tensor->type, qtype, input_data, f32_output, nelements); - return; - } - -@@ -269,14 +273,8 @@ static void llama_tensor_dequantize_chunk_impl( - size_t thr_elems = thr_blocks * block_size; // number of elements for this thread - size_t thr_block_bytes = thr_blocks * block_size_bytes; // number of input bytes for this thread - -- auto compute = [qtype] (ggml_type typ, uint8_t * inbuf, float * outbuf, int64_t nels) { -- if (typ == GGML_TYPE_F16) { -- ggml_fp16_to_fp32_row((ggml_fp16_t *)inbuf, outbuf, nels); -- } else if (typ == GGML_TYPE_BF16) { -- ggml_bf16_to_fp32_row((ggml_bf16_t *)inbuf, outbuf, nels); -- } else { -- qtype->to_float(inbuf, outbuf, nels); -- } -+ auto compute = [qtype] (ggml_type typ, const uint8_t * inbuf, float * outbuf, int64_t nels) { -+ llama_tensor_to_float_chunk(typ, qtype, inbuf, outbuf, nels); - }; - workers.emplace_back(compute, tensor->type, input_data + in_buff_offs, f32_output + out_buff_offs, thr_elems); - in_buff_offs += thr_block_bytes; --- -2.54.0 (Apple Git-156) - diff --git a/third_party/llama.cpp/patches/0093-Use-type-traits-for-tensor-float-conversion.patch b/third_party/llama.cpp/patches/0093-Use-type-traits-for-tensor-float-conversion.patch deleted file mode 100644 index 36780ae3bc..0000000000 --- a/third_party/llama.cpp/patches/0093-Use-type-traits-for-tensor-float-conversion.patch +++ /dev/null @@ -1,84 +0,0 @@ -From 3e3387508bc10e9953e849b025122aee7d501ce5 Mon Sep 17 00:00:00 2001 -From: James Dumay -Date: Mon, 15 Jun 2026 10:26:29 +1000 -Subject: [PATCH 04/10] Use type traits for tensor float conversion - ---- - src/llama-quant.cpp | 37 +++++++++---------------------------- - 1 file changed, 9 insertions(+), 28 deletions(-) - -diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp -index 5b14869b6..847292290 100644 ---- a/src/llama-quant.cpp -+++ b/src/llama-quant.cpp -@@ -209,16 +209,8 @@ struct tensor_metadata { - // dequantization - // - --static void llama_tensor_to_float_chunk( -- ggml_type type, const ggml_type_traits * qtype, const void * input, float * output, int64_t nelements --) { -- if (type == GGML_TYPE_F16) { -- ggml_fp16_to_fp32_row((const ggml_fp16_t *) input, output, nelements); -- } else if (type == GGML_TYPE_BF16) { -- ggml_bf16_to_fp32_row((const ggml_bf16_t *) input, output, nelements); -- } else { -- qtype->to_float(input, output, nelements); -- } -+static void llama_tensor_to_float_chunk(const ggml_type_traits * qtype, const void * input, float * output, int64_t nelements) { -+ qtype->to_float(input, output, nelements); - } - - static void llama_tensor_dequantize_chunk_impl( -@@ -231,22 +223,11 @@ static void llama_tensor_dequantize_chunk_impl( - float * f32_output = (float *) output.data(); - - const ggml_type_traits * qtype = ggml_get_type_traits(tensor->type); -- if (ggml_is_quantized(tensor->type)) { -- if (qtype->to_float == NULL) { -- throw std::runtime_error(format("type %s unsupported for integer quantization: no dequantization available", ggml_type_name(tensor->type))); -- } -- } else if (tensor->type != GGML_TYPE_F16 && -- tensor->type != GGML_TYPE_BF16) { -- throw std::runtime_error(format("cannot dequantize/convert tensor type %s", ggml_type_name(tensor->type))); -+ if (qtype->to_float == NULL) { -+ throw std::runtime_error(format("type %s unsupported for conversion to float", ggml_type_name(tensor->type))); - } - -- size_t block_size; -- if (tensor->type == GGML_TYPE_F16 || -- tensor->type == GGML_TYPE_BF16) { -- block_size = 1; -- } else { -- block_size = (size_t)ggml_blck_size(tensor->type); -- } -+ const size_t block_size = (size_t)ggml_blck_size(tensor->type); - - GGML_ASSERT(start >= 0); - GGML_ASSERT((size_t) start % block_size == 0); -@@ -257,7 +238,7 @@ static void llama_tensor_dequantize_chunk_impl( - uint8_t * input_data = (uint8_t *) tensor->data + input_offset; - - if (nthread < 2) { -- llama_tensor_to_float_chunk(tensor->type, qtype, input_data, f32_output, nelements); -+ llama_tensor_to_float_chunk(qtype, input_data, f32_output, nelements); - return; - } - -@@ -273,10 +254,10 @@ static void llama_tensor_dequantize_chunk_impl( - size_t thr_elems = thr_blocks * block_size; // number of elements for this thread - size_t thr_block_bytes = thr_blocks * block_size_bytes; // number of input bytes for this thread - -- auto compute = [qtype] (ggml_type typ, const uint8_t * inbuf, float * outbuf, int64_t nels) { -- llama_tensor_to_float_chunk(typ, qtype, inbuf, outbuf, nels); -+ auto compute = [qtype] (const uint8_t * inbuf, float * outbuf, int64_t nels) { -+ llama_tensor_to_float_chunk(qtype, inbuf, outbuf, nels); - }; -- workers.emplace_back(compute, tensor->type, input_data + in_buff_offs, f32_output + out_buff_offs, thr_elems); -+ workers.emplace_back(compute, input_data + in_buff_offs, f32_output + out_buff_offs, thr_elems); - in_buff_offs += thr_block_bytes; - out_buff_offs += thr_elems; - } --- -2.54.0 (Apple Git-156) - diff --git a/third_party/llama.cpp/patches/0094-Call-tensor-to_float-directly.patch b/third_party/llama.cpp/patches/0094-Call-tensor-to_float-directly.patch deleted file mode 100644 index 92cb285305..0000000000 --- a/third_party/llama.cpp/patches/0094-Call-tensor-to_float-directly.patch +++ /dev/null @@ -1,45 +0,0 @@ -From 3ce3f68c946031dcbabacba30c4cf935c154bc42 Mon Sep 17 00:00:00 2001 -From: James Dumay -Date: Mon, 15 Jun 2026 10:29:13 +1000 -Subject: [PATCH 05/10] Call tensor to_float directly - ---- - src/llama-quant.cpp | 8 ++------ - 1 file changed, 2 insertions(+), 6 deletions(-) - -diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp -index 847292290..cf1e5bce3 100644 ---- a/src/llama-quant.cpp -+++ b/src/llama-quant.cpp -@@ -209,10 +209,6 @@ struct tensor_metadata { - // dequantization - // - --static void llama_tensor_to_float_chunk(const ggml_type_traits * qtype, const void * input, float * output, int64_t nelements) { -- qtype->to_float(input, output, nelements); --} -- - static void llama_tensor_dequantize_chunk_impl( - ggml_tensor * tensor, std::vector> & output, std::vector & workers, - const int64_t start, const size_t nelements, const int nthread -@@ -238,7 +234,7 @@ static void llama_tensor_dequantize_chunk_impl( - uint8_t * input_data = (uint8_t *) tensor->data + input_offset; - - if (nthread < 2) { -- llama_tensor_to_float_chunk(qtype, input_data, f32_output, nelements); -+ qtype->to_float(input_data, f32_output, nelements); - return; - } - -@@ -255,7 +251,7 @@ static void llama_tensor_dequantize_chunk_impl( - size_t thr_block_bytes = thr_blocks * block_size_bytes; // number of input bytes for this thread - - auto compute = [qtype] (const uint8_t * inbuf, float * outbuf, int64_t nels) { -- llama_tensor_to_float_chunk(qtype, inbuf, outbuf, nels); -+ qtype->to_float(inbuf, outbuf, nels); - }; - workers.emplace_back(compute, input_data + in_buff_offs, f32_output + out_buff_offs, thr_elems); - in_buff_offs += thr_block_bytes; --- -2.54.0 (Apple Git-156) - diff --git a/third_party/llama.cpp/patches/0095-Keep-tensor-chunk-offset-assertions-signed.patch b/third_party/llama.cpp/patches/0095-Keep-tensor-chunk-offset-assertions-signed.patch deleted file mode 100644 index 1564200f7a..0000000000 --- a/third_party/llama.cpp/patches/0095-Keep-tensor-chunk-offset-assertions-signed.patch +++ /dev/null @@ -1,55 +0,0 @@ -From 37d97e1156a0162d0164531fb2b5a6e312297114 Mon Sep 17 00:00:00 2001 -From: James Dumay -Date: Mon, 15 Jun 2026 10:32:16 +1000 -Subject: [PATCH 06/10] Keep tensor chunk offset assertions signed - ---- - src/llama-quant.cpp | 14 ++++++++------ - 1 file changed, 8 insertions(+), 6 deletions(-) - -diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp -index cf1e5bce3..86089853a 100644 ---- a/src/llama-quant.cpp -+++ b/src/llama-quant.cpp -@@ -223,14 +223,16 @@ static void llama_tensor_dequantize_chunk_impl( - throw std::runtime_error(format("type %s unsupported for conversion to float", ggml_type_name(tensor->type))); - } - -- const size_t block_size = (size_t)ggml_blck_size(tensor->type); -+ const int64_t block_size = ggml_blck_size(tensor->type); - -+ GGML_ASSERT(block_size > 0); - GGML_ASSERT(start >= 0); -- GGML_ASSERT((size_t) start % block_size == 0); -- GGML_ASSERT(nelements % block_size == 0); -+ GGML_ASSERT(start % block_size == 0); - - size_t block_size_bytes = ggml_type_size(tensor->type); -- const size_t input_offset = ((size_t) start / block_size) * block_size_bytes; -+ const size_t block_size_size = (size_t) block_size; -+ GGML_ASSERT(nelements % block_size_size == 0); -+ const size_t input_offset = ((size_t) start / block_size_size) * block_size_bytes; - uint8_t * input_data = (uint8_t *) tensor->data + input_offset; - - if (nthread < 2) { -@@ -238,7 +240,7 @@ static void llama_tensor_dequantize_chunk_impl( - return; - } - -- size_t nblocks = nelements / block_size; -+ size_t nblocks = nelements / block_size_size; - size_t blocks_per_thread = nblocks / nthread; - size_t spare_blocks = nblocks - (blocks_per_thread * nthread); // if blocks aren't divisible by thread count - -@@ -247,7 +249,7 @@ static void llama_tensor_dequantize_chunk_impl( - - for (int tnum = 0; tnum < nthread; tnum++) { - size_t thr_blocks = blocks_per_thread + (tnum == nthread - 1 ? spare_blocks : 0); // num blocks for this thread -- size_t thr_elems = thr_blocks * block_size; // number of elements for this thread -+ size_t thr_elems = thr_blocks * block_size_size; // number of elements for this thread - size_t thr_block_bytes = thr_blocks * block_size_bytes; // number of input bytes for this thread - - auto compute = [qtype] (const uint8_t * inbuf, float * outbuf, int64_t nels) { --- -2.54.0 (Apple Git-156) - diff --git a/third_party/llama.cpp/patches/0096-Simplify-stream-chunk-row-calculation.patch b/third_party/llama.cpp/patches/0096-Simplify-stream-chunk-row-calculation.patch deleted file mode 100644 index dab15bfe51..0000000000 --- a/third_party/llama.cpp/patches/0096-Simplify-stream-chunk-row-calculation.patch +++ /dev/null @@ -1,25 +0,0 @@ -From f482b044622aee4b1ef1fe7c1a09c60adaf3a6c2 Mon Sep 17 00:00:00 2001 -From: James Dumay -Date: Mon, 15 Jun 2026 10:48:40 +1000 -Subject: [PATCH 07/10] Simplify stream chunk row calculation - ---- - src/llama-quant.cpp | 2 +- - 1 file changed, 1 insertion(+), 1 deletion(-) - -diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp -index 86089853a..fd695b388 100644 ---- a/src/llama-quant.cpp -+++ b/src/llama-quant.cpp -@@ -1222,7 +1222,7 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: - new_size = 0; - if (stream_source) { - static const int64_t stream_chunk_size = 32 * 1024 * 1024; -- const int64_t rows_per_stream_chunk = std::max(1, std::max(chunk_size, stream_chunk_size) / n_per_row); -+ const int64_t rows_per_stream_chunk = std::max(1, stream_chunk_size / n_per_row); - - for (int64_t i03 = 0; i03 < tensor->ne[2]; ++i03) { - void * new_data_03 = (char *)new_data + row_size * i03 * nrows; --- -2.54.0 (Apple Git-156) - diff --git a/third_party/llama.cpp/patches/0097-Add-llama-quantize-progress-watchdog.patch b/third_party/llama.cpp/patches/0097-Add-llama-quantize-progress-watchdog.patch deleted file mode 100644 index cfa97ce735..0000000000 --- a/third_party/llama.cpp/patches/0097-Add-llama-quantize-progress-watchdog.patch +++ /dev/null @@ -1,527 +0,0 @@ -From 5890fca0b695bf926fef701142ba2279ea03eef9 Mon Sep 17 00:00:00 2001 -From: James Dumay -Date: Mon, 15 Jun 2026 14:34:03 +1000 -Subject: [PATCH 08/10] Add llama-quantize progress watchdog - ---- - src/llama-quant.cpp | 414 ++++++++++++++++++++++++++++++++++++++++++-- - 1 file changed, 404 insertions(+), 10 deletions(-) - -diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp -index fd695b388..601f9dea2 100644 ---- a/src/llama-quant.cpp -+++ b/src/llama-quant.cpp -@@ -4,14 +4,24 @@ - #include "llama-ext.h" - - #include -+#include -+#include - #include -+#include - #include - #include -+#include - #include - #include - #include - #include - #include -+#include -+ -+#if defined(__APPLE__) -+#include -+#include -+#endif - - // result of parsing --tensor-type option - // (changes to this struct must be reflected in tools/quantize/quantize.cpp) -@@ -44,6 +54,206 @@ static void zeros(std::ofstream & file, size_t n) { - } - } - -+struct quantize_process_memory { -+ int64_t rss_kib = -1; -+ int64_t hwm_kib = -1; -+}; -+ -+static int64_t llama_quantize_parse_kib_line(const std::string & line, const char * prefix) { -+ if (line.rfind(prefix, 0) != 0) { -+ return -1; -+ } -+ -+ long long value = -1; -+ if (std::sscanf(line.c_str() + std::strlen(prefix), "%lld", &value) != 1) { -+ return -1; -+ } -+ return value; -+} -+ -+static quantize_process_memory llama_quantize_process_memory_usage() { -+ quantize_process_memory result; -+ -+#if defined(__linux__) -+ std::ifstream status("/proc/self/status"); -+ std::string line; -+ while (std::getline(status, line)) { -+ const int64_t rss = llama_quantize_parse_kib_line(line, "VmRSS:"); -+ if (rss >= 0) { -+ result.rss_kib = rss; -+ continue; -+ } -+ -+ const int64_t hwm = llama_quantize_parse_kib_line(line, "VmHWM:"); -+ if (hwm >= 0) { -+ result.hwm_kib = hwm; -+ } -+ } -+#elif defined(__APPLE__) -+ mach_task_basic_info info; -+ mach_msg_type_number_t count = MACH_TASK_BASIC_INFO_COUNT; -+ if (task_info(mach_task_self(), MACH_TASK_BASIC_INFO, (task_info_t) &info, &count) == KERN_SUCCESS) { -+ result.rss_kib = (int64_t) info.resident_size / 1024; -+ } -+ -+ struct rusage usage; -+ if (getrusage(RUSAGE_SELF, &usage) == 0) { -+ result.hwm_kib = (int64_t) usage.ru_maxrss / 1024; -+ } -+#endif -+ -+ return result; -+} -+ -+static int64_t llama_quantize_watchdog_interval_seconds() { -+ const char * env = std::getenv("LLAMA_QUANTIZE_WATCHDOG_SECONDS"); -+ if (env == nullptr || env[0] == '\0') { -+ return 180; -+ } -+ -+ char * end = nullptr; -+ const long long value = std::strtoll(env, &end, 10); -+ if (end == env) { -+ return 180; -+ } -+ return std::max(0, value); -+} -+ -+struct quantize_watchdog_snapshot { -+ std::string tensor_name; -+ std::string tensor_shape; -+ std::string phase; -+ ggml_type source_type = GGML_TYPE_COUNT; -+ ggml_type target_type = GGML_TYPE_COUNT; -+ size_t tensor_index = 0; -+ size_t tensor_count = 0; -+ int split_index = 0; -+ int split_count = 1; -+ size_t tensor_size = 0; -+ size_t output_size = 0; -+ int64_t expert_index = -1; -+ int64_t expert_count = -1; -+ int64_t first_row = -1; -+ int64_t rows = -1; -+ int64_t total_rows = -1; -+ bool quantize = false; -+ bool streaming = false; -+}; -+ -+class quantize_watchdog { -+public: -+ quantize_watchdog() : interval_seconds(llama_quantize_watchdog_interval_seconds()) { -+ if (interval_seconds <= 0) { -+ return; -+ } -+ -+ started_at = std::chrono::steady_clock::now(); -+ worker = std::thread([this]() { this->run(); }); -+ } -+ -+ ~quantize_watchdog() { -+ stop(); -+ } -+ -+ quantize_watchdog(const quantize_watchdog &) = delete; -+ quantize_watchdog & operator=(const quantize_watchdog &) = delete; -+ -+ void update(quantize_watchdog_snapshot next) { -+ if (interval_seconds <= 0) { -+ return; -+ } -+ -+ std::lock_guard lock(mutex); -+ current = std::move(next); -+ } -+ -+ void stop() { -+ if (interval_seconds <= 0) { -+ return; -+ } -+ -+ { -+ std::lock_guard lock(mutex); -+ stopping = true; -+ } -+ cv.notify_all(); -+ -+ if (worker.joinable()) { -+ worker.join(); -+ } -+ -+ interval_seconds = 0; -+ } -+ -+private: -+ void run() { -+ std::unique_lock lock(mutex); -+ while (!stopping) { -+ if (cv.wait_for(lock, std::chrono::seconds(interval_seconds), [this]() { return stopping; })) { -+ break; -+ } -+ -+ const auto snapshot = current; -+ const auto elapsed = std::chrono::duration_cast( -+ std::chrono::steady_clock::now() - started_at); -+ lock.unlock(); -+ log(snapshot, elapsed.count()); -+ lock.lock(); -+ } -+ } -+ -+ static void log(const quantize_watchdog_snapshot & snapshot, int64_t elapsed_seconds) { -+ if (snapshot.tensor_name.empty()) { -+ return; -+ } -+ -+ const quantize_process_memory mem = llama_quantize_process_memory_usage(); -+ const double rss_mib = mem.rss_kib >= 0 ? mem.rss_kib / 1024.0 : -1.0; -+ const double hwm_mib = mem.hwm_kib >= 0 ? mem.hwm_kib / 1024.0 : -1.0; -+ -+ LLAMA_LOG_INFO( -+ "\nllama_model_quantize_impl: watchdog elapsed=%" PRId64 "s phase=%s tensor=%zu/%zu split=%d/%d name=%s shape=[%s] type=%s->%s quantize=%s stream=%s tensor=%.2f MiB output=%.2f MiB rss=%.2f MiB hwm=%.2f MiB", -+ elapsed_seconds, -+ snapshot.phase.c_str(), -+ snapshot.tensor_index, -+ snapshot.tensor_count, -+ snapshot.split_index + 1, -+ snapshot.split_count, -+ snapshot.tensor_name.c_str(), -+ snapshot.tensor_shape.c_str(), -+ ggml_type_name(snapshot.source_type), -+ ggml_type_name(snapshot.target_type), -+ snapshot.quantize ? "yes" : "no", -+ snapshot.streaming ? "yes" : "no", -+ snapshot.tensor_size / 1024.0 / 1024.0, -+ snapshot.output_size / 1024.0 / 1024.0, -+ rss_mib, -+ hwm_mib); -+ -+ if (snapshot.expert_index >= 0) { -+ LLAMA_LOG_INFO( -+ " expert=%" PRId64 "/%" PRId64 " rows=%" PRId64 "+%" PRId64 "/%" PRId64, -+ snapshot.expert_index + 1, -+ snapshot.expert_count, -+ snapshot.first_row, -+ snapshot.rows, -+ snapshot.total_rows); -+ } -+ -+ LLAMA_LOG_INFO("\n"); -+ std::fflush(stdout); -+ std::fflush(stderr); -+ } -+ -+ std::mutex mutex; -+ std::condition_variable cv; -+ std::thread worker; -+ std::chrono::steady_clock::time_point started_at; -+ quantize_watchdog_snapshot current; -+ int64_t interval_seconds = 0; -+ bool stopping = false; -+}; -+ - static std::string remap_layer(const std::string & orig_name, const std::vector & prune, std::map & mapped, int & next_id) { - if (prune.empty()) { - return orig_name; -@@ -1097,6 +1307,8 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: - new_ofstream(0); - } - -+ quantize_watchdog watchdog; -+ - // - // main loop: iterate over all weights - // -@@ -1105,6 +1317,11 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: - const auto & weight = *tensors[i]; - const auto & tm = metadata[i]; - ggml_tensor * tensor = weight.tensor; -+ const size_t tensor_index = i + 1; -+ const std::string tensor_name = ggml_get_name(tensor); -+ const std::string tensor_shape = llama_format_tensor_shape(tensor); -+ const ggml_type cur_type = tensor->type; -+ const ggml_type new_type = tm.target_type; - - if (!params->dry_run && (weight.idx != cur_split && params->keep_split)) { - close_ofstream(); -@@ -1112,6 +1329,28 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: - } - - const size_t tensor_size = ggml_nbytes(tensor); -+ const bool quantize = cur_type != new_type; -+ -+ watchdog.update({ -+ tensor_name, -+ tensor_shape, -+ "loading", -+ cur_type, -+ new_type, -+ tensor_index, -+ tensors.size(), -+ params->keep_split ? weight.idx : 0, -+ n_split, -+ tensor_size, -+ 0, -+ -1, -+ -1, -+ -1, -+ -1, -+ -1, -+ quantize, -+ false, -+ }); - - if (!params->dry_run) { - if (!ml.use_mmap) { -@@ -1123,19 +1362,34 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: - ml.load_data_for(tensor); - } - -+ ++idx; -+ watchdog.update({ -+ tensor_name, -+ tensor_shape, -+ params->dry_run ? "dry-run" : "loaded", -+ cur_type, -+ new_type, -+ tensor_index, -+ tensors.size(), -+ params->keep_split ? weight.idx : 0, -+ n_split, -+ tensor_size, -+ 0, -+ -1, -+ -1, -+ -1, -+ -1, -+ -1, -+ quantize, -+ false, -+ }); -+ - LLAMA_LOG_INFO("[%4d/%4d] %-36s - [%s], type = %6s, ", -- ++idx, ml.n_tensors, -- ggml_get_name(tensor), -- llama_format_tensor_shape(tensor).c_str(), -+ idx, ml.n_tensors, -+ tensor_name.c_str(), -+ tensor_shape.c_str(), - ggml_type_name(tensor->type)); - -- const ggml_type cur_type = tensor->type; -- const ggml_type new_type = tm.target_type; -- -- // If we've decided to quantize to the same type the tensor is already -- // in then there's nothing to do. -- bool quantize = cur_type != new_type; -- - void * new_data; - size_t new_size; - -@@ -1156,12 +1410,52 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: - } - total_size_org += tensor_size; - total_size_new += new_size; -+ watchdog.update({ -+ tensor_name, -+ tensor_shape, -+ "done", -+ cur_type, -+ new_type, -+ tensor_index, -+ tensors.size(), -+ params->keep_split ? weight.idx : 0, -+ n_split, -+ tensor_size, -+ new_size, -+ -1, -+ -1, -+ -1, -+ -1, -+ -1, -+ quantize, -+ false, -+ }); - continue; - } else { - // no --dry-run, perform quantization - if (!quantize) { - new_data = tensor->data; - new_size = tensor_size; -+ watchdog.update({ -+ tensor_name, -+ tensor_shape, -+ "copying", -+ cur_type, -+ new_type, -+ tensor_index, -+ tensors.size(), -+ params->keep_split ? weight.idx : 0, -+ n_split, -+ tensor_size, -+ new_size, -+ -1, -+ -1, -+ -1, -+ -1, -+ -1, -+ quantize, -+ false, -+ }); - LLAMA_LOG_INFO("size = %8.3f MiB\n", tensor_size/1024.0/1024.0); - } else { - const int64_t nelements = ggml_nelements(tensor); -@@ -1232,6 +1526,26 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: - const int64_t this_nrow = std::min(nrows - first_row, rows_per_stream_chunk); - const int64_t start = i03 * nelements_matrix + first_row * n_per_row; - const size_t chunk_nelements = this_nrow * n_per_row; -+ watchdog.update({ -+ tensor_name, -+ tensor_shape, -+ "streaming", -+ cur_type, -+ new_type, -+ tensor_index, -+ tensors.size(), -+ params->keep_split ? weight.idx : 0, -+ n_split, -+ tensor_size, -+ new_size, -+ i03, -+ tensor->ne[2], -+ first_row, -+ this_nrow, -+ nrows, -+ quantize, -+ true, -+ }); - - llama_tensor_dequantize_chunk_impl(tensor, f32_conv_buf, workers, start, chunk_nelements, nthread); - -@@ -1250,6 +1564,26 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: - } else if (ggml_is_quantized(tensor->type) && !params->allow_requantize) { - throw std::runtime_error(format("requantizing from type %s is disabled", ggml_type_name(tensor->type))); - } else { -+ watchdog.update({ -+ tensor_name, -+ tensor_shape, -+ "dequantizing", -+ cur_type, -+ new_type, -+ tensor_index, -+ tensors.size(), -+ params->keep_split ? weight.idx : 0, -+ n_split, -+ tensor_size, -+ new_size, -+ -1, -+ -1, -+ -1, -+ -1, -+ -1, -+ quantize, -+ false, -+ }); - llama_tensor_dequantize_impl(tensor, f32_conv_buf, workers, nelements, nthread); - f32_data = (float *) f32_conv_buf.data(); - } -@@ -1258,6 +1592,26 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: - const float * f32_data_03 = f32_data + i03 * nelements_matrix; - void * new_data_03 = (char *)new_data + row_size * i03 * nrows; - const float * imatrix_03 = imatrix ? imatrix + i03 * n_per_row : nullptr; -+ watchdog.update({ -+ tensor_name, -+ tensor_shape, -+ "quantizing", -+ cur_type, -+ new_type, -+ tensor_index, -+ tensors.size(), -+ params->keep_split ? weight.idx : 0, -+ n_split, -+ tensor_size, -+ new_size, -+ i03, -+ tensor->ne[2], -+ 0, -+ nrows, -+ nrows, -+ quantize, -+ false, -+ }); - - new_size += llama_tensor_quantize_impl(new_type, f32_data_03, new_data_03, chunk_size, nrows, n_per_row, imatrix_03, workers, nthread_use); - } -@@ -1273,9 +1627,49 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: - gguf_set_tensor_data(ctx_outs[cur_split].get(), metadata[i].name.c_str(), new_data); - - // write tensor data + padding -+ watchdog.update({ -+ tensor_name, -+ tensor_shape, -+ "writing", -+ cur_type, -+ new_type, -+ tensor_index, -+ tensors.size(), -+ params->keep_split ? weight.idx : 0, -+ n_split, -+ tensor_size, -+ new_size, -+ -1, -+ -1, -+ -1, -+ -1, -+ -1, -+ quantize, -+ false, -+ }); - fout.write((const char *) new_data, new_size); - zeros(fout, GGML_PAD(new_size, align) - new_size); - ml.unmap_data_for(tensor); -+ watchdog.update({ -+ tensor_name, -+ tensor_shape, -+ "done", -+ cur_type, -+ new_type, -+ tensor_index, -+ tensors.size(), -+ params->keep_split ? weight.idx : 0, -+ n_split, -+ tensor_size, -+ new_size, -+ -1, -+ -1, -+ -1, -+ -1, -+ -1, -+ quantize, -+ false, -+ }); - } // no --dry-run - } // main loop - --- -2.54.0 (Apple Git-156) - diff --git a/third_party/llama.cpp/patches/0098-Drop-quantized-output-cache-after-tensor-writes.patch b/third_party/llama.cpp/patches/0098-Drop-quantized-output-cache-after-tensor-writes.patch deleted file mode 100644 index bc69cff7a2..0000000000 --- a/third_party/llama.cpp/patches/0098-Drop-quantized-output-cache-after-tensor-writes.patch +++ /dev/null @@ -1,170 +0,0 @@ -From e58fbeb30279cbb8bd9e2069f765438368ac4f36 Mon Sep 17 00:00:00 2001 -From: James Dumay -Date: Mon, 15 Jun 2026 15:09:14 +1000 -Subject: [PATCH 09/10] Drop quantized output cache after tensor writes - ---- - src/llama-quant.cpp | 105 +++++++++++++++++++++++++++++++++++++++++++- - 1 file changed, 104 insertions(+), 1 deletion(-) - -diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp -index 601f9dea2..46b174b35 100644 ---- a/src/llama-quant.cpp -+++ b/src/llama-quant.cpp -@@ -4,6 +4,7 @@ - #include "llama-ext.h" - - #include -+#include - #include - #include - #include -@@ -18,6 +19,11 @@ - #include - #include - -+#if defined(__linux__) -+#include -+#include -+#endif -+ - #if defined(__APPLE__) - #include - #include -@@ -105,6 +111,37 @@ static quantize_process_memory llama_quantize_process_memory_usage() { - return result; - } - -+static void llama_quantize_drop_output_cache(const std::string & path, std::streamoff offset, size_t size) { -+#if defined(__linux__) -+ if (path.empty() || offset < 0 || size == 0) { -+ return; -+ } -+ -+ const int fd = open(path.c_str(), O_RDWR | O_CLOEXEC); -+ if (fd < 0) { -+ LLAMA_LOG_WARN("%s: failed to open %s for output cache drop: %s\n", __func__, path.c_str(), std::strerror(errno)); -+ return; -+ } -+ -+ const off_t start = static_cast(offset); -+ const off_t len = static_cast(size); -+ if (fdatasync(fd) != 0) { -+ LLAMA_LOG_WARN("%s: fdatasync failed for %s: %s\n", __func__, path.c_str(), std::strerror(errno)); -+ } -+ -+ const int rc = posix_fadvise(fd, start, len, POSIX_FADV_DONTNEED); -+ if (rc != 0) { -+ LLAMA_LOG_WARN("%s: posix_fadvise DONTNEED failed for %s: %s\n", __func__, path.c_str(), std::strerror(rc)); -+ } -+ -+ close(fd); -+#else -+ (void) path; -+ (void) offset; -+ (void) size; -+#endif -+} -+ - static int64_t llama_quantize_watchdog_interval_seconds() { - const char * env = std::getenv("LLAMA_QUANTIZE_WATCHDOG_SECONDS"); - if (env == nullptr || env[0] == '\0') { -@@ -1275,6 +1312,7 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: - - int cur_split = -1; - std::ofstream fout; -+ std::string current_output_path; - auto close_ofstream = [&]() { - // Write metadata and close file handler - if (fout.is_open()) { -@@ -1295,6 +1333,7 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: - fname = std::string(split_path.data()); - } - -+ current_output_path = fname; - fout = std::ofstream(fname, std::ios::binary); - fout.exceptions(std::ofstream::failbit); // fail fast on write errors - const size_t meta_size = gguf_get_meta_size(ctx_outs[cur_split].get()); -@@ -1627,6 +1666,8 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: - gguf_set_tensor_data(ctx_outs[cur_split].get(), metadata[i].name.c_str(), new_data); - - // write tensor data + padding -+ const std::streamoff output_offset = fout.tellp(); -+ const size_t output_padding = GGML_PAD(new_size, align) - new_size; - watchdog.update({ - tensor_name, - tensor_shape, -@@ -1648,7 +1689,69 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: - false, - }); - fout.write((const char *) new_data, new_size); -- zeros(fout, GGML_PAD(new_size, align) - new_size); -+ zeros(fout, output_padding); -+ watchdog.update({ -+ tensor_name, -+ tensor_shape, -+ "flushing-output", -+ cur_type, -+ new_type, -+ tensor_index, -+ tensors.size(), -+ params->keep_split ? weight.idx : 0, -+ n_split, -+ tensor_size, -+ new_size, -+ -1, -+ -1, -+ -1, -+ -1, -+ -1, -+ quantize, -+ false, -+ }); -+ fout.flush(); -+ watchdog.update({ -+ tensor_name, -+ tensor_shape, -+ "dropping-output-cache", -+ cur_type, -+ new_type, -+ tensor_index, -+ tensors.size(), -+ params->keep_split ? weight.idx : 0, -+ n_split, -+ tensor_size, -+ new_size, -+ -1, -+ -1, -+ -1, -+ -1, -+ -1, -+ quantize, -+ false, -+ }); -+ llama_quantize_drop_output_cache(current_output_path, output_offset, new_size + output_padding); -+ watchdog.update({ -+ tensor_name, -+ tensor_shape, -+ "unmapping-source", -+ cur_type, -+ new_type, -+ tensor_index, -+ tensors.size(), -+ params->keep_split ? weight.idx : 0, -+ n_split, -+ tensor_size, -+ new_size, -+ -1, -+ -1, -+ -1, -+ -1, -+ -1, -+ quantize, -+ false, -+ }); - ml.unmap_data_for(tensor); - watchdog.update({ - tensor_name, --- -2.54.0 (Apple Git-156) - diff --git a/third_party/llama.cpp/patches/0099-Quantize-split-GGUF-windows.patch b/third_party/llama.cpp/patches/0099-Quantize-split-GGUF-windows.patch deleted file mode 100644 index 921e530c0e..0000000000 --- a/third_party/llama.cpp/patches/0099-Quantize-split-GGUF-windows.patch +++ /dev/null @@ -1,162 +0,0 @@ -From 21e71f17838dfd5d84f471d567e35994cff6c45a Mon Sep 17 00:00:00 2001 -From: James Dumay -Date: Mon, 15 Jun 2026 17:03:17 +1000 -Subject: [PATCH 10/10] Quantize split GGUF windows - ---- - include/llama.h | 2 ++ - src/llama-quant.cpp | 40 +++++++++++++++++++++++++++++++++---- - tools/quantize/quantize.cpp | 18 ++++++++++++++++- - 3 files changed, 55 insertions(+), 5 deletions(-) - -diff --git a/include/llama.h b/include/llama.h -index 27e480674..8f893592b 100644 ---- a/include/llama.h -+++ b/include/llama.h -@@ -421,6 +421,8 @@ extern "C" { - const struct llama_model_kv_override * kv_overrides; // pointer to kv overrides - const struct llama_model_tensor_override * tt_overrides; // pointer to tensor overrides - const int32_t * prune_layers; // pointer to layer indices to prune -+ int32_t first_split; // first 1-based split to quantize when keep_split is enabled, <= 0 means first -+ int32_t last_split; // last 1-based split to quantize when keep_split is enabled, <= 0 means last - } llama_model_quantize_params; - - typedef struct llama_logit_bias { -diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp -index 46b174b35..220b396a2 100644 ---- a/src/llama-quant.cpp -+++ b/src/llama-quant.cpp -@@ -1244,6 +1244,22 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: - n_split = std::max(uint16_t(it->idx + 1), n_split); - } - } -+ -+ if (!params->keep_split && (params->first_split > 0 || params->last_split > 0)) { -+ throw std::runtime_error("--first-split and --last-split require --keep-split"); -+ } -+ -+ const uint16_t first_split = params->keep_split && params->first_split > 0 ? uint16_t(params->first_split - 1) : 0; -+ const uint16_t last_split = params->keep_split && params->last_split > 0 ? uint16_t(params->last_split - 1) : uint16_t(n_split - 1); -+ if (first_split > last_split || last_split >= n_split) { -+ throw std::runtime_error(format("invalid split window: first=%d last=%d split_count=%d", -+ params->first_split, params->last_split, n_split)); -+ } -+ -+ auto split_in_window = [&](uint16_t split_idx) { -+ return !params->keep_split || (split_idx >= first_split && split_idx <= last_split); -+ }; -+ - std::vector ctx_outs(n_split); - ctx_outs[0] = std::move(ctx_out); - -@@ -1259,6 +1275,11 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: - const struct ggml_tensor * tensor = it->tensor; - - uint16_t i_split = params->keep_split ? it->idx : 0; -+ if (!split_in_window(i_split)) { -+ metadata[i].target_type = tensor->type; -+ continue; -+ } -+ - if (!ctx_outs[i_split]) { - ctx_outs[i_split].reset(gguf_init_empty()); - } -@@ -1294,6 +1315,9 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: - // Set split info if needed - if (n_split > 1) { - for (size_t i = 0; i < ctx_outs.size(); ++i) { -+ if (!ctx_outs[i]) { -+ continue; -+ } - gguf_set_val_u16(ctx_outs[i].get(), ml.llm_kv(LLM_KV_SPLIT_NO).c_str(), i); - gguf_set_val_u16(ctx_outs[i].get(), ml.llm_kv(LLM_KV_SPLIT_COUNT).c_str(), n_split); - gguf_set_val_i32(ctx_outs[i].get(), ml.llm_kv(LLM_KV_SPLIT_TENSORS_COUNT).c_str(), (int32_t)tensors.size()); -@@ -1342,7 +1366,7 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: - }; - - // no output file for --dry-run -- if (!params->dry_run) { -+ if (!params->dry_run && !params->keep_split) { - new_ofstream(0); - } - -@@ -1356,14 +1380,20 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: - const auto & weight = *tensors[i]; - const auto & tm = metadata[i]; - ggml_tensor * tensor = weight.tensor; -+ if (!split_in_window(params->keep_split ? weight.idx : 0)) { -+ continue; -+ } -+ - const size_t tensor_index = i + 1; - const std::string tensor_name = ggml_get_name(tensor); - const std::string tensor_shape = llama_format_tensor_shape(tensor); - const ggml_type cur_type = tensor->type; - const ggml_type new_type = tm.target_type; - -- if (!params->dry_run && (weight.idx != cur_split && params->keep_split)) { -- close_ofstream(); -+ if (!params->dry_run && params->keep_split && weight.idx != cur_split) { -+ if (cur_split >= 0) { -+ close_ofstream(); -+ } - new_ofstream(weight.idx); - } - -@@ -1814,7 +1844,9 @@ llama_model_quantize_params llama_model_quantize_default_params() { - /*.imatrix =*/ nullptr, - /*.kv_overrides =*/ nullptr, - /*.tensor_type =*/ nullptr, -- /*.prune_layers =*/ nullptr -+ /*.prune_layers =*/ nullptr, -+ /*.first_split =*/ 0, -+ /*.last_split =*/ 0, - }; - - return result; -diff --git a/tools/quantize/quantize.cpp b/tools/quantize/quantize.cpp -index 840eefc2f..b9ef9c0ca 100644 ---- a/tools/quantize/quantize.cpp -+++ b/tools/quantize/quantize.cpp -@@ -121,7 +121,7 @@ static bool try_parse_ftype(const std::string & ftype_str_in, llama_ftype & ftyp - static void usage(const char * executable) { - printf("usage: %s [--help] [--allow-requantize] [--leave-output-tensor] [--pure] [--imatrix] [--include-weights]\n", executable); - printf(" [--exclude-weights] [--output-tensor-type] [--token-embedding-type] [--tensor-type] [--tensor-type-file]\n"); -- printf(" [--prune-layers] [--keep-split] [--override-kv] [--dry-run]\n"); -+ printf(" [--prune-layers] [--keep-split] [--first-split N] [--last-split N] [--override-kv] [--dry-run]\n"); - printf(" model-f32.gguf [model-quant.gguf] type [nthreads]\n\n"); - printf(" --allow-requantize\n"); - printf(" allow requantizing tensors that have already been quantized\n"); -@@ -155,6 +155,10 @@ static void usage(const char * executable) { - printf(" WARNING: this is an advanced option, use with care.\n"); - printf(" --keep-split\n"); - printf(" generate quantized model in the same shards as input\n"); -+ printf(" --first-split N\n"); -+ printf(" first 1-based split to quantize with --keep-split\n"); -+ printf(" --last-split N\n"); -+ printf(" last 1-based split to quantize with --keep-split\n"); - printf(" --override-kv KEY=TYPE:VALUE\n"); - printf(" override model metadata by key in the quantized model. may be specified multiple times.\n"); - printf(" WARNING: this is an advanced option, use with care.\n"); -@@ -466,6 +470,18 @@ int llama_quantize(int argc, char ** argv) { - } - } else if (strcmp(argv[arg_idx], "--keep-split") == 0) { - params.keep_split = true; -+ } else if (strcmp(argv[arg_idx], "--first-split") == 0) { -+ if (arg_idx < argc-1) { -+ params.first_split = std::stoi(argv[++arg_idx]); -+ } else { -+ usage(argv[0]); -+ } -+ } else if (strcmp(argv[arg_idx], "--last-split") == 0) { -+ if (arg_idx < argc-1) { -+ params.last_split = std::stoi(argv[++arg_idx]); -+ } else { -+ usage(argv[0]); -+ } - } else { - usage(argv[0]); - } --- -2.54.0 (Apple Git-156) - diff --git a/third_party/llama.cpp/patches/0100-Add-Skippy-batched-stage-execution-ABI.patch b/third_party/llama.cpp/patches/0100-Add-Skippy-batched-stage-execution-ABI.patch new file mode 100644 index 0000000000..fe2dce409f --- /dev/null +++ b/third_party/llama.cpp/patches/0100-Add-Skippy-batched-stage-execution-ABI.patch @@ -0,0 +1,440 @@ +From 304072b6240ebf11246255b2138c8d201e0014ee Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Wed, 17 Jun 2026 07:36:49 +1000 +Subject: [PATCH 90/94] Add Skippy batched stage execution ABI + +--- + include/skippy.h | 24 ++++ + include/skippy/common.h | 2 +- + src/skippy.cpp | 296 ++++++++++++++++++++++++++++++++++++---- + 3 files changed, 294 insertions(+), 28 deletions(-) + +diff --git a/include/skippy.h b/include/skippy.h +index 075850e3..7cace0a4 100644 +--- a/include/skippy.h ++++ b/include/skippy.h +@@ -278,6 +278,15 @@ LLAMA_API enum skippy_status skippy_decode_step_sampled( + llama_token * out_predicted_token, + struct skippy_error ** out_error); + ++LLAMA_API enum skippy_status skippy_decode_batch_sampled( ++ struct skippy_session * const * sessions, ++ const llama_token * token_ids, ++ const struct skippy_sampling_config * const * sampling, ++ size_t request_count, ++ llama_token * out_predicted_tokens, ++ size_t predicted_token_capacity, ++ struct skippy_error ** out_error); ++ + LLAMA_API enum skippy_status skippy_prefill_chunk_frame( + struct skippy_session * session, + const llama_token * token_ids, +@@ -359,6 +368,21 @@ LLAMA_API enum skippy_status skippy_decode_step_frame_sampled( + llama_token * out_predicted_token, + struct skippy_error ** out_error); + ++LLAMA_API enum skippy_status skippy_decode_step_frame_batch_sampled( ++ struct skippy_session * const * sessions, ++ const llama_token * token_ids, ++ const struct skippy_sampling_config * const * sampling, ++ const struct skippy_activation_desc * const * input_descs, ++ const void * const * input_payloads, ++ struct skippy_activation_desc * output_descs, ++ void * const * output_payloads, ++ const size_t * output_payload_capacities, ++ size_t * out_output_payload_bytes, ++ llama_token * out_predicted_tokens, ++ size_t predicted_token_capacity, ++ size_t request_count, ++ struct skippy_error ** out_error); ++ + LLAMA_API enum skippy_status skippy_verify_tokens_frame( + struct skippy_session * session, + const llama_token * token_ids, +diff --git a/include/skippy/common.h b/include/skippy/common.h +index a6bd7fd9..7009d98c 100644 +--- a/include/skippy/common.h ++++ b/include/skippy/common.h +@@ -26,7 +26,7 @@ extern "C" { + + #define SKIPPY_ABI_VERSION_MAJOR 0 + #define SKIPPY_ABI_VERSION_MINOR 1 +-#define SKIPPY_ABI_VERSION_PATCH 26 ++#define SKIPPY_ABI_VERSION_PATCH 27 + + enum skippy_feature { + SKIPPY_FEATURE_RUNTIME_SLICE = 1 << 0, +diff --git a/src/skippy.cpp b/src/skippy.cpp +index 20cff23f..4707be19 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -1335,23 +1335,6 @@ static enum skippy_status skippy_verify_token_batch( + return status; + } + +-static llama_token skippy_greedy_sample(skippy_session * session) { +- const llama_vocab * vocab = llama_model_get_vocab(session->stage_model->model); +- const int32_t n_vocab = llama_vocab_n_tokens(vocab); +- const float * logits = llama_get_logits_ith(session->ctx, -1); +- +- llama_token best = 0; +- float best_logit = -std::numeric_limits::infinity(); +- for (int32_t token = 0; token < n_vocab; ++token) { +- if (logits[token] > best_logit) { +- best_logit = logits[token]; +- best = token; +- } +- } +- +- return best; +-} +- + static llama_token skippy_greedy_sample_ith(skippy_session * session, int32_t index) { + const llama_vocab * vocab = llama_model_get_vocab(session->stage_model->model); + const int32_t n_vocab = llama_vocab_n_tokens(vocab); +@@ -1372,6 +1355,10 @@ static llama_token skippy_greedy_sample_ith(skippy_session * session, int32_t in + return best; + } + ++static llama_token skippy_greedy_sample(skippy_session * session) { ++ return skippy_greedy_sample_ith(session, -1); ++} ++ + static bool skippy_compute_token_signal( + skippy_session * session, + int32_t logits_index, +@@ -1656,14 +1643,14 @@ static void skippy_sync_chat_sampling_history(skippy_session * session) { + } + } + +-static llama_token skippy_chat_sample_token(skippy_session * session) { ++static llama_token skippy_chat_sample_token(skippy_session * session, int32_t logits_index) { + skippy_sync_chat_sampling_history(session); + + const llama_vocab * vocab = llama_model_get_vocab(session->stage_model->model); + const int32_t n_vocab = llama_vocab_n_tokens(vocab); +- const float * logits = llama_get_logits_ith(session->ctx, -1); ++ const float * logits = llama_get_logits_ith(session->ctx, logits_index); + if (logits == nullptr) { +- return skippy_greedy_sample(session); ++ return skippy_greedy_sample_ith(session, logits_index); + } + + std::vector candidates; +@@ -1688,25 +1675,26 @@ static llama_token skippy_chat_sample_token(skippy_session * session) { + } + llama_sampler_apply(session->sampling_chain, &cur); + if (cur.selected < 0 || static_cast(cur.selected) >= cur.size) { +- return skippy_greedy_sample(session); ++ return skippy_greedy_sample_ith(session, logits_index); + } + return cur.data[cur.selected].id; + } + +-static llama_token skippy_sample_token( ++static llama_token skippy_sample_token_ith( + skippy_session * session, +- const skippy_sampling_config * sampling) { ++ const skippy_sampling_config * sampling, ++ int32_t logits_index) { + if (session != nullptr && session->sampling_chain != nullptr) { +- return skippy_chat_sample_token(session); ++ return skippy_chat_sample_token(session, logits_index); + } + if (!skippy_sampling_enabled(sampling)) { +- return skippy_greedy_sample(session); ++ return skippy_greedy_sample_ith(session, logits_index); + } + + llama_sampler_chain_params chain_params = llama_sampler_chain_default_params(); + llama_sampler * sampler = llama_sampler_chain_init(chain_params); + if (sampler == nullptr) { +- return skippy_greedy_sample(session); ++ return skippy_greedy_sample_ith(session, logits_index); + } + + const int32_t penalty_last_n = sampling->penalty_last_n == 0 ? -1 : sampling->penalty_last_n; +@@ -1752,11 +1740,17 @@ static llama_token skippy_sample_token( + for (const llama_token token : session->token_history) { + llama_sampler_accept(sampler, token); + } +- llama_token token = llama_sampler_sample(sampler, session->ctx, -1); ++ llama_token token = llama_sampler_sample(sampler, session->ctx, logits_index); + llama_sampler_free(sampler); + return token; + } + ++static llama_token skippy_sample_token( ++ skippy_session * session, ++ const skippy_sampling_config * sampling) { ++ return skippy_sample_token_ith(session, sampling, -1); ++} ++ + static enum skippy_status skippy_prepare_empty_activation_frame( + skippy_session * session, + size_t token_count, +@@ -2940,6 +2934,80 @@ enum skippy_status skippy_decode_step_sampled( + return skippy_success(out_error); + } + ++enum skippy_status skippy_decode_batch_sampled( ++ struct skippy_session * const * sessions, ++ const llama_token * token_ids, ++ const struct skippy_sampling_config * const * sampling, ++ size_t request_count, ++ llama_token * out_predicted_tokens, ++ size_t predicted_token_capacity, ++ struct skippy_error ** out_error) { ++ if (sessions == nullptr || token_ids == nullptr || request_count == 0 || out_predicted_tokens == nullptr) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "sessions, token_ids, request_count, and out_predicted_tokens are required"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ if (predicted_token_capacity < request_count) { ++ skippy_set_error(out_error, SKIPPY_STATUS_BUFFER_TOO_SMALL, "predicted token output buffer is too small"); ++ return SKIPPY_STATUS_BUFFER_TOO_SMALL; ++ } ++ if (request_count > static_cast(std::numeric_limits::max())) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "request_count exceeds int32_t range"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ ++ skippy_session * first = sessions[0]; ++ if (first == nullptr || first->ctx == nullptr || first->stage_model == nullptr) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "sessions must be active"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ if (skippy_is_filtered(first) && first->stage_model->config.layer_start > 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_UNSUPPORTED, "batched token decode requires the first runtime slice or a full model"); ++ return SKIPPY_STATUS_UNSUPPORTED; ++ } ++ ++ const int32_t n_tokens = static_cast(request_count); ++ llama_batch batch = llama_batch_init(n_tokens, 0, 1); ++ batch.n_tokens = n_tokens; ++ for (int32_t i = 0; i < n_tokens; ++i) { ++ skippy_session * session = sessions[i]; ++ if (session == nullptr || session->ctx != first->ctx || session->stage_model != first->stage_model) { ++ llama_batch_free(batch); ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "all sessions must belong to the same stage model"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ if (skippy_is_filtered(session) && session->stage_model->config.layer_start > 0) { ++ llama_batch_free(batch); ++ skippy_set_error(out_error, SKIPPY_STATUS_UNSUPPORTED, "batched token decode requires the first runtime slice or a full model"); ++ return SKIPPY_STATUS_UNSUPPORTED; ++ } ++ batch.token[i] = token_ids[i]; ++ batch.pos[i] = session->n_past; ++ batch.n_seq_id[i] = 1; ++ batch.seq_id[i][0] = session->seq_id; ++ batch.logits[i] = 1; ++ } ++ ++ skippy_graph_filter_scope graph_filter_scope(&first->stage_model->config); ++ const int32_t rc = llama_decode(first->ctx, batch); ++ if (rc != 0) { ++ llama_batch_free(batch); ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "llama_decode failed"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ llama_batch_free(batch); ++ ++ for (int32_t i = 0; i < n_tokens; ++i) { ++ skippy_session * session = sessions[i]; ++ session->n_past += 1; ++ skippy_record_tokens(session, &token_ids[i], 1); ++ skippy_record_signal(session, i); ++ const skippy_sampling_config * request_sampling = sampling != nullptr ? sampling[i] : nullptr; ++ out_predicted_tokens[i] = skippy_sample_token_ith(session, request_sampling, i); ++ } ++ ++ return skippy_success(out_error); ++} ++ + enum skippy_status skippy_verify_tokens( + struct skippy_session * session, + const llama_token * token_ids, +@@ -3246,6 +3314,181 @@ enum skippy_status skippy_decode_step_frame_sampled( + return skippy_copy_output_activation_frame(session, 1, output_payload, input_desc, input_payload, out_error); + } + ++enum skippy_status skippy_decode_step_frame_batch_sampled( ++ struct skippy_session * const * sessions, ++ const llama_token * token_ids, ++ const struct skippy_sampling_config * const * sampling, ++ const struct skippy_activation_desc * const * input_descs, ++ const void * const * input_payloads, ++ struct skippy_activation_desc * output_descs, ++ void * const * output_payloads, ++ const size_t * output_payload_capacities, ++ size_t * out_output_payload_bytes, ++ llama_token * out_predicted_tokens, ++ size_t predicted_token_capacity, ++ size_t request_count, ++ struct skippy_error ** out_error) { ++ if (sessions == nullptr || token_ids == nullptr || request_count == 0 || out_predicted_tokens == nullptr) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "sessions, token_ids, request_count, and out_predicted_tokens are required"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ if (predicted_token_capacity < request_count) { ++ skippy_set_error(out_error, SKIPPY_STATUS_BUFFER_TOO_SMALL, "predicted token output buffer is too small"); ++ return SKIPPY_STATUS_BUFFER_TOO_SMALL; ++ } ++ if (request_count > static_cast(std::numeric_limits::max())) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "request_count exceeds int32_t range"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ ++ skippy_session * first = sessions[0]; ++ if (first == nullptr || first->ctx == nullptr || first->stage_model == nullptr) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "sessions must be active"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ ++ const int32_t n_tokens = static_cast(request_count); ++ const int32_t n_embd = llama_model_n_embd(first->stage_model->model); ++ const int32_t n_embd_inp = llama_model_n_embd_inp(first->stage_model->model); ++ const int32_t n_pos_per_embd = first->stage_model->model->hparams.n_pos_per_embd(); ++ const bool activation_input = skippy_is_filtered(first) && first->stage_model->config.layer_start > 0; ++ const bool request_logits = first->stage_model->config.include_output; ++ const bool request_embeddings = skippy_emits_activation_frame(first); ++ ++ if (activation_input && (input_descs == nullptr || input_payloads == nullptr)) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "batched downstream decode requires activation frame inputs"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ ++ const size_t hidden_bytes_per_request = skippy_activation_hidden_bytes(first, 1); ++ for (int32_t i = 0; i < n_tokens; ++i) { ++ skippy_session * session = sessions[i]; ++ if (session == nullptr || session->ctx != first->ctx || session->stage_model != first->stage_model) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "all sessions must belong to the same stage model"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ const skippy_activation_desc * input_desc = input_descs != nullptr ? input_descs[i] : nullptr; ++ const void * input_payload = input_payloads != nullptr ? input_payloads[i] : nullptr; ++ enum skippy_status status = skippy_validate_frame_input(session, input_desc, input_payload, 1, out_error); ++ if (status != SKIPPY_STATUS_OK) { ++ return status; ++ } ++ if (input_desc != nullptr && input_desc->flags != 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_UNSUPPORTED, "batched activation decode does not support activation sidebands yet"); ++ return SKIPPY_STATUS_UNSUPPORTED; ++ } ++ if (skippy_output_activation_flags(session, input_desc) != 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_UNSUPPORTED, "batched activation decode does not support output activation sidebands yet"); ++ return SKIPPY_STATUS_UNSUPPORTED; ++ } ++ status = skippy_prepare_output_activation_frame( ++ session, ++ 1, ++ output_payloads != nullptr ? output_payloads[i] : nullptr, ++ output_payload_capacities != nullptr ? output_payload_capacities[i] : 0, ++ out_output_payload_bytes != nullptr ? &out_output_payload_bytes[i] : nullptr, ++ output_descs != nullptr ? &output_descs[i] : nullptr, ++ input_desc, ++ out_error); ++ if (status != SKIPPY_STATUS_OK) { ++ return status; ++ } ++ } ++ ++ std::vector embd_storage; ++ std::vector token_storage; ++ if (activation_input) { ++ embd_storage.resize(static_cast(n_tokens)*n_embd_inp); ++ for (int32_t i = 0; i < n_tokens; ++i) { ++ const skippy_activation_desc * input_desc = input_descs[i]; ++ const void * input_payload = input_payloads[i]; ++ float * dst = embd_storage.data() + static_cast(i)*n_embd_inp; ++ if ((input_desc->flags & SKIPPY_ACTIVATION_FLAG_GEMMA3N_ALTUP) != 0) { ++ return SKIPPY_STATUS_UNSUPPORTED; ++ } else if (n_embd_inp == n_embd) { ++ std::memcpy(dst, input_payload, hidden_bytes_per_request); ++ } else { ++ std::memcpy(dst, input_payload, static_cast(n_embd)*sizeof(float)); ++ std::memset(dst + n_embd, 0, static_cast(n_embd_inp - n_embd)*sizeof(float)); ++ } ++ } ++ } else { ++ token_storage.assign(token_ids, token_ids + request_count); ++ } ++ ++ std::vector pos_storage(static_cast(n_tokens)*n_pos_per_embd); ++ std::vector n_seq_id_storage(n_tokens, 1); ++ std::vector seq_id_values(n_tokens); ++ std::vector seq_id_storage(n_tokens, nullptr); ++ std::vector logits_storage(n_tokens); ++ for (int32_t i = 0; i < n_tokens; ++i) { ++ skippy_session * session = sessions[i]; ++ seq_id_values[i] = session->seq_id; ++ seq_id_storage[i] = &seq_id_values[i]; ++ logits_storage[i] = (request_logits || request_embeddings) ? 1 : 0; ++ if (n_pos_per_embd == 4) { ++ const llama_pos position = session->n_past; ++ pos_storage[ i] = position; ++ pos_storage[ n_tokens + i] = position; ++ pos_storage[2 * n_tokens + i] = position; ++ pos_storage[3 * n_tokens + i] = 0; ++ } else { ++ pos_storage[i] = session->n_past; ++ } ++ } ++ ++ llama_batch batch = { ++ /*n_tokens =*/ n_tokens, ++ /*token =*/ activation_input ? nullptr : token_storage.data(), ++ /*embd =*/ activation_input ? embd_storage.data() : nullptr, ++ /*pos =*/ pos_storage.data(), ++ /*n_seq_id =*/ n_seq_id_storage.data(), ++ /*seq_id =*/ seq_id_storage.data(), ++ /*logits =*/ logits_storage.data(), ++ }; ++ ++ skippy_activation_tokens_scope activation_tokens_scope(token_ids, request_count); ++ skippy_graph_filter_scope graph_filter_scope(&first->stage_model->config); ++ const int32_t rc = llama_decode(first->ctx, batch); ++ if (rc != 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "llama_decode failed"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ ++ float * embeddings = nullptr; ++ if (skippy_emits_activation_frame(first)) { ++ embeddings = llama_get_embeddings(first->ctx); ++ if (embeddings == nullptr) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "llama embeddings output was not available"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ } ++ ++ for (int32_t i = 0; i < n_tokens; ++i) { ++ skippy_session * session = sessions[i]; ++ session->n_past += 1; ++ if (!activation_input) { ++ skippy_record_tokens(session, &token_ids[i], 1); ++ } ++ if (request_logits) { ++ skippy_record_signal(session, i); ++ const skippy_sampling_config * request_sampling = sampling != nullptr ? sampling[i] : nullptr; ++ out_predicted_tokens[i] = skippy_sample_token_ith(session, request_sampling, i); ++ } else { ++ out_predicted_tokens[i] = -1; ++ } ++ ++ if (embeddings != nullptr && output_payloads != nullptr && output_payloads[i] != nullptr) { ++ std::memcpy( ++ output_payloads[i], ++ embeddings + static_cast(i)*n_embd, ++ hidden_bytes_per_request); ++ } ++ } ++ ++ return skippy_success(out_error); ++} ++ + enum skippy_status skippy_verify_tokens_frame( + struct skippy_session * session, + const llama_token * token_ids, +-- +2.54.0 (Apple Git-156) diff --git a/third_party/llama.cpp/patches/0101-Add-Skippy-native-MTP-n1-sidecar.patch b/third_party/llama.cpp/patches/0101-Add-Skippy-native-MTP-n1-sidecar.patch new file mode 100644 index 0000000000..5809d9ac6b --- /dev/null +++ b/third_party/llama.cpp/patches/0101-Add-Skippy-native-MTP-n1-sidecar.patch @@ -0,0 +1,566 @@ +From 0497272eb9b14bf25e3b9a64cce576efed009a15 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Wed, 17 Jun 2026 07:37:38 +1000 +Subject: [PATCH 91/94] Add Skippy native MTP n1 sidecar + +--- + include/skippy.h | 21 +++ + include/skippy/common.h | 1 + + src/skippy.cpp | 290 ++++++++++++++++++++++++++++++++++++++-- + 3 files changed, 301 insertions(+), 11 deletions(-) + +diff --git a/include/skippy.h b/include/skippy.h +index 7cace0a4..c4da8e28 100644 +--- a/include/skippy.h ++++ b/include/skippy.h +@@ -131,6 +131,13 @@ struct skippy_kv_page_desc { + uint64_t flags; + }; + ++struct skippy_native_mtp_draft { ++ uint32_t version; ++ bool available; ++ llama_token token_id; ++ int64_t proposal_compute_us; ++}; ++ + LLAMA_API enum skippy_status skippy_model_open( + const char * path, + const struct skippy_runtime_config * config, +@@ -368,6 +375,20 @@ LLAMA_API enum skippy_status skippy_decode_step_frame_sampled( + llama_token * out_predicted_token, + struct skippy_error ** out_error); + ++LLAMA_API enum skippy_status skippy_decode_step_frame_sampled_mtp_n1( ++ struct skippy_session * session, ++ llama_token token_id, ++ const struct skippy_sampling_config * sampling, ++ const struct skippy_activation_desc * input_desc, ++ const void * input_payload, ++ struct skippy_activation_desc * output_desc, ++ void * output_payload, ++ size_t output_payload_capacity, ++ size_t * out_output_payload_bytes, ++ llama_token * out_predicted_token, ++ struct skippy_native_mtp_draft * out_mtp_draft, ++ struct skippy_error ** out_error); ++ + LLAMA_API enum skippy_status skippy_decode_step_frame_batch_sampled( + struct skippy_session * const * sessions, + const llama_token * token_ids, +diff --git a/include/skippy/common.h b/include/skippy/common.h +index 7009d98c..cf97f7ab 100644 +--- a/include/skippy/common.h ++++ b/include/skippy/common.h +@@ -54,6 +54,7 @@ enum skippy_feature { + SKIPPY_FEATURE_CHAT_SAMPLING_GRAMMAR = 1 << 22, + SKIPPY_FEATURE_BACKEND_DEVICES = 1 << 23, + SKIPPY_FEATURE_RUNTIME_EVENTS = 1 << 24, ++ SKIPPY_FEATURE_NATIVE_MTP_N1 = 1 << 25, + }; + + enum skippy_status { +diff --git a/src/skippy.cpp b/src/skippy.cpp +index 4707be19..7d909a9e 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -6,6 +6,7 @@ + #include "gguf.h" + #include "llama-arch.h" + #include "llama-context.h" ++#include "llama-ext.h" + #include "llama-graph.h" + #include "llama-kv-cache.h" + #include "llama-memory-hybrid.h" +@@ -45,6 +46,7 @@ using json = nlohmann::ordered_json; + struct skippy_model { + llama_model * model = nullptr; + llama_context * ctx = nullptr; ++ llama_context * mtp_ctx = nullptr; + skippy_runtime_config config = {}; + bool executable = true; + uint32_t lane_count = 1; +@@ -73,6 +75,12 @@ struct skippy_session { + std::string chat_sampling_metadata; + uint64_t grammar_generated_start = 0; + size_t sampling_accepted_token_count = 0; ++ int32_t mtp_next_pos = 0; ++ bool mtp_has_pending_h = false; ++ std::vector mtp_pending_h; ++ bool mtp_has_pending_draft = false; ++ llama_pos mtp_pending_draft_pos = 0; ++ llama_token mtp_pending_draft_token = -1; + }; + + struct skippy_tensor_meta { +@@ -1227,6 +1232,135 @@ static enum skippy_status skippy_decode_batch( + return skippy_success(out_error); + } + ++static bool skippy_mtp_available(const skippy_session * session) { ++ return session != nullptr && ++ session->stage_model != nullptr && ++ session->stage_model->mtp_ctx != nullptr && ++ session->stage_model->config.include_output; ++} ++ ++static void skippy_mtp_clear_session_state(skippy_session * session) { ++ if (session == nullptr) { ++ return; ++ } ++ session->mtp_next_pos = session->n_past; ++ session->mtp_has_pending_h = false; ++ session->mtp_pending_h.clear(); ++ session->mtp_has_pending_draft = false; ++ session->mtp_pending_draft_pos = 0; ++ session->mtp_pending_draft_token = -1; ++ if (skippy_mtp_available(session)) { ++ if (llama_memory_t memory = session->stage_model->mtp_ctx->get_memory()) { ++ llama_memory_seq_rm(memory, session->seq_id, -1, -1); ++ } ++ } ++} ++ ++static enum skippy_status skippy_mtp_sync_target_tokens( ++ skippy_session * session, ++ const llama_token * token_ids, ++ size_t token_count, ++ llama_pos token_start, ++ struct skippy_error ** out_error) { ++ if (!skippy_mtp_available(session) || token_ids == nullptr || token_count == 0) { ++ return skippy_success(out_error); ++ } ++ if (token_count > static_cast(std::numeric_limits::max())) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "MTP token_count exceeds int32_t range"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ ++ llama_context * mtp_ctx = session->stage_model->mtp_ctx; ++ const int32_t n_embd = llama_model_n_embd(session->stage_model->model); ++ const size_t row_bytes = static_cast(n_embd)*sizeof(float); ++ if (session->mtp_pending_h.size() != static_cast(n_embd)) { ++ session->mtp_pending_h.assign(static_cast(n_embd), 0.0f); ++ session->mtp_has_pending_h = false; ++ } ++ ++ if (session->mtp_has_pending_draft) { ++ const llama_pos token_end = token_start + static_cast(token_count); ++ if (session->mtp_pending_draft_pos >= token_start && ++ session->mtp_pending_draft_pos < token_end) { ++ const size_t draft_index = static_cast(session->mtp_pending_draft_pos - token_start); ++ if (token_ids[draft_index] == session->mtp_pending_draft_token) { ++ session->mtp_next_pos = std::max(session->mtp_next_pos, session->mtp_pending_draft_pos + 1); ++ } else { ++ if (llama_memory_t memory = mtp_ctx->get_memory()) { ++ llama_memory_seq_rm(memory, session->seq_id, session->mtp_pending_draft_pos, -1); ++ } ++ session->mtp_next_pos = std::min(session->mtp_next_pos, static_cast(session->mtp_pending_draft_pos)); ++ } ++ } ++ session->mtp_has_pending_draft = false; ++ session->mtp_pending_draft_pos = 0; ++ session->mtp_pending_draft_token = -1; ++ } ++ ++ size_t first_index = 0; ++ if (session->mtp_next_pos > token_start) { ++ first_index = static_cast(std::min( ++ static_cast(token_count), ++ session->mtp_next_pos - token_start)); ++ } ++ ++ const int32_t n_decode = static_cast(token_count - first_index); ++ if (n_decode > 0) { ++ llama_batch batch = llama_batch_init(n_decode, n_embd, 1); ++ batch.token = static_cast(std::malloc(sizeof(llama_token)*static_cast(n_decode))); ++ if (batch.token == nullptr) { ++ llama_batch_free(batch); ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to allocate MTP token batch"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ batch.n_tokens = n_decode; ++ ++ for (int32_t out_i = 0; out_i < n_decode; ++out_i) { ++ const size_t src_i = first_index + static_cast(out_i); ++ batch.token[out_i] = token_ids[src_i]; ++ batch.pos[out_i] = token_start + static_cast(src_i); ++ batch.n_seq_id[out_i] = 1; ++ batch.seq_id[out_i][0] = session->seq_id; ++ batch.logits[out_i] = 0; ++ ++ float * dst = batch.embd + static_cast(out_i)*n_embd; ++ if (src_i == 0) { ++ std::memcpy(dst, session->mtp_pending_h.data(), row_bytes); ++ } else { ++ const float * h_prev = llama_get_embeddings_nextn_ith(session->ctx, static_cast(src_i - 1)); ++ if (h_prev == nullptr) { ++ std::free(batch.token); ++ batch.token = nullptr; ++ llama_batch_free(batch); ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "target pre-norm hidden row was not available"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ std::memcpy(dst, h_prev, row_bytes); ++ } ++ } ++ ++ skippy_graph_filter_scope graph_filter_scope(&session->stage_model->config); ++ const int32_t rc = llama_decode(mtp_ctx, batch); ++ std::free(batch.token); ++ batch.token = nullptr; ++ llama_batch_free(batch); ++ if (rc != 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "llama_decode failed for MTP sidecar sync"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ session->mtp_next_pos = token_start + static_cast(token_count); ++ } ++ ++ const float * h_last = llama_get_embeddings_nextn_ith(session->ctx, static_cast(token_count - 1)); ++ if (h_last == nullptr) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "target pre-norm hidden row was not available"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ std::memcpy(session->mtp_pending_h.data(), h_last, row_bytes); ++ session->mtp_has_pending_h = true; ++ return skippy_success(out_error); ++} ++ + static void skippy_record_tokens( + skippy_session * session, + const llama_token * token_ids, +@@ -1277,12 +1389,16 @@ static enum skippy_status skippy_decode_tokens( + /*logits =*/ &logits, + }; + enum skippy_status status = skippy_decode_batch(session, batch, 1, out_error); ++ if (status == SKIPPY_STATUS_OK) { ++ status = skippy_mtp_sync_target_tokens(session, token_ids, token_count, pos, out_error); ++ } + if (status == SKIPPY_STATUS_OK) { + skippy_record_tokens(session, token_ids, token_count); + } + return status; + } + ++ const llama_pos token_start = session->n_past; + llama_batch batch = llama_batch_init(n_tokens, 0, 1); + batch.n_tokens = n_tokens; + for (int32_t i = 0; i < n_tokens; ++i) { +@@ -1295,6 +1411,9 @@ static enum skippy_status skippy_decode_tokens( + + enum skippy_status status = skippy_decode_batch(session, batch, token_count, out_error); + llama_batch_free(batch); ++ if (status == SKIPPY_STATUS_OK) { ++ status = skippy_mtp_sync_target_tokens(session, token_ids, token_count, token_start, out_error); ++ } + if (status == SKIPPY_STATUS_OK) { + skippy_record_tokens(session, token_ids, token_count); + } +@@ -1335,10 +1454,16 @@ static enum skippy_status skippy_verify_token_batch( + return status; + } + +-static llama_token skippy_greedy_sample_ith(skippy_session * session, int32_t index) { +- const llama_vocab * vocab = llama_model_get_vocab(session->stage_model->model); ++static llama_token skippy_greedy_sample_context( ++ llama_model * model, ++ llama_context * ctx, ++ int32_t index) { ++ if (model == nullptr || ctx == nullptr) { ++ return 0; ++ } ++ const llama_vocab * vocab = llama_model_get_vocab(model); + const int32_t n_vocab = llama_vocab_n_tokens(vocab); +- const float * logits = llama_get_logits_ith(session->ctx, index); ++ const float * logits = llama_get_logits_ith(ctx, index); + if (logits == nullptr) { + return 0; + } +@@ -1355,6 +1480,10 @@ static llama_token skippy_greedy_sample_ith(skippy_session * session, int32_t in + return best; + } + ++static llama_token skippy_greedy_sample_ith(skippy_session * session, int32_t index) { ++ return skippy_greedy_sample_context(session->stage_model->model, session->ctx, index); ++} ++ + static llama_token skippy_greedy_sample(skippy_session * session) { + return skippy_greedy_sample_ith(session, -1); + } +@@ -1751,6 +1880,66 @@ static llama_token skippy_sample_token( + return skippy_sample_token_ith(session, sampling, -1); + } + ++static enum skippy_status skippy_mtp_propose_next( ++ skippy_session * session, ++ llama_token predicted_token, ++ skippy_native_mtp_draft * out_mtp_draft, ++ struct skippy_error ** out_error) { ++ if (out_mtp_draft != nullptr) { ++ *out_mtp_draft = { ++ 1, ++ false, ++ -1, ++ 0, ++ }; ++ } ++ if (!skippy_mtp_available(session) || predicted_token < 0 || !session->mtp_has_pending_h) { ++ return skippy_success(out_error); ++ } ++ ++ llama_context * mtp_ctx = session->stage_model->mtp_ctx; ++ const int32_t n_embd = llama_model_n_embd(session->stage_model->model); ++ if (session->mtp_pending_h.size() != static_cast(n_embd)) { ++ return skippy_success(out_error); ++ } ++ ++ llama_token token = predicted_token; ++ llama_pos pos = session->n_past; ++ int32_t n_seq_id = 1; ++ llama_seq_id seq_id = session->seq_id; ++ llama_seq_id * seq_ids = &seq_id; ++ int8_t logits = 1; ++ llama_batch batch = { ++ /*n_tokens =*/ 1, ++ /*token =*/ &token, ++ /*embd =*/ session->mtp_pending_h.data(), ++ /*pos =*/ &pos, ++ /*n_seq_id =*/ &n_seq_id, ++ /*seq_id =*/ &seq_ids, ++ /*logits =*/ &logits, ++ }; ++ ++ const int64_t t_start_us = ggml_time_us(); ++ skippy_graph_filter_scope graph_filter_scope(&session->stage_model->config); ++ const int32_t rc = llama_decode(mtp_ctx, batch); ++ const int64_t elapsed_us = ggml_time_us() - t_start_us; ++ if (rc != 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "llama_decode failed for MTP sidecar proposal"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ ++ if (out_mtp_draft != nullptr) { ++ out_mtp_draft->version = 1; ++ out_mtp_draft->available = true; ++ out_mtp_draft->token_id = skippy_greedy_sample_context(session->stage_model->model, mtp_ctx, -1); ++ out_mtp_draft->proposal_compute_us = elapsed_us; ++ session->mtp_has_pending_draft = true; ++ session->mtp_pending_draft_pos = session->n_past; ++ session->mtp_pending_draft_token = out_mtp_draft->token_id; ++ } ++ return skippy_success(out_error); ++} ++ + static enum skippy_status skippy_prepare_empty_activation_frame( + skippy_session * session, + size_t token_count, +@@ -1939,7 +2126,11 @@ static enum skippy_status skippy_decode_activation_frame( + skippy_activation_tokens_scope activation_tokens_scope(token_ids, token_count); + skippy_rwkv7_v_first_scope rwkv7_v_first_scope(input_desc, input_payload, hidden_bytes, n_embd); + skippy_gemma3n_altup_scope gemma3n_altup_scope(input_desc, input_payload, n_embd, n_altup); ++ const llama_pos token_start = pos_storage.empty() ? session->n_past : pos_storage[0]; + enum skippy_status status = skippy_decode_batch(session, batch, token_count, out_error); ++ if (status == SKIPPY_STATUS_OK && token_ids != nullptr) { ++ status = skippy_mtp_sync_target_tokens(session, token_ids, token_count, token_start, out_error); ++ } + return status; + } + +@@ -1947,6 +2138,7 @@ static enum skippy_status skippy_verify_activation_frame( + skippy_session * session, + const skippy_activation_desc * input_desc, + const void * input_payload, ++ const llama_token * token_ids, + size_t token_count, + struct skippy_error ** out_error) { + if (session == nullptr || session->ctx == nullptr || input_desc == nullptr || input_payload == nullptr || token_count == 0) { +@@ -1975,8 +2167,12 @@ static enum skippy_status skippy_verify_activation_frame( + + skippy_rwkv7_v_first_scope rwkv7_v_first_scope(input_desc, input_payload, hidden_bytes, n_embd); + skippy_gemma3n_altup_scope gemma3n_altup_scope(input_desc, input_payload, n_embd, n_altup); ++ const llama_pos token_start = session->n_past; + enum skippy_status status = skippy_decode_batch(session, batch, token_count, out_error); + llama_batch_free(batch); ++ if (status == SKIPPY_STATUS_OK && token_ids != nullptr) { ++ status = skippy_mtp_sync_target_tokens(session, token_ids, token_count, token_start, out_error); ++ } + return status; + } + +@@ -2011,7 +2207,8 @@ uint64_t skippy_abi_features(void) { + SKIPPY_FEATURE_GENERATION_SIGNALS | + SKIPPY_FEATURE_EXTERNAL_MEDIA_PREFILL | + SKIPPY_FEATURE_RUNTIME_EVENTS | +- SKIPPY_FEATURE_BACKEND_DEVICES; ++ SKIPPY_FEATURE_BACKEND_DEVICES | ++ SKIPPY_FEATURE_NATIVE_MTP_N1; + } + + const char * skippy_status_string(enum skippy_status status) { +@@ -2287,6 +2484,7 @@ static enum skippy_status skippy_finish_model_open( + params.embeddings = config != nullptr && config->filter_tensors_on_load && !config->include_output; + if (llm_arch_is_recurrent(model->arch) || llm_arch_is_hybrid(model->arch)) { + params.n_seq_max = std::max(2, stage_model->lane_count * 2); ++ params.n_rs_seq = std::max(params.n_rs_seq, 2); + params.kv_unified = true; + } + +@@ -2302,6 +2500,25 @@ static enum skippy_status skippy_finish_model_open( + skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, message); + return SKIPPY_STATUS_RUNTIME_ERROR; + } ++ ++ if (config != nullptr && ++ config->include_output && ++ model->hparams.n_layer_nextn > 0) { ++ llama_context_params mtp_params = params; ++ mtp_params.ctx_type = LLAMA_CONTEXT_TYPE_MTP; ++ mtp_params.embeddings = false; ++ { ++ skippy_graph_filter_scope graph_filter_scope(config); ++ stage_model->mtp_ctx = llama_init_from_model(model, mtp_params); ++ } ++ if (stage_model->mtp_ctx != nullptr) { ++ llama_set_embeddings_nextn(stage_model->ctx, true, false); ++ llama_set_embeddings_nextn(stage_model->mtp_ctx, true, true); ++ } else { ++ fprintf(stderr, "skippy: native MTP sidecar unavailable for this final stage; continuing without drafts\n"); ++ } ++ } ++ + stage_model->lane_in_use.assign(stage_model->lane_count, false); + stage_model->lane_resident_prefix_tokens.resize(stage_model->lane_count); + +@@ -2515,6 +2732,9 @@ enum skippy_status skippy_model_free( + struct skippy_model * model, + struct skippy_error ** out_error) { + if (model != nullptr) { ++ if (model->mtp_ctx != nullptr) { ++ llama_free(model->mtp_ctx); ++ } + if (model->ctx != nullptr) { + llama_free(model->ctx); + } +@@ -2562,6 +2782,11 @@ enum skippy_status skippy_session_create( + session->n_past = 0; + session->checkpoint_valid = false; + session->checkpoint_n_past = 0; ++ session->mtp_next_pos = 0; ++ if (skippy_mtp_available(session)) { ++ session->mtp_pending_h.assign(static_cast(llama_model_n_embd(model->model)), 0.0f); ++ skippy_mtp_clear_session_state(session); ++ } + *out_session = session; + if (static_cast(seq_id) < model->lane_resident_prefix_tokens.size()) { + model->lane_resident_prefix_tokens[static_cast(seq_id)].clear(); +@@ -2645,6 +2870,7 @@ enum skippy_status skippy_session_set_position( + if (session->signal_history.size() > static_cast(n_past)) { + session->signal_history.resize(static_cast(n_past)); + } ++ skippy_mtp_clear_session_state(session); + return skippy_success(out_error); + } + +@@ -2738,6 +2964,7 @@ enum skippy_status skippy_session_reset( + } + session->signal_history.clear(); + skippy_clear_chat_sampling(session); ++ skippy_mtp_clear_session_state(session); + session->ctx->synchronize(); + return skippy_success(out_error); + } +@@ -2774,6 +3001,7 @@ enum skippy_status skippy_session_free( + session->stage_model->lane_in_use[lane] = false; + } + } ++ skippy_mtp_clear_session_state(session); + delete session; + } + return skippy_success(out_error); +@@ -3314,6 +3542,47 @@ enum skippy_status skippy_decode_step_frame_sampled( + return skippy_copy_output_activation_frame(session, 1, output_payload, input_desc, input_payload, out_error); + } + ++enum skippy_status skippy_decode_step_frame_sampled_mtp_n1( ++ struct skippy_session * session, ++ llama_token token_id, ++ const struct skippy_sampling_config * sampling, ++ const struct skippy_activation_desc * input_desc, ++ const void * input_payload, ++ struct skippy_activation_desc * output_desc, ++ void * output_payload, ++ size_t output_payload_capacity, ++ size_t * out_output_payload_bytes, ++ llama_token * out_predicted_token, ++ struct skippy_native_mtp_draft * out_mtp_draft, ++ struct skippy_error ** out_error) { ++ if (out_mtp_draft != nullptr) { ++ *out_mtp_draft = { ++ 1, ++ false, ++ -1, ++ 0, ++ }; ++ } ++ ++ enum skippy_status status = skippy_decode_step_frame_sampled( ++ session, ++ token_id, ++ sampling, ++ input_desc, ++ input_payload, ++ output_desc, ++ output_payload, ++ output_payload_capacity, ++ out_output_payload_bytes, ++ out_predicted_token, ++ out_error); ++ if (status != SKIPPY_STATUS_OK || out_predicted_token == nullptr || *out_predicted_token < 0) { ++ return status; ++ } ++ ++ return skippy_mtp_propose_next(session, *out_predicted_token, out_mtp_draft, out_error); ++} ++ + enum skippy_status skippy_decode_step_frame_batch_sampled( + struct skippy_session * const * sessions, + const llama_token * token_ids, +@@ -3555,7 +3824,7 @@ enum skippy_status skippy_verify_tokens_frame( + + if (skippy_is_filtered(session) && session->stage_model->config.layer_start > 0) { + status = session->stage_model->config.include_output ? +- skippy_verify_activation_frame(session, input_desc, input_payload, token_count, out_error) : ++ skippy_verify_activation_frame(session, input_desc, input_payload, token_ids, token_count, out_error) : + skippy_decode_activation_frame(session, input_desc, input_payload, token_ids, nullptr, 0, token_count, false, out_error); + } else { + status = session->stage_model->config.include_output ? +@@ -4077,15 +4346,13 @@ enum skippy_status skippy_trim_session( + } + const llama_pos p0 = static_cast(token_count); + if (auto * hybrid = dynamic_cast(memory)) { +- llama_kv_cache * kv = hybrid->get_mem_attn(); +- if (kv != nullptr && !kv->seq_rm(session->seq_id, p0, -1)) { +- skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to trim hybrid attention KV suffix"); ++ if (!hybrid->seq_rm(session->seq_id, p0, -1)) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to trim hybrid memory suffix"); + return SKIPPY_STATUS_RUNTIME_ERROR; + } + } else if (auto * hybrid_iswa = dynamic_cast(memory)) { +- llama_kv_cache_iswa * kv = hybrid_iswa->get_mem_attn(); +- if (kv != nullptr && !kv->seq_rm(session->seq_id, p0, -1)) { +- skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to trim hybrid ISWA attention KV suffix"); ++ if (!hybrid_iswa->seq_rm(session->seq_id, p0, -1)) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to trim hybrid ISWA memory suffix"); + return SKIPPY_STATUS_RUNTIME_ERROR; + } + } else if (auto * kv = dynamic_cast(memory)) { +@@ -4105,6 +4372,7 @@ enum skippy_status skippy_trim_session( + if (session->signal_history.size() > token_count) { + session->signal_history.resize(static_cast(token_count)); + } ++ skippy_mtp_clear_session_state(session); + session->ctx->synchronize(); + + return skippy_success(out_error); +-- +2.54.0 (Apple Git-156) diff --git a/third_party/llama.cpp/patches/0102-Add-Skippy-batched-MTP-verification-drafts.patch b/third_party/llama.cpp/patches/0102-Add-Skippy-batched-MTP-verification-drafts.patch new file mode 100644 index 0000000000..3fda6f217c --- /dev/null +++ b/third_party/llama.cpp/patches/0102-Add-Skippy-batched-MTP-verification-drafts.patch @@ -0,0 +1,154 @@ +From d1d28b4b904b0b870ac1160a0a59cae90b6529c8 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Wed, 17 Jun 2026 07:39:00 +1000 +Subject: [PATCH 92/94] Add Skippy batched MTP verification drafts + +--- + include/skippy.h | 16 +++++++ + src/skippy.cpp | 82 ++++++++++++++++++++++++++++++++--- + 2 files changed, 92 insertions(+), 6 deletions(-) + +diff --git a/include/skippy.h b/include/skippy.h +index c4da8e28..8360f40d 100644 +--- a/include/skippy.h ++++ b/include/skippy.h +@@ -419,6 +419,22 @@ LLAMA_API enum skippy_status skippy_verify_tokens_frame( + size_t * out_token_count, + struct skippy_error ** out_error); + ++LLAMA_API enum skippy_status skippy_verify_tokens_frame_sampled( ++ struct skippy_session * session, ++ const llama_token * token_ids, ++ size_t token_count, ++ const struct skippy_sampling_config * sampling, ++ const struct skippy_activation_desc * input_desc, ++ const void * input_payload, ++ struct skippy_activation_desc * output_desc, ++ void * output_payload, ++ size_t output_payload_capacity, ++ size_t * out_output_payload_bytes, ++ llama_token * output_tokens, ++ size_t output_token_capacity, ++ size_t * out_token_count, ++ struct skippy_error ** out_error); ++ + LLAMA_API enum skippy_status skippy_session_copy_output_activation_frame( + struct skippy_session * session, + size_t token_count, +diff --git a/src/skippy.cpp b/src/skippy.cpp +index f8c07856..8ecb890c 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -1488,17 +1488,18 @@ static llama_token skippy_greedy_sample(skippy_session * session) { + return skippy_greedy_sample_ith(session, -1); + } + +-static bool skippy_compute_token_signal( +- skippy_session * session, ++static bool skippy_compute_token_signal_context( ++ const llama_model * model, ++ llama_context * ctx, + int32_t logits_index, + skippy_token_signal * out_signal) { +- if (session == nullptr || session->ctx == nullptr || session->stage_model == nullptr || session->stage_model->model == nullptr || out_signal == nullptr) { ++ if (model == nullptr || ctx == nullptr || out_signal == nullptr) { + return false; + } + +- const llama_vocab * vocab = llama_model_get_vocab(session->stage_model->model); ++ const llama_vocab * vocab = llama_model_get_vocab(model); + const int32_t n_vocab = llama_vocab_n_tokens(vocab); +- const float * logits = llama_get_logits_ith(session->ctx, logits_index); ++ const float * logits = llama_get_logits_ith(ctx, logits_index); + if (logits == nullptr || n_vocab <= 0) { + return false; + } +@@ -1555,6 +1556,20 @@ static bool skippy_compute_token_signal( + return true; + } + ++static bool skippy_compute_token_signal( ++ skippy_session * session, ++ int32_t logits_index, ++ skippy_token_signal * out_signal) { ++ if (session == nullptr || session->stage_model == nullptr) { ++ return false; ++ } ++ return skippy_compute_token_signal_context( ++ session->stage_model->model, ++ session->ctx, ++ logits_index, ++ out_signal); ++} ++ + static void skippy_record_signal(skippy_session * session, int32_t logits_index) { + skippy_token_signal signal = {}; + if (skippy_compute_token_signal(session, logits_index, &signal)) { +@@ -3771,6 +3787,38 @@ enum skippy_status skippy_verify_tokens_frame( + size_t output_token_capacity, + size_t * out_token_count, + struct skippy_error ** out_error) { ++ return skippy_verify_tokens_frame_sampled( ++ session, ++ token_ids, ++ token_count, ++ nullptr, ++ input_desc, ++ input_payload, ++ output_desc, ++ output_payload, ++ output_payload_capacity, ++ out_output_payload_bytes, ++ output_tokens, ++ output_token_capacity, ++ out_token_count, ++ out_error); ++} ++ ++enum skippy_status skippy_verify_tokens_frame_sampled( ++ struct skippy_session * session, ++ const llama_token * token_ids, ++ size_t token_count, ++ const struct skippy_sampling_config * sampling, ++ const struct skippy_activation_desc * input_desc, ++ const void * input_payload, ++ struct skippy_activation_desc * output_desc, ++ void * output_payload, ++ size_t output_payload_capacity, ++ size_t * out_output_payload_bytes, ++ llama_token * output_tokens, ++ size_t output_token_capacity, ++ size_t * out_token_count, ++ struct skippy_error ** out_error) { + if (out_token_count != nullptr) { + *out_token_count = 0; + } +@@ -3843,7 +3891,25 @@ enum skippy_status skippy_verify_tokens_frame( + if (session->stage_model->config.include_output) { + const int32_t n_tokens = static_cast(token_count); + for (int32_t i = 0; i < n_tokens; ++i) { +- output_tokens[i] = skippy_greedy_sample_ith(session, i); ++ output_tokens[i] = skippy_sample_token_ith(session, sampling, i); ++ } ++ if (output_token_capacity >= token_count + 2 && token_count > 0) { ++ skippy_native_mtp_draft mtp_draft = {}; ++ status = skippy_mtp_propose_next( ++ session, ++ output_tokens[token_count - 1], ++ &mtp_draft, ++ out_error); ++ if (status != SKIPPY_STATUS_OK) { ++ return status; ++ } ++ if (mtp_draft.available) { ++ output_tokens[token_count] = mtp_draft.token_id; ++ output_tokens[token_count + 1] = static_cast(std::min( ++ std::max(mtp_draft.proposal_compute_us, 0), ++ std::numeric_limits::max())); ++ *out_token_count = token_count + 2; ++ } + } + } + +-- +2.54.0 (Apple Git-156) diff --git a/third_party/llama.cpp/patches/0103-Align-Skippy-sampling-and-detokenization-parity.patch b/third_party/llama.cpp/patches/0103-Align-Skippy-sampling-and-detokenization-parity.patch new file mode 100644 index 0000000000..4d5fcf1006 --- /dev/null +++ b/third_party/llama.cpp/patches/0103-Align-Skippy-sampling-and-detokenization-parity.patch @@ -0,0 +1,256 @@ +From 7927825c5fb5eb96563f2f9dfd0149ee7960422d Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Wed, 17 Jun 2026 07:39:00 +1000 +Subject: [PATCH 93/94] Align Skippy sampling and detokenization parity + +--- + src/skippy.cpp | 155 ++++++++++++++++++++++++++++++++++++++----------- + 1 file changed, 121 insertions(+), 34 deletions(-) + +diff --git a/src/skippy.cpp b/src/skippy.cpp +index 04c0afb4..e4a854e2 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -72,6 +72,8 @@ struct skippy_session { + std::vector signal_history; + llama_sampler * sampling_chain = nullptr; + llama_sampler * grammar_sampler = nullptr; ++ skippy_sampling_config sampling_config = {}; ++ bool sampling_config_valid = false; + std::string chat_sampling_metadata; + uint64_t grammar_generated_start = 0; + size_t sampling_accepted_token_count = 0; +@@ -1474,6 +1476,7 @@ static llama_token skippy_greedy_sample_context( + } + const llama_vocab * vocab = llama_model_get_vocab(model); + const int32_t n_vocab = llama_vocab_n_tokens(vocab); ++ llama_synchronize(ctx); + const float * logits = llama_get_logits_ith(ctx, index); + if (logits == nullptr) { + return 0; +@@ -1596,6 +1599,8 @@ static void skippy_clear_chat_sampling(skippy_session * session) { + llama_sampler_free(session->grammar_sampler); + session->sampling_chain = nullptr; + session->grammar_sampler = nullptr; ++ session->sampling_config = {}; ++ session->sampling_config_valid = false; + session->chat_sampling_metadata.clear(); + session->grammar_generated_start = 0; + session->sampling_accepted_token_count = 0; +@@ -1634,6 +1639,34 @@ static bool skippy_sampling_enabled(const skippy_sampling_config * sampling) { + return true; + } + ++static bool skippy_sampling_configs_equal( ++ const skippy_sampling_config & lhs, ++ const skippy_sampling_config & rhs) { ++ if (lhs.version != rhs.version || ++ lhs.flags != rhs.flags || ++ lhs.seed != rhs.seed || ++ lhs.top_k != rhs.top_k || ++ lhs.penalty_last_n != rhs.penalty_last_n || ++ lhs.temperature != rhs.temperature || ++ lhs.top_p != rhs.top_p || ++ lhs.presence_penalty != rhs.presence_penalty || ++ lhs.frequency_penalty != rhs.frequency_penalty || ++ lhs.repeat_penalty != rhs.repeat_penalty || ++ lhs.logit_bias_count != rhs.logit_bias_count || ++ lhs.min_p != rhs.min_p) { ++ return false; ++ } ++ ++ const uint32_t logit_bias_count = std::min(lhs.logit_bias_count, SKIPPY_MAX_LOGIT_BIAS); ++ for (uint32_t i = 0; i < logit_bias_count; ++i) { ++ if (lhs.logit_bias[i].token != rhs.logit_bias[i].token || ++ lhs.logit_bias[i].bias != rhs.logit_bias[i].bias) { ++ return false; ++ } ++ } ++ return true; ++} ++ + static llama_sampler * skippy_build_sampling_chain( + skippy_session * session, + const skippy_sampling_config * sampling) { +@@ -1673,21 +1706,40 @@ static llama_sampler * skippy_build_sampling_chain( + llama_sampler_chain_add(sampler, llama_sampler_init_top_k(sampling->top_k)); + } + if (sampling->top_p > 0.0f && sampling->top_p < 1.0f) { +- llama_sampler_chain_add(sampler, llama_sampler_init_top_p(sampling->top_p, 1)); ++ llama_sampler_chain_add(sampler, llama_sampler_init_top_p(sampling->top_p, 0)); + } + if (sampling->min_p > 0.0f && sampling->min_p < 1.0f) { +- llama_sampler_chain_add(sampler, llama_sampler_init_min_p(sampling->min_p, 1)); ++ llama_sampler_chain_add(sampler, llama_sampler_init_min_p(sampling->min_p, 0)); + } +- if (sampling->temperature != 1.0f) { +- llama_sampler_chain_add(sampler, llama_sampler_init_temp(sampling->temperature)); ++ llama_sampler_chain_add(sampler, llama_sampler_init_temp_ext(sampling->temperature, 0.0f, 1.0f)); ++ const uint32_t seed = sampling->seed == 0 ? LLAMA_DEFAULT_SEED : sampling->seed; ++ llama_sampler_chain_add(sampler, llama_sampler_init_dist(seed)); ++ return sampler; ++} ++ ++static bool skippy_ensure_plain_sampling_chain( ++ skippy_session * session, ++ const skippy_sampling_config * sampling) { ++ if (session == nullptr || !skippy_sampling_enabled(sampling)) { ++ return false; + } +- if (sampling->temperature <= 0.0f) { +- llama_sampler_chain_add(sampler, llama_sampler_init_greedy()); +- } else { +- const uint32_t seed = sampling->seed == 0 ? LLAMA_DEFAULT_SEED : sampling->seed + static_cast(session->n_past); +- llama_sampler_chain_add(sampler, llama_sampler_init_dist(seed)); ++ ++ if (session->sampling_chain != nullptr && ++ session->chat_sampling_metadata.empty() && ++ session->sampling_config_valid && ++ skippy_sampling_configs_equal(session->sampling_config, *sampling)) { ++ return true; + } +- return sampler; ++ ++ skippy_clear_chat_sampling(session); ++ session->sampling_chain = skippy_build_sampling_chain(session, sampling); ++ if (session->sampling_chain == nullptr) { ++ return false; ++ } ++ session->sampling_config = *sampling; ++ session->sampling_config_valid = true; ++ session->sampling_accepted_token_count = 0; ++ return true; + } + + static bool skippy_init_chat_grammar_sampler( +@@ -1800,6 +1852,7 @@ static void skippy_sync_chat_sampling_history(skippy_session * session) { + + static llama_token skippy_chat_sample_token(skippy_session * session, int32_t logits_index) { + skippy_sync_chat_sampling_history(session); ++ llama_synchronize(session->ctx); + + const llama_vocab * vocab = llama_model_get_vocab(session->stage_model->model); + const int32_t n_vocab = llama_vocab_n_tokens(vocab); +@@ -1840,11 +1893,21 @@ static llama_token skippy_sample_token_ith( + const skippy_sampling_config * sampling, + int32_t logits_index) { + if (session != nullptr && session->sampling_chain != nullptr) { +- return skippy_chat_sample_token(session, logits_index); ++ if (!session->chat_sampling_metadata.empty() || ++ (skippy_sampling_enabled(sampling) && ++ session->sampling_config_valid && ++ skippy_sampling_configs_equal(session->sampling_config, *sampling))) { ++ return skippy_chat_sample_token(session, logits_index); ++ } ++ skippy_clear_chat_sampling(session); + } + if (!skippy_sampling_enabled(sampling)) { + return skippy_greedy_sample_ith(session, logits_index); + } ++ if (skippy_ensure_plain_sampling_chain(session, sampling)) { ++ return skippy_chat_sample_token(session, logits_index); ++ } ++ llama_synchronize(session->ctx); + + llama_sampler_chain_params chain_params = llama_sampler_chain_default_params(); + llama_sampler * sampler = llama_sampler_chain_init(chain_params); +@@ -1877,20 +1940,14 @@ static llama_token skippy_sample_token_ith( + llama_sampler_chain_add(sampler, llama_sampler_init_top_k(sampling->top_k)); + } + if (sampling->top_p > 0.0f && sampling->top_p < 1.0f) { +- llama_sampler_chain_add(sampler, llama_sampler_init_top_p(sampling->top_p, 1)); ++ llama_sampler_chain_add(sampler, llama_sampler_init_top_p(sampling->top_p, 0)); + } + if (sampling->min_p > 0.0f && sampling->min_p < 1.0f) { +- llama_sampler_chain_add(sampler, llama_sampler_init_min_p(sampling->min_p, 1)); +- } +- if (sampling->temperature != 1.0f) { +- llama_sampler_chain_add(sampler, llama_sampler_init_temp(sampling->temperature)); +- } +- if (sampling->temperature <= 0.0f) { +- llama_sampler_chain_add(sampler, llama_sampler_init_greedy()); +- } else { +- const uint32_t seed = sampling->seed == 0 ? LLAMA_DEFAULT_SEED : sampling->seed + static_cast(session->n_past); +- llama_sampler_chain_add(sampler, llama_sampler_init_dist(seed)); ++ llama_sampler_chain_add(sampler, llama_sampler_init_min_p(sampling->min_p, 0)); + } ++ llama_sampler_chain_add(sampler, llama_sampler_init_temp_ext(sampling->temperature, 0.0f, 1.0f)); ++ const uint32_t seed = sampling->seed == 0 ? LLAMA_DEFAULT_SEED : sampling->seed; ++ llama_sampler_chain_add(sampler, llama_sampler_init_dist(seed)); + + for (const llama_token token : session->token_history) { + llama_sampler_accept(sampler, token); +@@ -2951,6 +3008,10 @@ enum skippy_status skippy_session_configure_chat_sampling( + return SKIPPY_STATUS_RUNTIME_ERROR; + } + session->chat_sampling_metadata = metadata_json; ++ if (sampling != nullptr) { ++ session->sampling_config = *sampling; ++ session->sampling_config_valid = true; ++ } + session->grammar_generated_start = prompt_token_count; + session->sampling_accepted_token_count = 0; + if (!skippy_init_chat_grammar_sampler(session, metadata, out_error)) { +@@ -4904,21 +4965,47 @@ enum skippy_status skippy_detokenize( + } + + const llama_vocab * vocab = llama_model_get_vocab(model->model); +- const int32_t result = llama_detokenize( +- vocab, +- tokens, +- static_cast(token_count), +- output_text, +- static_cast(output_text_capacity), +- true, +- false); +- if (result < 0) { +- *out_text_bytes = static_cast(-result); ++ std::string text; ++ for (size_t i = 0; i < token_count; ++i) { ++ std::string piece; ++ piece.resize(piece.capacity()); ++ int32_t piece_bytes = llama_token_to_piece( ++ vocab, ++ tokens[i], ++ piece.data(), ++ static_cast(piece.size()), ++ 0, ++ true); ++ if (piece_bytes < 0) { ++ piece.resize(static_cast(-piece_bytes)); ++ piece_bytes = llama_token_to_piece( ++ vocab, ++ tokens[i], ++ piece.data(), ++ static_cast(piece.size()), ++ 0, ++ true); ++ } ++ if (piece_bytes < 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_BUFFER_TOO_SMALL, "token piece output buffer is too small"); ++ return SKIPPY_STATUS_BUFFER_TOO_SMALL; ++ } ++ piece.resize(static_cast(piece_bytes)); ++ text += piece; ++ } ++ ++ *out_text_bytes = text.size(); ++ if (output_text_capacity < text.size()) { + skippy_set_error(out_error, SKIPPY_STATUS_BUFFER_TOO_SMALL, "text output buffer is too small"); + return SKIPPY_STATUS_BUFFER_TOO_SMALL; + } +- +- *out_text_bytes = static_cast(result); ++ if (!text.empty() && output_text == nullptr) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "output_text is required when output capacity is sufficient"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ if (!text.empty()) { ++ std::memcpy(output_text, text.data(), text.size()); ++ } + return skippy_success(out_error); + } + +-- +2.54.0 (Apple Git-156) + diff --git a/third_party/llama.cpp/patches/0104-Remove-unused-Skippy-ABI-exports.patch b/third_party/llama.cpp/patches/0104-Remove-unused-Skippy-ABI-exports.patch new file mode 100644 index 0000000000..a71d3025d5 --- /dev/null +++ b/third_party/llama.cpp/patches/0104-Remove-unused-Skippy-ABI-exports.patch @@ -0,0 +1,238 @@ +From 7a2e03abe085ebbd12d429d54e54c51dbb227b23 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Wed, 17 Jun 2026 07:39:07 +1000 +Subject: [PATCH 94/94] Remove unused Skippy ABI exports + +--- + include/skippy.h | 41 ----------------- + include/skippy/common.h | 2 - + src/skippy.cpp | 99 ----------------------------------------- + 3 files changed, 142 deletions(-) + +diff --git a/include/skippy.h b/include/skippy.h +index 8360f40d..a74ad680 100644 +--- a/include/skippy.h ++++ b/include/skippy.h +@@ -193,9 +193,6 @@ LLAMA_API struct llama_context * skippy_session_llama_context( + LLAMA_API int32_t skippy_session_position( + const struct skippy_session * session); + +-LLAMA_API int32_t skippy_session_native_seq_id( +- const struct skippy_session * session); +- + LLAMA_API int32_t skippy_session_batch_size( + const struct skippy_session * session); + +@@ -254,17 +251,6 @@ LLAMA_API enum skippy_status skippy_prefill_chunk( + size_t * out_output_activation_bytes, + struct skippy_error ** out_error); + +-LLAMA_API enum skippy_status skippy_decode_step( +- struct skippy_session * session, +- llama_token token_id, +- const void * input_activation, +- size_t input_activation_bytes, +- void * output_activation, +- size_t output_activation_capacity, +- size_t * out_output_activation_bytes, +- llama_token * out_predicted_token, +- struct skippy_error ** out_error); +- + LLAMA_API enum skippy_status skippy_verify_tokens( + struct skippy_session * session, + const llama_token * token_ids, +@@ -351,18 +337,6 @@ LLAMA_API enum skippy_status skippy_prefill_chunk_frame_sampled_with_positions( + llama_token * out_predicted_token, + struct skippy_error ** out_error); + +-LLAMA_API enum skippy_status skippy_decode_step_frame( +- struct skippy_session * session, +- llama_token token_id, +- const struct skippy_activation_desc * input_desc, +- const void * input_payload, +- struct skippy_activation_desc * output_desc, +- void * output_payload, +- size_t output_payload_capacity, +- size_t * out_output_payload_bytes, +- llama_token * out_predicted_token, +- struct skippy_error ** out_error); +- + LLAMA_API enum skippy_status skippy_decode_step_frame_sampled( + struct skippy_session * session, + llama_token token_id, +@@ -405,21 +379,6 @@ LLAMA_API enum skippy_status skippy_decode_step_frame_batch_sampled( + size_t request_count, + struct skippy_error ** out_error); + +-LLAMA_API enum skippy_status skippy_verify_tokens_frame( +- struct skippy_session * session, +- const llama_token * token_ids, +- size_t token_count, +- const struct skippy_activation_desc * input_desc, +- const void * input_payload, +- struct skippy_activation_desc * output_desc, +- void * output_payload, +- size_t output_payload_capacity, +- size_t * out_output_payload_bytes, +- llama_token * output_tokens, +- size_t output_token_capacity, +- size_t * out_token_count, +- struct skippy_error ** out_error); +- + LLAMA_API enum skippy_status skippy_verify_tokens_frame_sampled( + struct skippy_session * session, + const llama_token * token_ids, +diff --git a/include/skippy/common.h b/include/skippy/common.h +index cf97f7ab..cfb4a8fa 100644 +--- a/include/skippy/common.h ++++ b/include/skippy/common.h +@@ -164,8 +164,6 @@ LLAMA_API struct skippy_abi_version skippy_abi_version(void); + + LLAMA_API uint64_t skippy_abi_features(void); + +-LLAMA_API const char * skippy_status_string(enum skippy_status status); +- + LLAMA_API void skippy_error_free(struct skippy_error * error); + + #ifdef __cplusplus +diff --git a/src/skippy.cpp b/src/skippy.cpp +index e4a854e2..01cfd3c9 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -2308,21 +2308,6 @@ uint64_t skippy_abi_features(void) { + SKIPPY_FEATURE_NATIVE_MTP_N1; + } + +-const char * skippy_status_string(enum skippy_status status) { +- switch (status) { +- case SKIPPY_STATUS_OK: return "ok"; +- case SKIPPY_STATUS_ERROR: return "error"; +- case SKIPPY_STATUS_INVALID_ARGUMENT: return "invalid_argument"; +- case SKIPPY_STATUS_UNSUPPORTED: return "unsupported"; +- case SKIPPY_STATUS_BUFFER_TOO_SMALL: return "buffer_too_small"; +- case SKIPPY_STATUS_IO_ERROR: return "io_error"; +- case SKIPPY_STATUS_MODEL_ERROR: return "model_error"; +- case SKIPPY_STATUS_RUNTIME_ERROR: return "runtime_error"; +- } +- +- return "unknown"; +-} +- + void skippy_error_free(struct skippy_error * error) { + if (error == nullptr) { + return; +@@ -2904,11 +2889,6 @@ int32_t skippy_session_position( + return session != nullptr ? session->n_past : -1; + } + +-int32_t skippy_session_native_seq_id( +- const struct skippy_session * session) { +- return session != nullptr ? session->seq_id : -1; +-} +- + int32_t skippy_session_batch_size( + const struct skippy_session * session) { + return session != nullptr && session->ctx != nullptr ? llama_n_batch(session->ctx) : 0; +@@ -3208,29 +3188,6 @@ enum skippy_status skippy_prefill_chunk( + return skippy_decode_tokens(session, token_ids, token_count, false, out_error); + } + +-enum skippy_status skippy_decode_step( +- struct skippy_session * session, +- llama_token token_id, +- const void * input_activation, +- size_t input_activation_bytes, +- void * output_activation, +- size_t output_activation_capacity, +- size_t * out_output_activation_bytes, +- llama_token * out_predicted_token, +- struct skippy_error ** out_error) { +- return skippy_decode_step_sampled( +- session, +- token_id, +- nullptr, +- input_activation, +- input_activation_bytes, +- output_activation, +- output_activation_capacity, +- out_output_activation_bytes, +- out_predicted_token, +- out_error); +-} +- + enum skippy_status skippy_decode_step_sampled( + struct skippy_session * session, + llama_token token_id, +@@ -3567,31 +3524,6 @@ enum skippy_status skippy_prefill_chunk_frame_sampled_with_positions( + out_error); + } + +-enum skippy_status skippy_decode_step_frame( +- struct skippy_session * session, +- llama_token token_id, +- const struct skippy_activation_desc * input_desc, +- const void * input_payload, +- struct skippy_activation_desc * output_desc, +- void * output_payload, +- size_t output_payload_capacity, +- size_t * out_output_payload_bytes, +- llama_token * out_predicted_token, +- struct skippy_error ** out_error) { +- return skippy_decode_step_frame_sampled( +- session, +- token_id, +- nullptr, +- input_desc, +- input_payload, +- output_desc, +- output_payload, +- output_payload_capacity, +- out_output_payload_bytes, +- out_predicted_token, +- out_error); +-} +- + enum skippy_status skippy_decode_step_frame_sampled( + struct skippy_session * session, + llama_token token_id, +@@ -3859,37 +3791,6 @@ enum skippy_status skippy_decode_step_frame_batch_sampled( + return skippy_success(out_error); + } + +-enum skippy_status skippy_verify_tokens_frame( +- struct skippy_session * session, +- const llama_token * token_ids, +- size_t token_count, +- const struct skippy_activation_desc * input_desc, +- const void * input_payload, +- struct skippy_activation_desc * output_desc, +- void * output_payload, +- size_t output_payload_capacity, +- size_t * out_output_payload_bytes, +- llama_token * output_tokens, +- size_t output_token_capacity, +- size_t * out_token_count, +- struct skippy_error ** out_error) { +- return skippy_verify_tokens_frame_sampled( +- session, +- token_ids, +- token_count, +- nullptr, +- input_desc, +- input_payload, +- output_desc, +- output_payload, +- output_payload_capacity, +- out_output_payload_bytes, +- output_tokens, +- output_token_capacity, +- out_token_count, +- out_error); +-} +- + enum skippy_status skippy_verify_tokens_frame_sampled( + struct skippy_session * session, + const llama_token * token_ids, +-- +2.54.0 (Apple Git-156) + diff --git a/third_party/llama.cpp/patches/0105-Fix-Skippy-MTP-stage-layer-count.patch b/third_party/llama.cpp/patches/0105-Fix-Skippy-MTP-stage-layer-count.patch new file mode 100644 index 0000000000..5fb42cec65 --- /dev/null +++ b/third_party/llama.cpp/patches/0105-Fix-Skippy-MTP-stage-layer-count.patch @@ -0,0 +1,68 @@ +From d207849787a221123cf1ad195e23f4cc5454e365 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Wed, 17 Jun 2026 09:05:21 +1000 +Subject: [PATCH] Fix Skippy MTP stage layer count + +--- + src/models/deepseek2.cpp | 4 +++- + src/skippy.cpp | 15 +++++++++++++-- + 2 files changed, 16 insertions(+), 3 deletions(-) + +diff --git a/src/models/deepseek2.cpp b/src/models/deepseek2.cpp +index d0ddb3c5..180fb082 100644 +--- a/src/models/deepseek2.cpp ++++ b/src/models/deepseek2.cpp +@@ -194,7 +194,9 @@ llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_p + const skippy_graph_filter & stage_filter = skippy_graph_get_filter(); + const bool stage_filtered = stage_filter.enabled; + const int il_start = stage_filtered ? stage_filter.layer_start : 0; +- const int il_end = stage_filtered ? stage_filter.layer_end : effective_n_layers; ++ const int il_end = stage_filtered ? ++ std::min(stage_filter.layer_end, effective_n_layers) : ++ effective_n_layers; + + // {n_embd, n_tokens} + inpL = build_inp_embd(stage_filtered && il_start > 0 ? nullptr : model.tok_embd); +diff --git a/src/skippy.cpp b/src/skippy.cpp +index d49946cf..9f05153f 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -181,6 +181,17 @@ static enum skippy_status skippy_success(skippy_error ** out_error) { + return SKIPPY_STATUS_OK; + } + ++static int32_t skippy_stage_layer_count(const llama_model * model) { ++ if (model == nullptr) { ++ return 0; ++ } ++ const auto & hparams = model->hparams; ++ if (hparams.n_layer_nextn > 0) { ++ return static_cast(hparams.n_layer_all); ++ } ++ return llama_model_n_layer(model); ++} ++ + static enum skippy_backend_device_type skippy_backend_device_type_from_ggml( + enum ggml_backend_dev_type type) { + switch (type) { +@@ -2417,7 +2428,7 @@ static enum skippy_status skippy_finish_model_open( + struct skippy_model ** out_model, + struct skippy_error ** out_error) { + if (config != nullptr && config->filter_tensors_on_load) { +- const int32_t n_layer = llama_model_n_layer(model); ++ const int32_t n_layer = skippy_stage_layer_count(model); + if (model->arch != LLM_ARCH_LLAMA && + model->arch != LLM_ARCH_AFMOE && + model->arch != LLM_ARCH_APERTUS && +@@ -3976,7 +3987,7 @@ static enum skippy_status skippy_validate_state_range( + const skippy_runtime_config & config = session->stage_model->config; + const int32_t expected_layer_start = config.filter_tensors_on_load ? config.layer_start : 0; + const int32_t expected_layer_end = config.filter_tensors_on_load ? +- config.layer_end : llama_model_n_layer(session->stage_model->model); ++ config.layer_end : skippy_stage_layer_count(session->stage_model->model); + if (layer_start != expected_layer_start || layer_end != expected_layer_end) { + skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "state range must match the session layer range"); + return SKIPPY_STATUS_INVALID_ARGUMENT; +-- +2.54.0 (Apple Git-156) + diff --git a/third_party/llama.cpp/patches/0106-Avoid-copying-VerifySpan-activation-inputs.patch b/third_party/llama.cpp/patches/0106-Avoid-copying-VerifySpan-activation-inputs.patch new file mode 100644 index 0000000000..96a454c911 --- /dev/null +++ b/third_party/llama.cpp/patches/0106-Avoid-copying-VerifySpan-activation-inputs.patch @@ -0,0 +1,66 @@ +From bd53d862a053cceea05d8df0c196fae4bafc1f2c Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Wed, 17 Jun 2026 20:21:16 +1000 +Subject: [PATCH] Avoid copying VerifySpan activation inputs + +--- + src/skippy.cpp | 35 +++++++++++++++++++++++++++++++---- + 1 file changed, 31 insertions(+), 4 deletions(-) + +diff --git a/src/skippy.cpp b/src/skippy.cpp +index 56bb1ba2..b357620e 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -2293,10 +2293,35 @@ static enum skippy_status skippy_verify_activation_frame( + const int32_t n_tokens = static_cast(token_count); + const int32_t n_embd = llama_model_n_embd(session->stage_model->model); + const int32_t n_altup = static_cast(session->stage_model->model->hparams.n_altup); +- llama_batch batch = llama_batch_init(n_tokens, n_embd, 1); +- batch.n_tokens = n_tokens; + const size_t hidden_bytes = skippy_activation_hidden_bytes(session, token_count); +- std::memcpy(batch.embd, input_payload, hidden_bytes); ++ const bool alias_input_payload = input_desc->flags == 0; ++ ++ llama_batch batch = {}; ++ std::vector pos_storage; ++ std::vector n_seq_id_storage; ++ std::vector seq_id_0; ++ std::vector seq_id_storage; ++ std::vector logits_storage; ++ if (alias_input_payload) { ++ pos_storage.resize(n_tokens); ++ n_seq_id_storage.resize(n_tokens); ++ seq_id_0.assign(1, session->seq_id); ++ seq_id_storage.resize(n_tokens, seq_id_0.data()); ++ logits_storage.resize(n_tokens, 1); ++ batch = { ++ /*n_tokens =*/ n_tokens, ++ /*token =*/ nullptr, ++ /*embd =*/ const_cast(static_cast(input_payload)), ++ /*pos =*/ pos_storage.data(), ++ /*n_seq_id =*/ n_seq_id_storage.data(), ++ /*seq_id =*/ seq_id_storage.data(), ++ /*logits =*/ logits_storage.data(), ++ }; ++ } else { ++ batch = llama_batch_init(n_tokens, n_embd, 1); ++ std::memcpy(batch.embd, input_payload, hidden_bytes); ++ } ++ batch.n_tokens = n_tokens; + + for (int32_t i = 0; i < n_tokens; ++i) { + batch.pos[i] = session->n_past + i; +@@ -2309,7 +2334,9 @@ static enum skippy_status skippy_verify_activation_frame( + skippy_gemma3n_altup_scope gemma3n_altup_scope(input_desc, input_payload, n_embd, n_altup); + const llama_pos token_start = session->n_past; + enum skippy_status status = skippy_decode_batch(session, batch, token_count, out_error); +- llama_batch_free(batch); ++ if (!alias_input_payload) { ++ llama_batch_free(batch); ++ } + if (status == SKIPPY_STATUS_OK && token_ids != nullptr) { + status = skippy_mtp_sync_target_tokens(session, token_ids, token_count, token_start, out_error); + } +-- +2.54.0 (Apple Git-156) + diff --git a/third_party/llama.cpp/patches/0107-Add-Skippy-greedy-sampling-fast-path.patch b/third_party/llama.cpp/patches/0107-Add-Skippy-greedy-sampling-fast-path.patch new file mode 100644 index 0000000000..4b14d8af2d --- /dev/null +++ b/third_party/llama.cpp/patches/0107-Add-Skippy-greedy-sampling-fast-path.patch @@ -0,0 +1,72 @@ +From 2d45dd40bbf9cd30c54182a548fafe1fb08d32e8 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Wed, 17 Jun 2026 21:29:11 +1000 +Subject: [PATCH] Add Skippy greedy sampling fast path + +--- + src/skippy.cpp | 36 ++++++++++++++++++++++++++++++++++++ + 1 file changed, 36 insertions(+) + +diff --git a/src/skippy.cpp b/src/skippy.cpp +index b357620e..dce41495 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -1239,6 +1239,21 @@ static bool skippy_mtp_available(const skippy_session * session) { + session->stage_model->config.include_output; + } + ++static bool skippy_env_enabled(const char * name) { ++ const char * value = std::getenv(name); ++ if (value == nullptr || value[0] == '\0') { ++ return false; ++ } ++ return std::strcmp(value, "0") != 0 && ++ std::strcmp(value, "false") != 0 && ++ std::strcmp(value, "off") != 0 && ++ std::strcmp(value, "no") != 0; ++} ++ ++static bool skippy_mtp_greedy_sampling_fastpath_enabled() { ++ return skippy_env_enabled("SKIPPY_NATIVE_MTP_GREEDY_SAMPLING_FASTPATH"); ++} ++ + static void skippy_mtp_clear_session_state(skippy_session * session) { + if (session == nullptr) { + return; +@@ -1760,6 +1764,18 @@ static bool skippy_sampling_enabled(const skippy_sampling_config * sampling) { + return true; + } + ++static bool skippy_sampling_is_greedy_equivalent(const skippy_sampling_config * sampling) { ++ if (!skippy_sampling_enabled(sampling)) { ++ return true; ++ } ++ const float repeat_penalty = sampling->repeat_penalty == 0.0f ? 1.0f : sampling->repeat_penalty; ++ return sampling->temperature <= 0.0f && ++ sampling->presence_penalty == 0.0f && ++ sampling->frequency_penalty == 0.0f && ++ repeat_penalty == 1.0f && ++ sampling->logit_bias_count == 0; ++} ++ + static bool skippy_sampling_configs_equal( + const skippy_sampling_config & lhs, + const skippy_sampling_config & rhs) { +@@ -2013,6 +2029,15 @@ static llama_token skippy_sample_token_ith( + skippy_session * session, + const skippy_sampling_config * sampling, + int32_t logits_index) { ++ if (skippy_mtp_greedy_sampling_fastpath_enabled() && ++ session != nullptr && ++ session->grammar_sampler == nullptr && ++ skippy_sampling_is_greedy_equivalent(sampling)) { ++ if (session->sampling_chain != nullptr) { ++ skippy_clear_chat_sampling(session); ++ } ++ return skippy_greedy_sample_ith(session, logits_index); ++ } + if (session != nullptr && session->sampling_chain != nullptr) { + if (!session->chat_sampling_metadata.empty() || + (skippy_sampling_enabled(sampling) && +-- +2.54.0 (Apple Git-156)