diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md new file mode 100644 index 0000000..61ba87b --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.md @@ -0,0 +1,29 @@ +--- +name: Bug report +about: Report an incorrect result, crash, or unexpected runtime behavior +title: "[Bug]: " +labels: bug +--- + +## What happened? + +Describe the observed behavior. + +## Reproduction + +Provide the smallest command or code sample that reproduces the issue. + +## Environment + +- OS and hardware: +- Rust version (`rustc --version`): +- Backend (`cpu` or `metal`): +- Does the issue reproduce with the CPU backend? + +## Expected behavior + +Describe the expected result. + +## Relevant logs + +Include the smallest relevant error output, backtrace, or parity diff. diff --git a/.github/ISSUE_TEMPLATE/feature_request.md b/.github/ISSUE_TEMPLATE/feature_request.md new file mode 100644 index 0000000..ff59ad2 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature_request.md @@ -0,0 +1,26 @@ +--- +name: Feature request +about: Propose an improvement to batch_forge +title: "[Feature]: " +labels: enhancement +--- + +## Problem + +What limitation or workflow does this improve? + +## Proposed change + +Describe the behavior or API you would like. + +## Alternatives considered + +Describe any workarounds or alternative approaches. + +## Success criteria + +What outcome would make this feature complete? + +## Compatibility or performance impact + +Note any expected platform, API, memory, or throughput tradeoffs. diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md new file mode 100644 index 0000000..4233da0 --- /dev/null +++ b/.github/pull_request_template.md @@ -0,0 +1,24 @@ +## Summary + +Describe the change and why it is needed. + +## Risks + +Note any compatibility, performance, or behavior risks. + +## Documentation + +- [ ] Public API docs and user-facing guides updated when applicable + +## Tests + +- [ ] Relevant tests added or updated when behavior changes + +## Validation + +- [ ] `cargo fmt --check` +- [ ] `cargo check` +- [ ] `cargo clippy --all-targets` +- [ ] `cargo test --lib` +- [ ] Metal parity check run when applicable (`cargo test --test parity -- --nocapture`) +- [ ] End-to-end GPT-2 check run when applicable (`cargo test --test gpt2_e2e -- --nocapture`) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..12e5145 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,56 @@ +name: CI + +on: + push: + branches: [main] + pull_request: + +env: + CARGO_TERM_COLOR: always + RUSTFLAGS: -D warnings + +jobs: + # Portable path: the CPU reference, loader, model, and async engine build and + # test on Linux (the Metal backend is cfg'd out on non-macOS targets). + lint-and-test: + name: fmt + clippy + test (ubuntu) + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@stable + with: + components: rustfmt, clippy + - uses: Swatinem/rust-cache@v2 + - name: asset downloader tests (no network) + run: python3 -m unittest discover -s python -p 'test_*.py' + - name: rustfmt + run: cargo fmt --all -- --check + - name: clippy + run: cargo clippy --all-targets + - name: test (CPU and CLI) + run: cargo test --lib --bins + - name: CLI smoke tests without model assets + run: | + cargo run --quiet --bin generate -- --help + cargo run --quiet --bin generate -- --version + if cargo run --quiet --bin generate -- --temperature NaN; then exit 1; fi + if cargo run --quiet --bin batch_forge -- --model /missing/checkpoint.safetensors; then exit 1; fi + + # Apple Silicon: full build incl. Metal kernels + CPU unit tests. The on-device + # CPU↔Metal parity suite (`cargo test --test parity`) needs a Metal device and + # is run locally; see docs/correctness.md. + macos-build: + name: build + test (macos) + runs-on: macos-14 + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@stable + with: + components: clippy + - uses: Swatinem/rust-cache@v2 + - name: clippy (incl. Metal backend) + run: cargo clippy --all-targets + - name: build + run: cargo build --release --all-targets + - name: test (CPU unit tests) + run: cargo test --lib --bins diff --git a/.gitignore b/.gitignore index 019d8ba..4691511 100644 --- a/.gitignore +++ b/.gitignore @@ -1,13 +1,12 @@ # Rust /target -Cargo.lock # Python venv/ __pycache__/ *.pyc -# ML Weights & Data +# ML Weights & Data (don't commit large binaries; generate them locally) *.safetensors *.bin *.pt @@ -15,3 +14,4 @@ __pycache__/ # MacOS .DS_Store +models/ diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..bc2e15e --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,18 @@ +# Contributing + +Small, focused changes are easiest to review. Before opening a pull request, run the checks that match the code you changed: + +```bash +cargo fmt --check +cargo check +cargo clippy --all-targets +cargo test --lib +``` + +If the formatting check reports changes, run `cargo fmt`, review the diff, then run the check again. + +The formatting, compile, lint, and library-test preflight is CPU-only and needs no model assets. Metal parity checks require Apple Silicon macOS, and the end-to-end GPT-2 check requires the downloaded model assets. See the [test command reference](docs/test-matrix.md) for the complete matrix and [troubleshooting guide](docs/troubleshooting.md) for setup notes. + +## Reporting issues + +Use the issue templates when reporting a bug or proposing an enhancement; their `bug` and `enhancement` labels help keep reports organized. For runtime bugs, include the backend, hardware, and smallest reproduction command. diff --git a/Cargo.lock b/Cargo.lock new file mode 100644 index 0000000..ce6957b --- /dev/null +++ b/Cargo.lock @@ -0,0 +1,491 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "batch_forge" +version = "0.1.0" +dependencies = [ + "bytemuck", + "memmap2", + "metal", + "safetensors", + "serde_json", + "thiserror", + "tokio", + "tracing", + "tracing-subscriber", +] + +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" + +[[package]] +name = "block" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d8c1fef690941d3e7788d328517591fecc684c084084702d6ff1641e993699a" + +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "core-foundation" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e195e091a93c46f7102ec7818a2aa394e1e1771c3ab4825963fa03e45afb8f" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + +[[package]] +name = "core-graphics-types" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45390e6114f68f718cc7a830514a96f903cccd70d02a8f6d9f643ac4ba45afaf" +dependencies = [ + "bitflags 1.3.2", + "core-foundation", + "libc", +] + +[[package]] +name = "foreign-types" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d737d9aa519fb7b749cbc3b962edcf310a8dd1f4b67c91c4f83975dbdd17d965" +dependencies = [ + "foreign-types-macros", + "foreign-types-shared", +] + +[[package]] +name = "foreign-types-macros" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a5c6c585bc94aaf2c7b51dd4c2ba22680844aba4c687be581871a6f518c5742" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "foreign-types-shared" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa9a19cbb55df58761df49b23516a86d432839add4af60fc256da840f66ed35b" + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[package]] +name = "malloc_buf" +version = "0.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62bb907fe88d54d8d9ce32a3cceab4218ed2f6b7d35617cafe9adf84e43919cb" +dependencies = [ + "libc", +] + +[[package]] +name = "matchers" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9" +dependencies = [ + "regex-automata", +] + +[[package]] +name = "memchr" +version = "2.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4" + +[[package]] +name = "memmap2" +version = "0.9.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1219ed1b7f229ee7104d281dd01d6802fe28bb6e95d292942c4daacdeb798c0" +dependencies = [ + "libc", +] + +[[package]] +name = "metal" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5637e166ea14be6063a3f8ba5ccb9a4159df7d8f6d61c02fc3d480b1f90dcfcb" +dependencies = [ + "bitflags 2.13.0", + "block", + "core-graphics-types", + "foreign-types", + "log", + "objc", + "paste", +] + +[[package]] +name = "nu-ansi-term" +version = "0.50.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "objc" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "915b1b472bc21c53464d6c8461c9d3af805ba1ef837e1cac254428f4a77177b1" +dependencies = [ + "malloc_buf", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "safetensors" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44560c11236a6130a46ce36c836a62936dc81ebf8c36a37947423571be0e55b6" +dependencies = [ + "serde", + "serde_json", +] + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.150" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "syn" +version = "2.0.118" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "thread_local" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f60246a4944f24f6e018aa17cdeffb7818b76356965d03b07d6a9886e8962185" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "tokio" +version = "1.52.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +dependencies = [ + "pin-project-lite", + "tokio-macros", +] + +[[package]] +name = "tokio-macros" +version = "2.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" +dependencies = [ + "matchers", + "nu-ansi-term", + "once_cell", + "regex-automata", + "sharded-slab", + "smallvec", + "thread_local", + "tracing", + "tracing-core", + "tracing-log", +] + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" diff --git a/Cargo.toml b/Cargo.toml index 60e4a12..852eb46 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,17 +2,39 @@ name = "batch_forge" version = "0.1.0" edition = "2021" -description = "High-performance inference engine for JAX/Equinox models in Rust" +rust-version = "1.75" +description = "Correctness-first Metal inference runtime for JAX/Equinox models, written in Rust" +license = "MIT OR Apache-2.0" +repository = "https://github.com/yash27-lab/batch_forge" +keywords = ["inference", "metal", "apple-silicon", "jax", "llm"] +categories = ["science", "hardware-support"] + +[lib] +name = "batch_forge" +path = "src/lib.rs" + +[[bin]] +name = "batch_forge" +path = "src/main.rs" + +[[bin]] +name = "bench" +path = "src/bin/bench.rs" [dependencies] safetensors = "0.4" memmap2 = "0.9" +serde_json = "1" bytemuck = { version = "1.14", features = ["derive"] } thiserror = "1.0" tracing = "0.1" -tracing-subscriber = "0.3" -tokio = { version = "1", features = ["full"] } -futures = "0.3" +tracing-subscriber = { version = "0.3", features = ["env-filter"] } +tokio = { version = "1", features = ["rt-multi-thread", "macros", "sync"] } [target.'cfg(target_os = "macos")'.dependencies] metal = "0.28" + +[profile.release] +opt-level = 3 +lto = "thin" +codegen-units = 1 diff --git a/LICENSE-APACHE b/LICENSE-APACHE new file mode 100644 index 0000000..4d10712 --- /dev/null +++ b/LICENSE-APACHE @@ -0,0 +1,31 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + + Full license text: http://www.apache.org/licenses/LICENSE-2.0 + +Copyright 2026 batch_forge contributors + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. diff --git a/LICENSE-MIT b/LICENSE-MIT new file mode 100644 index 0000000..ec2b541 --- /dev/null +++ b/LICENSE-MIT @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 batch_forge contributors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md index f488084..09967a6 100644 --- a/README.md +++ b/README.md @@ -1,67 +1,125 @@ # batch_forge -batch_forge is a high-performance inference engine written in Rust, designed for large-scale Transformer, Diffusion, and State-Space Models (SSMs) authored in JAX and Equinox. It provides a bare-metal, zero-Python runtime for executing complex models on edge devices and consumer hardware using Metal and Vulkan compute kernels. +A from-scratch **GPT-2 inference engine** for Apple Silicon, written in Rust with hand-written **Metal** compute kernels. -## Key Features +No PyTorch, no Python runtime, no `tokenizers` library. batch_forge loads real HuggingFace `gpt2` weights, tokenizes with its own byte-level BPE, runs the transformer on custom Metal kernels, and generates text — and **every GPU kernel is checked against a pure-Rust CPU reference**, so its output is verifiable, not asserted. -- **Asynchronous Request Management**: Built on `tokio`, featuring a non-blocking `RequestManager` for high-concurrency token generation and batching. -- **KV-Cache System**: Session-based Key-Value cache management in Metal's unified memory for efficient autoregressive generation without redundant re-computation. -- **Hybrid Metal Backend**: Combines Apple's highly optimized **Metal Performance Shaders (MPS)** for standard precision operations with custom **MSL (Metal Shading Language)** kernels for specialized tasks. -- **JAX/Equinox Integration**: Direct zero-copy loading of Equinox PyTrees via Safetensors, bypassing heavyweight XLA or TFLite runtimes. -- **Quantized Inference**: On-the-fly INT8 and INT4 dequantization kernels to minimize memory bandwidth and footprint on edge devices. +[![CI](https://github.com/yash27-lab/batch_forge/actions/workflows/ci.yml/badge.svg)](https://github.com/yash27-lab/batch_forge/actions/workflows/ci.yml) -## Architecture +``` +$ cargo run --release --bin generate -- --prompt "The meaning of life is" --greedy -The engine is built on four core pillars: -1. **Async Runner**: A `tokio`-based server architecture that manages concurrent inference requests via mpsc channels and oneshot responses. -2. **Stateful Session Storage**: Persistent KV-cache buffers pre-allocated per request ID to support large-scale language generation. -3. **Zero-Copy Loader**: Uses memory-mapped I/O (`mmap`) to load weights instantly from Safetensors files without additional memory overhead. -4. **Optimized Dispatch**: Dynamically branches between Apple MPS for standard matmuls and custom hand-written kernels for fused attention and quantized states. +The meaning of life is not the same as the meaning of death. -## Getting Started +[25 tokens in 3.06s = 8.2 tok/s on metal] +``` -### 1. Exporting Models from JAX/Equinox +``` +$ cargo run --release --bin generate -- \ + --prompt "In a shocking turn of events, scientists discovered" --temperature 0.7 -Install the required Python utilities: -```bash -pip install jax equinox safetensors numpy +In a shocking turn of events, scientists discovered that when they were forced to +eat a "clean meal" each morning, they found that they were nearly three times more +likely to lose weight. ``` -Export your Equinox model (PyTree) to the Safetensors format: -```bash -python python/export_eqx.py --out model.safetensors +That text is produced by a 124M-parameter GPT-2 running on the Metal backend. The next-token predictions are **bit-for-bit rank-identical to HuggingFace `transformers`**, verified by [`tests/gpt2_e2e.rs`](tests/gpt2_e2e.rs) and the NumPy reference in [`python/gpt2_reference.py`](python/gpt2_reference.py) — CPU and Metal logits agree to **9e-5**. + +## Why this is more than a toy + +The hard part of an inference engine isn't the architecture — it's being *correct*. batch_forge keeps a pure-Rust CPU implementation of every operator as the ground truth, and validates each Metal kernel against it on randomized inputs ([`tests/parity.rs`](tests/parity.rs)). That discipline caught a real bug during development: + +> GPT-2's GELU drove the tanh argument past ~70. Metal's fast-math `tanh` evaluates `exp(2·arg)`, which **overflows f32 to inf → NaN**, while Rust's CPU `tanh` saturates correctly. Every individual kernel passed parity at small magnitudes; only the composed forward produced `NaN`. The CPU reference + a magnitude-scaled parity test pinned it to GELU in minutes. Fix: clamp the tanh argument (it's saturated to ±1 by |arg|=15 anyway). That regression test now lives in the suite. + +That is the whole point of the design: a reviewer can trust the GPU path without owning a Mac, because the tests prove CPU≡Metal. + +## What works today + +| Component | Status | Verified by | +|-----------|--------|-------------| +| **GPT-2 (124M) text generation, CPU + Metal** | Done | `tests/gpt2_e2e.rs` (rank-identical to HF) | +| From-scratch byte-level **BPE tokenizer** | Done | `encode("Hello world") == [15496, 995]`, round-trips | +| **Metal kernels**: tiled matmul, multi-head attention, layernorm, rmsnorm, rope, gelu, int8 dequant | Done | `cargo test --test parity` (CPU↔Metal on-device) | +| Pure-Rust CPU reference for every op | Done | `cargo test --lib` | +| **Tiled matmul** (threadgroup shared memory) | Done | 1.8× over naive @ 512, 296 GF/s @ 1024 | +| Zero-copy `mmap` safetensors loader (unaligned-safe) | Done | unit + e2e | +| Sampling: greedy, temperature, top-k | Done | demo | +| Async request engine (`tokio` mpsc/oneshot) | Done | `--requests N` on the MLP path | + +## Roadmap (not yet built — stated honestly) + +- **KV cache for generation.** Today each step recomputes the full sequence (`O(n²)` over the context). The cache kernels exist (`update_kv_cache`, `kv_attention`); wiring them into the GPT-2 loop is next and is the biggest generation speedup available. +- **Resident weights.** The ergonomic op API re-uploads weights to the GPU each call; pooling/persisting them is a large, easy win. +- **FP16/BF16 compute**, **larger GPT-2 / Llama**, **INT4**, **flash-attention-style fused kernel**, **Vulkan/WebGPU**. + +This is not competing with [MLX](https://github.com/ml-explore/mlx) / [llama.cpp](https://github.com/ggerganov/llama.cpp) / [candle](https://github.com/huggingface/candle). It's a correctness-first engine that runs a real LLM end-to-end and proves it. + +## Architecture + +``` + prompt ──► BPE tokenizer (from scratch) ──► token ids + │ + gpt2 safetensors ──► mmap loader ──► Gpt2 ◄────┘ + │ forward + ┌──────────────┴───────────────┐ + ▼ ▼ + CpuBackend (ops.rs) MetalBackend (compute.metal) + the ground truth ◄─ parity ─► tiled matmul · MHA · + layernorm · gelu · … ``` -### 2. Building the Rust Engine +- `src/gpt2.rs` — model, the `LlmOps` backend trait, forward pass, sampling. +- `src/tokenizer.rs` — byte-level BPE (`bytes_to_unicode`, merges, pre-tokenizer). +- `src/ops.rs` — pure-Rust reference numerics (the spec). +- `src/metal_backend.rs` + `src/shaders/compute.metal` — the Metal kernels. +- `tests/parity.rs`, `tests/gpt2_e2e.rs` — CPU↔Metal + end-to-end verification. + +## Quickstart + +### Requirements + +- Rust 1.75 or later +- Apple Silicon and macOS for the Metal backend. The CPU backend builds on other platforms; use `--backend cpu` there. + +The default binary runs exported MLP checkpoints. The separate generate binary runs GPT-2: -Ensure you have the Rust toolchain installed. Build the project in release mode: ```bash -cargo build --release +python python/fetch_gpt2.py +cargo run --release --bin generate -- --backend cpu --prompt "Hello" --greedy +cargo run --release --bin generate -- --help ``` -### 3. Running Inference & Benchmarks +Generation writes text to stdout and diagnostics to stderr, so text can be redirected to a file. Use `--model-dir PATH` to select another asset directory. An empty prompt starts from GPT-2's end-of-text token; an explicitly requested unavailable backend returns an error. -Run the engine to start the async inference loop and execute the built-in performance benchmarks: ```bash -./target/release/batch_forge -``` +# 1. Build (Apple Silicon for Metal; the CPU path builds anywhere) +cargo build --release -## Performance Comparison +# 2. Create a deterministic demo MLP checkpoint + reference +python python/make_demo_model.py -batch_forge includes built-in benchmarking to compare custom compute kernels against native Apple Silicon hardware acceleration (MPS): -- **Custom Kernel**: Hand-written MSL shaders for specialized operations (INT8, Fused Attention). -- **Apple MPS**: Assembly-tuned matrix multiplication for standard FP32/FP16 precision. +# 3. Run the portable CPU reference +cargo run --release -- --model model.safetensors --backend cpu + +# 4. Verify the model output against the generated reference +cargo run --release -- --model model.safetensors --verify reference.safetensors + +# 5. Tests + benchmarks +cargo check # fast compile-only preflight +cargo clippy --all-targets # lint binaries and tests +cargo test --lib # portable CPU reference tests +cargo test --test parity -- --nocapture # CPU↔Metal parity (Apple Silicon) +cargo run --release --bin bench # matmul/gelu/MLP numbers on your machine +``` -## Supported Architectures +`batch_forge` flags: `--model/-m`, `--verify/-v`, `--requests/-r`, `--backend cpu|metal|both`. -- **Transformers**: Autoregressive LLMs with KV-Cache support. -- **Vision Language Models (VLMs)**: High-throughput vision feature extraction. -- **State-Space Models (SSMs)**: Selective scan operations (e.g., Mamba). -- **Diffusion**: Fast UNet/DiT inference for image generation. +## Performance & correctness -## Contributing +Both measured and reproducible — see [docs/benchmarks.md](docs/benchmarks.md) (real M2 numbers, tiled vs naive matmul) and [docs/correctness.md](docs/correctness.md) (per-op CPU↔Metal deviations + the GPT-2 end-to-end check). For common setup and platform questions, see the [troubleshooting guide](docs/troubleshooting.md); browse the [documentation index](docs/README.md) for the complete guide list. Contributors can start with [CONTRIBUTING.md](CONTRIBUTING.md). -Contributions focusing on new compute kernels, quantization techniques (FP8/NF4), or additional hardware backends (Vulkan/WebGPU) are welcome. Please ensure all new kernels include numerical parity tests against JAX references. +For hands-on use, read the [generation guide](docs/generation.md), +[tensor API guide](docs/tensors.md), and [async engine guide](docs/engine.md). ## License diff --git a/docs/README.md b/docs/README.md new file mode 100644 index 0000000..cfbd9f4 --- /dev/null +++ b/docs/README.md @@ -0,0 +1,13 @@ +# Documentation + +Start with the [test command reference](test-matrix.md) to choose the smallest validation check for your change. + +- [Benchmarks](benchmarks.md) — measured M2 performance and tiled-versus-naive matmul results. +- [Correctness](correctness.md) — CPU-to-Metal parity and end-to-end GPT-2 validation. +- [Contributing](../CONTRIBUTING.md) — local formatting, compile, lint, and test checks before proposing a change. +- [Issue templates](../.github/ISSUE_TEMPLATE/) — structured bug reports and enhancement proposals. +- [Test command reference](test-matrix.md) — formatting, compile, lint, and test commands for each task. +- [Troubleshooting](troubleshooting.md) — platform, model-asset, and test setup notes. +- [Generation guide](generation.md) — GPT-2 setup, command-line controls, output, and context behavior. +- [Tensor API guide](tensors.md) — Shapes, fallible allocation, validation, parity, and borrowed views. +- [Async engine guide](engine.md) — Request submission, queue bounds, cancellation, shutdown, and errors. diff --git a/docs/benchmarks.md b/docs/benchmarks.md new file mode 100644 index 0000000..e3294fa --- /dev/null +++ b/docs/benchmarks.md @@ -0,0 +1,103 @@ +# batch_forge Benchmarks + +All numbers here are produced by `cargo run --release --bin bench` — there are **no +hard-coded results**. Re-run it on your machine to get your own. The table below +was captured on the reference device described next. + +## Environment + +| | | +|---|---| +| **Device** | Apple M2 (8-core GPU), unified memory | +| **OS** | macOS (Metal 4) | +| **Toolchain** | Rust 1.96, `--release` (LTO thin, codegen-units = 1) | +| **Compute** | Custom MSL kernels (no MPS), FP32 | + +## Methodology + +- Each op is warmed up, then timed over many iterations; the mean is reported. +- Metal timings are **end-to-end**: buffer allocation + kernel dispatch + + `wait_until_completed` + readback. On Apple Silicon's unified memory there is + no discrete host↔device transfer, but per-call allocation overhead **is** + included — which is why the GPU loses at tiny sizes and wins as work grows. +- CPU is the same naive reference used for correctness (single-threaded, no SIMD + intrinsics, simple loop-order blocking only). + +The numbers below are a reference point, not a cross-device ranking; Apple GPU performance varies with chip generation, OS version, and thermal conditions. + +## Results (Apple M2) + +### GPT-2 (124M) text generation + +End-to-end autoregressive generation on the Metal backend, no KV cache yet +(each step recomputes the full context): + +| Backend | Throughput | +|---------|-----------:| +| Metal (M2) | ~8 tok/s | + +The dominant cost today is (a) recomputing the whole sequence each step and +(b) re-uploading weights to the GPU per call. Both are on the roadmap; a KV cache +alone removes the `O(n²)` blowup. + +### Square matmul, FP32 — naive vs tiled Metal + +The tiled kernel stages 16×16 tiles into threadgroup memory; the speedup over the +naive one-thread-per-output kernel grows with size: + +| N | CPU GF/s | naive GF/s | tiled GF/s | tiled / naive | +|------:|--------:|--------:|--------:|------:| +| 128 | 33.4 | 15.0 | 16.1 | 1.1× | +| 256 | 27.6 | 38.5 | 52.3 | 1.4× | +| 512 | 26.2 | 113.6 | 202.4 | **1.8×** | +| 1024 | 28.2 | 206.9 | 295.7 | 1.4× | + +At N=128 the GPU is *slower* than CPU — dispatch + allocation overhead dominates a +tiny problem. Tiled matmul reaches ~296 GF/s at N=1024 (still well below the M2's +FP32 peak; a `simdgroup_matrix` kernel is the next step). + +### GELU (elementwise, 2²⁰ elements) + +| CPU (ms) | Metal (ms) | Speedup | +|------:|------:|------:| +| 6.136 | 1.178 | 5.2× | + +### MLP forward (256 → 1024 → 1024 → 256) + +| Batch | CPU (ms) | Metal (ms) | Speedup | +|------:|------:|------:|------:| +| 1 | 1.122 | 1.826 | 0.6× | +| 8 | 8.966 | 2.270 | 3.9× | +| 32 | 35.985 | 5.556 | 6.5× | + +Batching amortizes per-dispatch overhead: at batch 1 the GPU trails, by batch 32 +it is 6.5× ahead. + +## Honest limitations (a.k.a. headroom) + +These results are from **deliberately simple kernels**. They are correct first; +fast second. Known gaps, in rough priority order: + +1. **Naive matmul.** One thread per output element, no threadgroup tiling, no + `simdgroup_matrix`. ~108 GFLOP/s is a small fraction of the M2 GPU's FP32 + peak (~3.6 TFLOP/s). A tiled kernel should close most of that gap. +2. **Per-call buffer allocation.** The ergonomic `Vec`-in/`Vec`-out methods + allocate buffers every call. The buffer-level API (used by the KV-cache path) + avoids this; the high-level API should pool buffers. +3. **Attention is single-thread-per-query and recomputes scores** (O(M·S·D), down + from the previous O(M·S·D²) bug, but not yet flash-attention style with + threadgroup reductions). +4. **No FP16/BF16 compute path yet** — everything runs in FP32. + +## Reproducing + +```bash +cargo run --release --bin bench +``` + +For end-to-end model latency with a real checkpoint: + +```bash +python python/make_demo_model.py +cargo run --release --bin batch_forge -- --requests 256 # prints req/s +``` diff --git a/docs/correctness.md b/docs/correctness.md new file mode 100644 index 0000000..4ce17e1 --- /dev/null +++ b/docs/correctness.md @@ -0,0 +1,122 @@ +# batch_forge Correctness + +An inference engine is only useful if you can trust its output. batch_forge backs +that trust with a concrete, runnable method rather than a promise: + +> **The pure-Rust CPU implementation in `src/ops.rs` is the ground truth. Every +> Metal kernel is tested against it on randomized inputs.** When the GPU and the +> reference agree to within tolerance, the kernel meets the test's numerical +> contract. + +This is the same discipline ggml and candle use, and it means a reviewer can +believe the Metal path without owning a Mac — the test either passes in CI-visible +output or it doesn't. + +## How verification works + +1. **CPU unit tests** (`cargo test --lib`) pin the reference ops to known values + (e.g. GELU reference points, softmax sums to 1, RoPE preserves norm, causal + masking hides future keys). +2. **CPU↔Metal parity tests** (`cargo test --test parity`) generate random inputs, + run the reference and the Metal kernel, and assert the maximum absolute + deviation is within tolerance. Run on Apple Silicon. +3. **End-to-end verification** (`--verify`) compares the full model forward pass + against a reference `output` exported alongside the weights by the Python + tooling (JAX in `export_eqx.py`, NumPy in `make_demo_model.py`). + +Tensor and Metal parity comparisons treat non-finite differences as failures. +NaN outputs cannot disappear from the maximum-difference calculation and be +reported as an exact match. + +```bash +cargo test --lib +cargo test --test parity -- --nocapture +python python/make_demo_model.py +cargo run --release --bin batch_forge -- --verify reference.safetensors +``` + +## Observed deviations (Apple M2, FP32) + +Measured by `cargo test --test parity -- --nocapture`. These are the *actual* +maximum absolute differences between the Metal kernel and the CPU reference, not +the (looser) thresholds the tests assert against. + +| Operator | Test tolerance | Observed max\|Δ\| | Notes | +|----------|---------------:|-------------------:|-------| +| MatMul (32×64×48) | 1e-3 | 1.4e-6 | FP32 accumulation order differs | +| Linear (16×64×40) | 1e-3 | 1.4e-6 | `y = x·Wᵀ + b` | +| GELU (4096) | 1e-5 | 7.5e-8 | tanh approximation | +| LayerNorm (16×64) | 1e-4 | 2.4e-7 | population variance | +| RMSNorm (16×64) | 1e-4 | 1.2e-7 | | +| RoPE (8×64) | 1e-3 | 6.0e-7 | rotate-half, `sin/cos/pow` | +| Attention, full (m=4,s=12,d=32) | 1e-3 | 8.9e-8 | | +| Attention, causal (q_offset=3) | 1e-3 | 1.2e-7 | masking + softmax | +| INT8 quant matmul (24×48×32) | 1e-3 | 7.6e-6 | vs dequantize+matmul | +| Tiled matmul (40×72×56) | 1e-3 | 1.9e-6 | shared-memory GEMM | +| Multi-head attention (s=7,h=3,d=8) | 1e-3 | 1.2e-7 | causal, GPT-2 layout | +| GELU at scale (5×3072, ±12) | 1e-2 | 4.8e-7 | regression test for the tanh-overflow fix | +| Cached attention (e2e) | 1e-3 | 6.0e-8 | `update_kv_cache` → `kv_attention` | +| KV-cache write | exact | 0 | bitwise copy check | +| **MLP forward vs reference** | 1e-3 | **1.3e-6** | full model, CPU & Metal | +| **GPT-2 logits, CPU vs Metal** | 1e-2 | **9.2e-5** | full 12-layer forward | +| **GPT-2 top-5 vs HuggingFace** | exact | match | `tests/gpt2_e2e.rs` | + +### The GELU / tanh-overflow bug + +A worked example of why the CPU reference matters. GPT-2 activations drove the +GELU tanh argument past ~70. Metal's fast-math `tanh` evaluates `exp(2·arg)`, +which overflows f32 to `inf` and yields `NaN`; Rust's CPU `tanh` saturates. Every +op passed parity at small magnitudes, so the bug only surfaced in the composed +forward as all-`NaN` logits. A magnitude-scaled GELU parity test localized it +immediately, and the fix (clamping the tanh argument, exact since tanh saturates +by |arg|=15) is now a permanent regression test (the "GELU at scale" row above). + +Deviations are dominated by FP32 summation-order differences between the +sequential CPU loop and the parallel GPU kernel — i.e. the kernels are doing the +same math, just associating the additions differently. Tolerances are set with +generous margin above the observed values. + +## Tolerance rationale + +| Precision | Typical bound | Why | +|-----------|---------------|-----| +| FP32 | 1e-3 abs (observed ~1e-6) | accumulation-order differences only | +| INT8 (W8A16) | 1e-3 abs (observed ~7e-6) | per-row scale dequant is exact up to FP32 rounding | +| FP16 / BF16 | n/a | no half-precision compute path yet (roadmap) | + +## What is *not* yet covered + +Being explicit, since the previous version of this doc claimed tests that did not +exist: + +- No FP16/BF16 numerics (no half-precision kernels yet). +- No SSM `scan` op (not implemented). +- No full transformer-LM end-to-end parity (no LM yet) — only the MLP is verified + end-to-end, plus every individual operator above. + +## Invalid-input contracts and sampling + +CPU reference operators check dimension products before allocation. Normalization +requires a positive width and a finite non-negative epsilon; rotary embeddings +require a finite positive theta. Attention validates head widths and bounds +absolute causal offsets without wrapping. + +Use Sampler::sample_token for checked single-token sampling, and Sampler::validate +to check temperature. Empty or non-finite logits are rejected. Probabilities use +f64 intermediates; top-k ties prefer smaller token IDs and the retained candidates +have a stable ID order. Seeded sampled sequences may differ from earlier versions +because their candidate ordering and arithmetic were corrected. Greedy finite +argmax retains its earliest-ID tie behavior. + +Tokenizer merge parsing retains hash-prefixed merge tokens, rejects malformed or +duplicate pairs, and checks both operands. Repeated whitespace preserves GPT-2's +lookahead boundary before the next token. The additive try_decode and +try_decode_bytes APIs report unknown token IDs; existing decode APIs retain their +permissive behavior. + +The download script accepts --out, --timeout, and --force. Each transfer uses a +temporary file beside its destination, checks any advertised Content-Length, and +atomically replaces the destination only after completion. Existing nonempty +assets are skipped by default; --force refreshes them. Network-free downloader +regressions run in CI. These tests do not download or validate pretrained GPT-2 +weights and do not replace on-device Metal parity tests. diff --git a/docs/engine.md b/docs/engine.md new file mode 100644 index 0000000..d431660 --- /dev/null +++ b/docs/engine.md @@ -0,0 +1,46 @@ +# Async inference engine guide + +The engine serves MLP requests through a bounded Tokio channel. Create a +submitter while a Tokio runtime is active, then await each request result: + +```rust +use std::sync::Arc; +use batch_forge::engine::spawn; +use batch_forge::model::{CpuBackend, Mlp, ModelError}; +use batch_forge::tensor::Tensor; + +async fn infer_once(model: Arc, input: Tensor) -> Result { + let client = spawn(Arc::new(CpuBackend), model, 8); + client.infer(1, input).await +} +``` + +The input is owned by the submitted request. The request ID is used for +tracing; it is not a deduplication key or a guarantee of result persistence. + +## Queue capacity and execution + +`queue_depth` sets the bounded channel capacity; zero is normalized to one +slot. When the queue is full, sending waits for room. This bounds queued +requests, but not the number of caller futures retaining inputs while waiting +to submit. + +The manager runs one forward pass at a time. The forward call is synchronous +inside its Tokio task, so CPU work occupies a runtime worker during inference. +Async channels do not fuse requests into a batch or make the numerical work +itself asynchronous. + +## Cancellation, shutdown, and errors + +If a response receiver is dropped while its request is queued, the manager +skips that request before computing. A forward pass that has already started +is not interrupted by this queue check. + +Dropping all submitter handles closes the queue. The manager drains available +requests and exits after receiving channel closure. `infer` returns +`ModelError::EngineClosed` when sending or receiving can no longer complete; +model-validation errors are returned through the same result channel. + +The submitter does not retry automatically. A caller should distinguish an +invalid input from a closed engine before deciding whether another attempt is +appropriate. diff --git a/docs/generation.md b/docs/generation.md new file mode 100644 index 0000000..fd53d6e --- /dev/null +++ b/docs/generation.md @@ -0,0 +1,130 @@ +# GPT-2 generation guide + +The `generate` binary runs GPT-2 small (124M). The default `batch_forge` binary +runs exported MLP checkpoints. Invoke the generation binary explicitly: + +```bash +python3 python/fetch_gpt2.py +cargo run --release --bin generate -- --backend cpu --prompt "Hello" --greedy +``` + +Run commands from the repository root. Downloads are setup work; generation +runs in Rust and does not require a Python runtime. + +## Choose an asset directory + +Keep `model.safetensors`, `vocab.json`, and `merges.txt` together. The download +script's `--out` directory and the generator's `--model-dir` must refer to the +same location. Quote paths that contain spaces: + +```bash +python3 python/fetch_gpt2.py --out "/path/to/gpt2 assets" +cargo run --release --bin generate -- --model-dir "/path/to/gpt2 assets" --backend cpu +``` + +Existing nonempty assets are skipped by the download script. Its `--force` +option replaces them after a completed transfer; `--timeout SECONDS` sets the +socket timeout. Failed transfers leave an existing final asset intact. + +## Inspect options without downloading a model + +```bash +cargo run --bin generate -- --help +cargo run --bin generate -- --version +``` + +Help and version return before loading model assets. Unknown arguments, +missing values, invalid integers, and invalid temperatures return argument-error +exit status 2. Asset and execution errors return failure. The default MLP +binary has its own `--help`; its flags differ from the generator's flags. + +## Select the compute backend + +`--backend cpu` selects the portable Rust reference. On macOS, the generator +defaults to `metal`; elsewhere it defaults to `cpu`. An explicit Metal request +fails on a non-macOS platform or when Metal initialization fails. It does not +silently select a different backend. + +```bash +cargo run --release --bin generate -- --backend cpu --prompt "Once upon a time" +cargo run --release --bin generate -- --backend metal --prompt "Once upon a time" +``` + +See [platform troubleshooting](troubleshooting.md) for Metal requirements. + +## Redirect generated text + +Generation writes text to stdout and setup, backend, and timing information +to stderr. Redirect them independently: + +```bash +cargo run --release --bin generate -- --backend cpu --prompt "Hello" > answer.txt 2> run.log +``` + +The output stream ends with a newline. Characters that span multiple tokenizer +tokens are emitted after their UTF-8 bytes are complete. Closing an output pipe +early stops the token loop and treats a broken pipe as normal termination; +other output errors report failure. + +## Empty prompts and end-of-text + +An empty CLI prompt is seeded internally with GPT-2's end-of-text token. This +provides the first input token required by the model without printing that seed. +The loop stops when it samples end-of-text and does not print the marker. + +```bash +cargo run --release --bin generate -- --backend cpu --prompt "" --max-new 20 +``` + +`--max-new` must be a positive integer. Timing statistics count sampled tokens, +including a sampled stop token, so the token count is not a character count. + +## Temperature and greedy decoding + +`--temperature 0` selects greedy argmax. A positive temperature enables +sampling; the value must be finite and non-negative. `--greedy` forces zero +temperature regardless of where the temperature flag appears: + +```bash +cargo run --release --bin generate -- --backend cpu --greedy --temperature 0.8 --prompt "Hello" +``` + +Greedy ties retain the earliest token ID. NaN, infinity, and negative CLI +temperatures are rejected before asset loading. + +## Top-k candidate filtering + +`--top-k K` restricts sampling to the K largest logits. Zero disables this +filter, and a value at least as large as the vocabulary retains all candidates. +With positive temperature, `--top-k 1` leaves only the highest candidate. + +Equal-logit cutoff ties prefer smaller token IDs, including signed-zero ties. +Retained candidates use a stable token-ID order when drawing from their +probabilities. A greedy run uses argmax directly and ignores top-k filtering. + +## Reproduce a sampling run + +Choose a seed explicitly when comparing runs: + +```bash +cargo run --release --bin generate -- --backend cpu --seed 42 --temperature 0.8 --top-k 40 --prompt "Hello" +``` + +The generation loop retains its RNG across tokens. Record the seed, backend, +model and tokenizer assets, prompt, sampling flags, and code revision. The same +seed alone does not guarantee identical output across different assets, +backends, or revisions: numerical logits and sampling implementation changes +can affect the result. + +## Context limits and loading errors + +The default GPT-2 small configuration has a 1024-token context. Generation +forwards the most recent 1024 tokens when a sequence grows beyond that limit, +and uses positions within that forwarded slice. Each step recomputes its +sequence; the loop does not use a KV cache. + +Loading rejects checkpoint shapes that differ from the configured model, +structurally invalid tensors, and mismatched tokenizer vocabulary size. Keep +weights and tokenizer files from the same model. For a report, include the +failing command, error text, backend/platform, and asset source; see +[troubleshooting](troubleshooting.md) and [correctness](correctness.md). diff --git a/docs/tensors.md b/docs/tensors.md new file mode 100644 index 0000000..6cfe74a --- /dev/null +++ b/docs/tensors.md @@ -0,0 +1,87 @@ +# Tensor API guide + +`Tensor` is an owned row-major F32 buffer. Its shape records dimensions; its +data vector records elements. `Tensor::new` checks that these agree: + +```rust +use batch_forge::tensor::{Tensor, TensorError}; + +fn matrix() -> Result { + Tensor::new(vec![1.0, 2.0, 3.0, 4.0], vec![2, 2]) +} +``` + +`dims2()` checks rank and returns `(rows, cols)`. A 2×2 tensor contains four +F32 elements; byte-backed views count bytes separately. + +## Scalars and empty dimensions + +An empty shape `vec![]` describes a scalar containing one element: + +```rust +let scalar = Tensor::new(vec![3.0], vec![])?; +``` + +A shape containing a zero dimension describes zero elements. For example, +`vec![2, 0, 3]` requires an empty data vector. A zero dimension makes the count +zero regardless of its position in the shape. These are shape semantics; +model layers and reduction operators can impose additional positive-width +requirements. + +## Fallible zero allocation + +Use `try_zeros` when dimensions come from a caller or checkpoint: + +```rust +fn empty_batch() -> Result { + Tensor::try_zeros(vec![0, 768]) +} +``` + +An overflowing element count returns `TensorError::BufferOverflow`. A failed +buffer reservation returns `TensorError::AllocationFailed`. The convenience +`zeros` constructor panics on those failures. The fallible API handles errors +reported by buffer reservation; it is not a process-wide memory limit. + +## Validate after changing public fields + +`shape` and `data` are public. If a caller edits either after construction, +call `tensor.validate()` before using that structure: + +```rust +let mut tensor = Tensor::new(vec![1.0, 2.0], vec![2])?; +tensor.shape = vec![1, 2]; +tensor.validate()?; +``` + +Validation checks structural consistency and shape overflow. It does not +require every F32 value to be finite. MLP inference revalidates its public +layer tensors and the supplied input before forwarding. + +## Interpret parity differences + +`max_abs_diff` returns the largest absolute difference for valid, equal-shaped +buffers. It returns infinity for shape mismatch, invalid buffers, or a +non-finite element difference; even two equally truncated tensors fail. + +```rust +let difference = actual.max_abs_diff(&reference); +let passes = difference.is_finite() && difference <= 1e-3; +``` + +Choose a tolerance appropriate to the operation and test. A matching numeric +buffer is evidence about that comparison, not proof that a whole model or +asset set has been verified. See [correctness](correctness.md) for the test +boundaries used by this repository. + +## Borrowed views and portable F32 materialization + +`TensorView::new` checks shape, dtype, and byte-buffer size while borrowing the +bytes. `to_tensor_f32()` copies F32 values by reading little-endian bytes and +works with an unaligned source offset. It rejects a non-F32 dtype and +revalidates public fields before conversion. + +`as_slice::()` attempts a byte reinterpretation with alignment and length +checks. It does not verify that `T` matches the semantic dtype or convert byte +order. Callers must check those conditions themselves; use F32 materialization +when a portable owned tensor is needed. diff --git a/docs/test-matrix.md b/docs/test-matrix.md new file mode 100644 index 0000000..53dc7a2 --- /dev/null +++ b/docs/test-matrix.md @@ -0,0 +1,25 @@ +# Test command reference + +Use the smallest relevant check while iterating: + +| Goal | Command | Platform | +| --- | --- | --- | +| Check formatting without changing files | `cargo fmt --check` | Any supported Rust platform with `rustfmt` | +| Check for compile errors quickly | `cargo check` | Any supported Rust platform | +| Run lints across binaries and tests | `cargo clippy --all-targets` | Any supported Rust platform with `clippy` | +| Validate CPU reference operations | `cargo test --lib --bins` | Any supported Rust platform | +| Test asset downloads without network access | `python3 -m unittest discover -s python -p 'test_*.py'` | Python 3; standard library only | +| Check Metal-to-CPU kernel parity | `cargo test --test parity -- --nocapture` | Apple Silicon macOS | +| Check end-to-end GPT-2 predictions | `cargo test --test gpt2_e2e -- --nocapture` | Requires downloaded GPT-2 assets | + +Run these commands from the repository root. For a lightweight CPU-only preflight, run `cargo fmt --check`, `cargo check`, and `cargo test --lib --bins`; add `cargo clippy --all-targets` when the Clippy component is installed. + +Downloader regressions use mocked HTTP responses to check completed, failed, empty, and truncated transfers, timeout validation, and forced refresh. They do not download pretrained model files or require the Python ML dependencies. + +Before the end-to-end check, download the model weights and tokenizer: + +```bash +python python/fetch_gpt2.py +``` + +For platform and asset setup notes, see the [troubleshooting guide](troubleshooting.md). diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md new file mode 100644 index 0000000..140ea92 --- /dev/null +++ b/docs/troubleshooting.md @@ -0,0 +1,78 @@ +# Troubleshooting + +## Metal backend is unavailable + +The Metal backend requires macOS on Apple Silicon. On other platforms, run the portable CPU reference with a demo checkpoint: + +```bash +python python/make_demo_model.py +cargo run --release -- --backend cpu +``` + +## GPT-2 reference assets are missing + +Model weights and tokenizer files are intentionally not committed. Download them before running the GPT-2 end-to-end reference check: + +```bash +python python/fetch_gpt2.py +``` + +The files are placed under `models/gpt2/` and require roughly 550 MB of disk space. + +## Parity tests need a Mac with Metal + +`cargo test --test parity -- --nocapture` validates Metal kernels against the pure-Rust CPU reference, so it must run on an Apple Silicon Mac. The regular library test suite remains useful for CPU-only environments: + +```bash +cargo test --lib +``` + +## `cargo fmt` is unavailable + +Install Rust's formatting component for the active toolchain, then rerun the formatting check: + +```bash +rustup component add rustfmt +cargo fmt --check +``` + +## `cargo clippy` is unavailable + +Install the Clippy component for the active toolchain, then rerun the lint check: + +```bash +rustup component add clippy +cargo clippy --all-targets +``` + +## Generation arguments and checkpoints + +Run cargo run --bin generate -- --help for generation options. Unknown flags, +missing values, invalid integers, and non-finite or negative temperatures return +an error before loading assets. A zero top-k disables filtering; --greedy takes +precedence over temperature regardless of option order. + +A checkpoint must match the default GPT-2 small configuration. Incorrect shapes, +inconsistent tokenizer IDs, and mismatched vocabulary sizes are reported at load +time. Keep model.safetensors, vocab.json, and merges.txt from the same model +under the directory selected by --model-dir. + +Generation streams complete UTF-8 characters even when a character spans tokens. +Diagnostics go to stderr. A consumer that closes a pipe early stops generation +without a panic. Invalid bytes are replaced, and an unfinished final character +is replaced only at the end of the stream. + +## MLP verification exit status + +The default MLP binary returns failure for missing checkpoints, invalid input +buffers, failed inference requests, and failed numerical checks. A reference +file selected with --verify must contain both input and output; a +synthetic fallback is not used to claim verification. + +## Detailed usage references + +See the [generation guide](generation.md) for backend, sampling, output, and +context-window behavior. The [tensor guide](tensors.md) explains shape and +buffer errors; the [engine guide](engine.md) explains queue bounds and shutdown +behavior. Use the [test command reference](test-matrix.md) to choose a local +check for a reproduction. diff --git a/python/export_eqx.py b/python/export_eqx.py index 0529253..0eca7f7 100644 --- a/python/export_eqx.py +++ b/python/export_eqx.py @@ -1,50 +1,66 @@ -import jax -import jax.numpy as jnp -import equinox as eqx -from safetensors.numpy import save_file +"""Export an Equinox MLP to safetensors with named weights and a reference I/O pair. + +The Rust engine loads weights by name (`layers.{i}.weight` / `layers.{i}.bias`), +so this exporter walks the model's `Linear` layers and names them explicitly +instead of emitting opaque `leaf_{i}` keys. It also writes a `reference.safetensors` +containing a sample `input` and the model's `output`, which `batch_forge --verify` +uses to confirm the Rust forward pass matches JAX numerically. +""" + import argparse + +import equinox as eqx +import jax import numpy as np +from safetensors.numpy import save_file + class SimpleMLP(eqx.Module): layers: list - def __init__(self, key): - keys = jax.random.split(key, 3) + def __init__(self, key, width=256, hidden=1024): + k1, k2, k3 = jax.random.split(key, 3) self.layers = [ - eqx.nn.Linear(256, 1024, key=keys[0]), - eqx.nn.Linear(1024, 1024, key=keys[1]), - eqx.nn.Linear(1024, 256, key=keys[2]) + eqx.nn.Linear(width, hidden, key=k1), + eqx.nn.Linear(hidden, hidden, key=k2), + eqx.nn.Linear(hidden, width, key=k3), ] def __call__(self, x): for layer in self.layers[:-1]: - x = jax.nn.gelu(layer(x)) + # approximate=True (tanh GELU) matches batch_forge::ops::gelu. + x = jax.nn.gelu(layer(x), approximate=True) return self.layers[-1](x) -def export_model(model: eqx.Module, output_path: str): - print(f"Exporting Equinox model to {output_path}...") - - # Flatten the PyTree into leaves - leaves, treedef = jax.tree_util.tree_flatten(model) - - # Filter out non-array leaves (e.g., callables, strings) if necessary - tensor_dict = {} - for i, leaf in enumerate(leaves): - if hasattr(leaf, 'shape') and hasattr(leaf, 'dtype'): - # Convert JAX array to NumPy array for safetensors - name = f"leaf_{i}" - tensor_dict[name] = np.array(leaf) - print(f"Exported {name}: shape {leaf.shape}, dtype {leaf.dtype}") - - save_file(tensor_dict, output_path) - print("Export complete.") + +def export_weights(model: SimpleMLP, path: str): + tensors = {} + for i, layer in enumerate(model.layers): + # eqx.nn.Linear stores weight as [out, in] and bias as [out]. + tensors[f"layers.{i}.weight"] = np.asarray(layer.weight, dtype=np.float32) + tensors[f"layers.{i}.bias"] = np.asarray(layer.bias, dtype=np.float32) + print(f" layers.{i}: weight {tensors[f'layers.{i}.weight'].shape}") + save_file(tensors, path) + print(f"Wrote weights -> {path}") + + +def export_reference(model: SimpleMLP, path: str, seed: int): + width = model.layers[0].weight.shape[1] + rng = np.random.default_rng(seed) + x = rng.standard_normal(width).astype(np.float32) + y = np.asarray(jax.vmap(model)(x[None, :]), dtype=np.float32) + save_file({"input": x[None, :], "output": y}, path) + print(f"Wrote reference (input {x[None, :].shape} -> output {y.shape}) -> {path}") + if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Export an Equinox model to Safetensors.") - parser.add_argument("--out", type=str, default="../model.safetensors", help="Output path") + parser = argparse.ArgumentParser(description="Export an Equinox MLP to safetensors.") + parser.add_argument("--out", default="model.safetensors", help="weights output path") + parser.add_argument("--ref", default="reference.safetensors", help="reference I/O output path") + parser.add_argument("--seed", type=int, default=42) args = parser.parse_args() - key = jax.random.PRNGKey(42) - model = SimpleMLP(key) - - export_model(model, args.out) + model = SimpleMLP(jax.random.PRNGKey(args.seed)) + export_weights(model, args.out) + export_reference(model, args.ref, args.seed) + print("Done. Verify with: cargo run --release -- --verify", args.ref) diff --git a/python/fetch_gpt2.py b/python/fetch_gpt2.py new file mode 100644 index 0000000..116987d --- /dev/null +++ b/python/fetch_gpt2.py @@ -0,0 +1,69 @@ +"""Download GPT-2 weights and tokenizer with atomic, bounded transfers.""" +import argparse +import math +import os +import shutil +import sys +import tempfile +import urllib.request + +BASE = "https://huggingface.co/openai-community/gpt2/resolve/main" +FILES = ["model.safetensors", "vocab.json", "merges.txt", "config.json"] +OUT = os.path.join("models", "gpt2") + + +def positive_timeout(value): + try: + seconds = float(value) + except ValueError as error: + raise argparse.ArgumentTypeError("timeout must be a number") from error + if not math.isfinite(seconds) or seconds <= 0: + raise argparse.ArgumentTypeError("timeout must be finite and positive") + return seconds + + +def download_file(url, destination, timeout=60.0): + """Install a completed transfer atomically, preserving any previous asset.""" + directory = os.path.dirname(os.path.abspath(destination)) + descriptor, temporary = tempfile.mkstemp(prefix=".gpt2-", suffix=".part", dir=directory) + os.close(descriptor) + try: + with urllib.request.urlopen(url, timeout=timeout) as response, open(temporary, "wb") as output: + shutil.copyfileobj(response, output) + expected = response.headers.get("Content-Length") + size = os.path.getsize(temporary) + if size == 0: + raise OSError("downloaded file is empty") + if expected is not None and size != int(expected): + raise OSError(f"incomplete transfer: expected {expected} bytes, received {size}") + os.replace(temporary, destination) + finally: + if os.path.exists(temporary): + os.unlink(temporary) + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out", default=OUT, help="asset directory (default: models/gpt2)") + parser.add_argument("--timeout", type=positive_timeout, default=60.0, help="socket timeout in seconds") + parser.add_argument("--force", action="store_true", help="replace existing assets") + args = parser.parse_args(argv) + try: + os.makedirs(args.out, exist_ok=True) + for name in FILES: + destination = os.path.join(args.out, name) + if not args.force and os.path.isfile(destination) and os.path.getsize(destination) > 0: + print(f" have {name}") + continue + print(f" downloading {name} …") + download_file(f"{BASE}/{name}", destination, args.timeout) + except (OSError, ValueError) as error: + print(f"Download failed: {error}", file=sys.stderr) + return 1 + print(f"Done. Weights in {args.out}/") + print(f'Try: cargo run --release --bin generate -- --model-dir "{args.out}" --prompt "Once upon a time"') + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/python/gpt2_reference.py b/python/gpt2_reference.py new file mode 100644 index 0000000..e922b7f --- /dev/null +++ b/python/gpt2_reference.py @@ -0,0 +1,66 @@ +"""NumPy reference forward for GPT-2 — ground truth for verifying the Rust engine. + +Loads the same safetensors weights and computes next-token logits for a fixed +token sequence. Prints the top predictions and dumps the full last-position +logits to reference_logits.npy so the Rust side can assert numerical parity. +""" + +import sys + +import numpy as np +from safetensors.numpy import load_file + +W = load_file("models/gpt2/model.safetensors") +NH, HD = 12, 64 + + +def ln(x, g, b, eps=1e-5): + mu = x.mean(-1, keepdims=True) + var = x.var(-1, keepdims=True) + return (x - mu) / np.sqrt(var + eps) * g + b + + +def gelu(x): + return 0.5 * x * (1 + np.tanh(0.7978845608 * (x + 0.044715 * x**3))) + + +def softmax(x): + x = x - x.max(-1, keepdims=True) + e = np.exp(x) + return e / e.sum(-1, keepdims=True) + + +def forward(tokens): + T = len(tokens) + x = W["wte.weight"][tokens] + W["wpe.weight"][:T] + for i in range(12): + p = f"h.{i}." + a = ln(x, W[p + "ln_1.weight"], W[p + "ln_1.bias"]) + qkv = a @ W[p + "attn.c_attn.weight"] + W[p + "attn.c_attn.bias"] + q, k, v = np.split(qkv, 3, axis=-1) + heads = lambda m: m.reshape(T, NH, HD).transpose(1, 0, 2) + qh, kh, vh = heads(q), heads(k), heads(v) + att = qh @ kh.transpose(0, 2, 1) / np.sqrt(HD) + att = att + np.triu(np.ones((T, T)), 1) * -1e10 + o = softmax(att) @ vh + o = o.transpose(1, 0, 2).reshape(T, NH * HD) + x = x + o @ W[p + "attn.c_proj.weight"] + W[p + "attn.c_proj.bias"] + m = ln(x, W[p + "ln_2.weight"], W[p + "ln_2.bias"]) + h = gelu(m @ W[p + "mlp.c_fc.weight"] + W[p + "mlp.c_fc.bias"]) + x = x + h @ W[p + "mlp.c_proj.weight"] + W[p + "mlp.c_proj.bias"] + x = ln(x, W["ln_f.weight"], W["ln_f.bias"]) + return x @ W["wte.weight"].T # [T, vocab] + + +if __name__ == "__main__": + # "The meaning of life is" + tokens = [int(t) for t in (sys.argv[1:] or ["464", "3616", "286", "1204", "318"])] + logits = forward(tokens) + last = logits[-1] + top = np.argsort(last)[-5:][::-1] + print("tokens:", tokens) + print("top-5 next ids:", top.tolist()) + print("top-5 logits:", [round(float(last[i]), 3) for i in top]) + print("logit[0] ('!'):", round(float(last[0]), 3)) + np.save("models/gpt2/reference_logits.npy", last.astype(np.float32)) + print("saved last-position logits -> models/gpt2/reference_logits.npy") diff --git a/python/make_demo_model.py b/python/make_demo_model.py new file mode 100644 index 0000000..21511d9 --- /dev/null +++ b/python/make_demo_model.py @@ -0,0 +1,72 @@ +"""Generate a demo model + reference WITHOUT JAX (NumPy only). + +This mirrors `export_eqx.py`'s SimpleMLP exactly — same layer shapes, same +`y = x @ Wᵀ + b` convention, same tanh-GELU — so the Rust engine's output matches +the `output` tensor written here. It exists so the end-to-end demo and CI can run +on machines without the (heavy) JAX/Equinox stack installed. + + python python/make_demo_model.py + cargo run --release -- --verify reference.safetensors +""" + +import argparse + +import numpy as np +from safetensors.numpy import save_file + +SQRT_2_OVER_PI = 0.7978845608 + + +def gelu(x): + # Tanh approximation; matches batch_forge::ops::gelu and jax.nn.gelu(approximate=True). + return 0.5 * x * (1.0 + np.tanh(SQRT_2_OVER_PI * (x + 0.044715 * x**3))) + + +def linear(x, w, b): + # eqx.nn.Linear convention: w is [out, in], so y = x @ wᵀ + b. + return x @ w.T + b + + +def main(): + parser = argparse.ArgumentParser(description="NumPy-only demo model generator.") + parser.add_argument("--out", default="model.safetensors") + parser.add_argument("--ref", default="reference.safetensors") + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--width", type=int, default=256) + parser.add_argument("--hidden", type=int, default=1024) + args = parser.parse_args() + + rng = np.random.default_rng(args.seed) + dims = [ + (args.hidden, args.width), + (args.hidden, args.hidden), + (args.width, args.hidden), + ] + + tensors = {} + weights = [] + for i, (out_f, in_f) in enumerate(dims): + # Kaiming-ish init so activations stay in a sane range. + w = (rng.standard_normal((out_f, in_f)) / np.sqrt(in_f)).astype(np.float32) + b = np.zeros(out_f, dtype=np.float32) + tensors[f"layers.{i}.weight"] = w + tensors[f"layers.{i}.bias"] = b + weights.append((w, b)) + print(f" layers.{i}: weight {w.shape}, bias {b.shape}") + save_file(tensors, args.out) + print(f"Wrote weights -> {args.out}") + + # Reference forward (matches the Rust/JAX forward exactly). + x = rng.standard_normal((1, args.width)).astype(np.float32) + h = x + for i, (w, b) in enumerate(weights): + h = linear(h, w, b) + if i < len(weights) - 1: + h = gelu(h) + save_file({"input": x, "output": h.astype(np.float32)}, args.ref) + print(f"Wrote reference (input {x.shape} -> output {h.shape}) -> {args.ref}") + print("Verify with: cargo run --release -- --verify", args.ref) + + +if __name__ == "__main__": + main() diff --git a/python/test_fetch_gpt2.py b/python/test_fetch_gpt2.py new file mode 100644 index 0000000..2185a7a --- /dev/null +++ b/python/test_fetch_gpt2.py @@ -0,0 +1,74 @@ +"""Network-free regression tests for GPT-2 asset installation.""" +import argparse +import io +import os +import tempfile +import unittest +from unittest.mock import patch + +import fetch_gpt2 + + +class Response(io.BytesIO): + def __init__(self, payload=b"complete", length=None): + super().__init__(payload) + self.headers = {} if length is None else {"Content-Length": str(length)} + + +class DownloadTests(unittest.TestCase): + def test_complete_download_replaces_asset(self): + with tempfile.TemporaryDirectory() as directory: + destination = os.path.join(directory, "model") + with patch("fetch_gpt2.urllib.request.urlopen", return_value=Response(length=8)) as opener: + fetch_gpt2.download_file("https://example.test/model", destination, 12.0) + opener.assert_called_once_with("https://example.test/model", timeout=12.0) + with open(destination, "rb") as stream: + self.assertEqual(stream.read(), b"complete") + self.assertEqual(os.listdir(directory), ["model"]) + + def test_failure_preserves_previous_asset(self): + with tempfile.TemporaryDirectory() as directory: + destination = os.path.join(directory, "model") + with open(destination, "wb") as stream: + stream.write(b"previous") + with patch("fetch_gpt2.urllib.request.urlopen", side_effect=OSError("interrupted")): + with self.assertRaises(OSError): + fetch_gpt2.download_file("https://example.test/model", destination) + with open(destination, "rb") as stream: + self.assertEqual(stream.read(), b"previous") + self.assertEqual(os.listdir(directory), ["model"]) + + def test_empty_and_truncated_downloads_never_become_final_assets(self): + for payload, length in [(b"", None), (b"short", 100)]: + with tempfile.TemporaryDirectory() as directory: + destination = os.path.join(directory, "model") + with patch("fetch_gpt2.urllib.request.urlopen", return_value=Response(payload, length)): + with self.assertRaises(OSError): + fetch_gpt2.download_file("https://example.test/model", destination) + self.assertEqual(os.listdir(directory), []) + + def test_timeout_validation(self): + for value in ["0", "-1", "NaN", "inf", "invalid"]: + with self.assertRaises(argparse.ArgumentTypeError): + fetch_gpt2.positive_timeout(value) + self.assertEqual(fetch_gpt2.positive_timeout("5"), 5.0) + + def test_existing_assets_are_skipped_unless_forced(self): + with tempfile.TemporaryDirectory() as directory: + for name in fetch_gpt2.FILES: + with open(os.path.join(directory, name), "wb") as stream: + stream.write(b"previous") + with patch("builtins.print"), patch("fetch_gpt2.download_file") as transfer: + self.assertEqual(fetch_gpt2.main(["--out", directory]), 0) + transfer.assert_not_called() + self.assertEqual(fetch_gpt2.main(["--out", directory, "--force", "--timeout", "3"]), 0) + self.assertEqual(transfer.call_count, len(fetch_gpt2.FILES)) + + def test_main_reports_network_failure_as_unsuccessful(self): + with tempfile.TemporaryDirectory() as directory: + with patch("builtins.print"), patch("fetch_gpt2.download_file", side_effect=OSError("network")): + self.assertEqual(fetch_gpt2.main(["--out", directory]), 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/bin/bench.rs b/src/bin/bench.rs new file mode 100644 index 0000000..3c9837c --- /dev/null +++ b/src/bin/bench.rs @@ -0,0 +1,173 @@ +//! Microbenchmarks: CPU reference vs Metal for matmul, GELU, and MLP forward. +//! +//! Numbers are produced on *your* machine — there are no hard-coded results. +//! Metal timings are end-to-end (buffer allocation + dispatch + readback); +//! on Apple Silicon's unified memory there is no discrete host↔device copy, but +//! per-call allocation overhead is included and dominates at small sizes. +//! +//! Run with: `cargo run --release --bin bench` + +use std::time::{Duration, Instant}; + +use batch_forge::model::{CpuBackend, Mlp}; +use batch_forge::ops; +use batch_forge::tensor::Tensor; + +struct Rng(u64); +impl Rng { + fn new(seed: u64) -> Self { + Rng(seed | 1) + } + fn f32(&mut self) -> f32 { + let mut x = self.0; + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + self.0 = x; + ((x >> 40) as f32 / (1u64 << 24) as f32) * 2.0 - 1.0 + } + fn vec(&mut self, n: usize) -> Vec { + (0..n).map(|_| self.f32()).collect() + } +} + +/// Times `f` over `iters` runs after a warmup, returning mean duration. +fn bench(warmup: usize, iters: usize, mut f: impl FnMut()) -> Duration { + for _ in 0..warmup { + f(); + } + let start = Instant::now(); + for _ in 0..iters { + f(); + } + start.elapsed() / iters as u32 +} + +fn gflops(m: usize, k: usize, n: usize, d: Duration) -> f64 { + (2.0 * m as f64 * k as f64 * n as f64) / d.as_secs_f64() / 1e9 +} + +fn main() { + println!("batch_forge bench — {}", std::env::consts::ARCH); + + #[cfg(target_os = "macos")] + let metal = batch_forge::metal_backend::MetalBackend::new(batch_forge::SHADER_SOURCE).ok(); + #[cfg(target_os = "macos")] + if let Some(m) = &metal { + println!("Metal device: {}\n", m.device.name()); + } + + let mut rng = Rng::new(0xBEEF); + + // ---- matmul: CPU vs naive Metal vs tiled Metal ---- + println!("== matmul (square, f32): CPU vs naive Metal vs tiled Metal =="); + println!( + "{:>6} | {:>9} | {:>9} {:>8} | {:>9} {:>8} | {:>9}", + "N", "cpu GF/s", "naive ms", "GF/s", "tiled ms", "GF/s", "tiled/naive" + ); + for n in [128usize, 256, 512, 1024] { + let a = rng.vec(n * n); + let b = rng.vec(n * n); + let cpu_iters = if n >= 512 { 2 } else { 5 }; + let cpu = bench(1, cpu_iters, || { + std::hint::black_box(ops::matmul(&a, &b, n, n, n)); + }); + #[cfg(target_os = "macos")] + { + if let Some(m) = metal.as_ref() { + let gpu_iters = if n >= 512 { 15 } else { 50 }; + let naive = bench(3, gpu_iters, || { + std::hint::black_box(m.matmul(&a, &b, n, n, n)); + }); + let tiled = bench(3, gpu_iters, || { + std::hint::black_box(m.matmul_tiled(&a, &b, n, n, n)); + }); + println!( + "{n:>6} | {:>9.1} | {:>9.3} {:>8.1} | {:>9.3} {:>8.1} | {:>8.1}x", + gflops(n, n, n, cpu), + naive.as_secs_f64() * 1e3, + gflops(n, n, n, naive), + tiled.as_secs_f64() * 1e3, + gflops(n, n, n, tiled), + naive.as_secs_f64() / tiled.as_secs_f64(), + ); + continue; + } + } + println!( + "{n:>6} | {:>9.1} | {:>9} {:>8} | {:>9} {:>8} | {:>9}", + gflops(n, n, n, cpu), + "-", + "-", + "-", + "-", + "-" + ); + } + + // ---- GELU ---- + println!("\n== gelu (elementwise) =="); + let n = 1 << 20; + let x = rng.vec(n); + let cpu = bench(1, 20, || { + let mut y = x.clone(); + ops::gelu_inplace(&mut y); + std::hint::black_box(y); + }); + println!("{:>10} elems | cpu {:>8.3} ms", n, cpu.as_secs_f64() * 1e3); + #[cfg(target_os = "macos")] + if let Some(m) = &metal { + let gpu = bench(3, 50, || { + std::hint::black_box(m.gelu(&x)); + }); + println!( + "{:>10} elems | metal {:>6.3} ms ({:.1}x)", + n, + gpu.as_secs_f64() * 1e3, + cpu.as_secs_f64() / gpu.as_secs_f64() + ); + } + + // ---- MLP forward ---- + println!("\n== MLP forward (256→1024→1024→256) =="); + let model = random_mlp(&mut rng, 256, 1024); + println!("{:>6} | {:>12} | {:>12}", "batch", "cpu (ms)", "metal (ms)"); + for batch in [1usize, 8, 32] { + let input = Tensor::new(rng.vec(batch * 256), vec![batch, 256]).unwrap(); + let cpu = bench(1, 10, || { + std::hint::black_box(model.forward(&CpuBackend, &input).unwrap()); + }); + let cpu_ms = cpu.as_secs_f64() * 1e3; + #[cfg(target_os = "macos")] + { + if let Some(m) = &metal { + let gpu = bench(3, 50, || { + std::hint::black_box(model.forward(m, &input).unwrap()); + }); + println!( + "{batch:>6} | {cpu_ms:>12.3} | {:>12.3}", + gpu.as_secs_f64() * 1e3 + ); + continue; + } + } + println!("{batch:>6} | {cpu_ms:>12.3} | {:>12}", "-"); + } +} + +fn random_mlp(rng: &mut Rng, width: usize, hidden: usize) -> Mlp { + let dims = [(hidden, width), (hidden, hidden), (width, hidden)]; + let layers = dims + .iter() + .map(|&(out_f, in_f)| { + let scale = 1.0 / (in_f as f32).sqrt(); + let w: Vec = rng.vec(out_f * in_f).iter().map(|v| v * scale).collect(); + let b = vec![0.0f32; out_f]; + ( + Tensor::new(w, vec![out_f, in_f]).unwrap(), + Tensor::new(b, vec![out_f]).unwrap(), + ) + }) + .collect(); + Mlp { layers } +} diff --git a/src/bin/generate.rs b/src/bin/generate.rs new file mode 100644 index 0000000..6ccdfe3 --- /dev/null +++ b/src/bin/generate.rs @@ -0,0 +1,357 @@ +//! GPT-2 text generation on batch_forge. +//! +//! cargo run --release --bin generate -- --prompt "The meaning of life is" +//! +//! Loads HuggingFace `gpt2` weights, tokenizes with the from-scratch BPE +//! tokenizer, runs the transformer on the Metal backend (or CPU), and streams +//! decoded text. + +use std::io::{self, Write}; +use std::path::PathBuf; +use std::process::ExitCode; +use std::time::Instant; + +use batch_forge::gpt2::{Config, Gpt2, LlmOps, Sampler}; +use batch_forge::loader; +use batch_forge::tokenizer::Tokenizer; + +const EOT: usize = 50256; // <|endoftext|> +const MODEL_DIR: &str = "models/gpt2"; + +#[derive(Debug)] +struct Args { + model_dir: PathBuf, + prompt: String, + max_new: usize, + backend: String, + temperature: f32, + top_k: usize, + seed: u64, +} + +enum Command { + Run(Args), + Help, + Version, +} + +fn next_value(it: &mut impl Iterator, flag: &str) -> Result { + it.next().ok_or_else(|| format!("{flag} needs a value")) +} + +fn parse_args(arguments: impl IntoIterator) -> Result { + let mut a = Args { + model_dir: MODEL_DIR.into(), + prompt: "The meaning of life is".to_string(), + max_new: 40, + backend: if cfg!(target_os = "macos") { + "metal" + } else { + "cpu" + } + .to_string(), + temperature: 0.8, + top_k: 40, + seed: 42, + }; + let mut greedy = false; + let mut it = arguments.into_iter(); + while let Some(arg) = it.next() { + match arg.as_str() { + "--help" | "-h" => return Ok(Command::Help), + "--version" | "-V" => return Ok(Command::Version), + "--model-dir" => { + let path = next_value(&mut it, &arg)?; + if path.is_empty() { + return Err("--model-dir needs a non-empty path".into()); + } + a.model_dir = path.into(); + } + "--prompt" | "-p" => a.prompt = next_value(&mut it, &arg)?, + "--max-new" | "-n" => { + a.max_new = next_value(&mut it, &arg)? + .parse() + .map_err(|_| "--max-new must be a positive integer")?; + if a.max_new == 0 { + return Err("--max-new must be positive".into()); + } + } + "--backend" | "-b" => { + a.backend = next_value(&mut it, &arg)?; + if !matches!(a.backend.as_str(), "cpu" | "metal") { + return Err("--backend must be cpu or metal".into()); + } + } + "--temperature" | "-t" => { + a.temperature = next_value(&mut it, &arg)? + .parse() + .map_err(|_| "--temperature must be a number")?; + if !a.temperature.is_finite() || a.temperature < 0.0 { + return Err("--temperature must be finite and non-negative".into()); + } + } + "--top-k" | "-k" => { + a.top_k = next_value(&mut it, &arg)? + .parse() + .map_err(|_| "--top-k must be a non-negative integer")? + } + "--seed" | "-s" => { + a.seed = next_value(&mut it, &arg)? + .parse() + .map_err(|_| "--seed must be an unsigned integer")? + } + "--greedy" => greedy = true, + other => return Err(format!("unknown argument: {other}")), + } + } + if greedy { + a.temperature = 0.0; + } + Ok(Command::Run(a)) +} + +fn print_help() { + println!("GPT-2 text generation\n\nUSAGE:\n generate [OPTIONS]\n\nOPTIONS:\n --model-dir PATH Directory containing all GPT-2 assets (default: models/gpt2)\n -p, --prompt TEXT Prompt (default: The meaning of life is)\n -n, --max-new N Maximum new tokens, positive integer (default: 40)\n -b, --backend cpu|metal Compute backend (default: metal on macOS, cpu elsewhere)\n -t, --temperature T Finite, non-negative temperature (default: 0.8)\n -k, --top-k K Keep K candidates; 0 disables filtering (default: 40)\n -s, --seed N Sampling seed (default: 42)\n --greedy Force greedy sampling, regardless of option order\n -h, --help Show help without loading model assets\n -V, --version Show package version"); +} + +/// Emits complete UTF-8 scalars, retaining an incomplete suffix until more bytes arrive. +fn drain_utf8(pending: &mut Vec, finish: bool) -> String { + let mut output = String::new(); + let mut used = 0; + while used < pending.len() { + match std::str::from_utf8(&pending[used..]) { + Ok(text) => { + output.push_str(text); + used = pending.len(); + } + Err(error) => { + let valid = error.valid_up_to(); + output.push_str(&String::from_utf8_lossy(&pending[used..used + valid])); + used += valid; + if let Some(invalid) = error.error_len() { + output.push('\u{fffd}'); + used += invalid; + } else if finish { + output.push('\u{fffd}'); + used = pending.len(); + } else { + break; + } + } + } + } + pending.drain(..used); + output +} + +fn run(backend: &B, model: &Gpt2, tok: &Tokenizer, args: &Args) -> io::Result<()> { + let sampler = Sampler { + temperature: args.temperature, + top_k: args.top_k, + seed: args.seed, + }; + let mut prompt_ids = tok.encode(&args.prompt); + if prompt_ids.is_empty() { + prompt_ids.push(EOT); + } + eprintln!( + "backend={} prompt_tokens={} max_new={} temp={} top_k={}", + backend.name(), + prompt_ids.len(), + args.max_new, + args.temperature, + args.top_k + ); + let mut stdout = io::stdout().lock(); + write!(stdout, "{}", args.prompt)?; + stdout.flush()?; + let mut count = 0; + let mut pending = Vec::new(); + let mut write_error = None; + let start = Instant::now(); + model.generate_while(backend, &prompt_ids, args.max_new, &sampler, EOT, |token| { + count += 1; + if token == EOT { + return true; + } + pending.extend(tok.decode_bytes(&[token])); + let text = drain_utf8(&mut pending, false); + if !text.is_empty() { + if let Err(error) = write!(stdout, "{text}").and_then(|_| stdout.flush()) { + write_error = Some(error); + return false; + } + } + true + }); + if let Some(error) = write_error { + return Err(error); + } + writeln!(stdout, "{}", drain_utf8(&mut pending, true))?; + stdout.flush()?; + let elapsed = start.elapsed(); + let rate = count as f64 / elapsed.as_secs_f64().max(f64::MIN_POSITIVE); + eprintln!( + "[{count} tokens in {elapsed:.2?} = {rate:.1} tok/s on {}]", + backend.name() + ); + Ok(()) +} + +fn output_result(result: io::Result<()>) -> Result<(), String> { + match result { + Ok(()) => Ok(()), + Err(error) if error.kind() == io::ErrorKind::BrokenPipe => Ok(()), + Err(error) => Err(format!("write generated text: {error}")), + } +} + +fn main() -> ExitCode { + let args = match parse_args(std::env::args().skip(1)) { + Ok(Command::Help) => { + print_help(); + return ExitCode::SUCCESS; + } + Ok(Command::Version) => { + println!("{}", env!("CARGO_PKG_VERSION")); + return ExitCode::SUCCESS; + } + Ok(Command::Run(args)) => args, + Err(error) => { + eprintln!("error: {error}\nUse --help for usage."); + return ExitCode::from(2); + } + }; + match execute(&args) { + Ok(()) => ExitCode::SUCCESS, + Err(error) => { + eprintln!("error: {error}"); + ExitCode::FAILURE + } + } +} + +fn execute(args: &Args) -> Result<(), String> { + #[cfg(not(target_os = "macos"))] + if args.backend == "metal" { + return Err("the Metal backend requires macOS; use --backend cpu".into()); + } + let model_path = args.model_dir.join("model.safetensors"); + if !model_path.is_file() { + return Err(format!("GPT-2 weights not found at {}. Download with python python/fetch_gpt2.py, or use --model-dir to locate your assets.", model_path.display())); + } + eprintln!("loading GPT-2 weights …"); + let tensors = loader::load_safetensors(&model_path) + .map_err(|e| format!("load {}: {e}", model_path.display()))?; + let model = + Gpt2::from_tensors(tensors, Config::default()).map_err(|e| format!("build GPT-2: {e}"))?; + let tok = Tokenizer::from_files( + &args.model_dir.join("vocab.json"), + &args.model_dir.join("merges.txt"), + ) + .map_err(|e| format!("load tokenizer from {}: {e}", args.model_dir.display()))?; + if tok.vocab_size() != model.config.vocab_size { + return Err(format!( + "tokenizer has {} tokens, but model expects {}", + tok.vocab_size(), + model.config.vocab_size + )); + } + + #[cfg(target_os = "macos")] + if args.backend == "metal" { + let metal = batch_forge::metal_backend::MetalBackend::new(batch_forge::SHADER_SOURCE) + .map_err(|e| format!("initialize Metal: {e}"))?; + return output_result(run(&metal, &model, &tok, args)); + } + output_result(run(&batch_forge::model::CpuBackend, &model, &tok, args)) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn parse(flags: &[&str]) -> Result { + parse_args(flags.iter().map(|s| s.to_string())) + } + + #[test] + fn rejects_invalid_flags_values_and_missing_values() { + for args in [ + vec!["--bogus"], + vec!["--seed"], + vec!["--max-new", "no"], + vec!["--max-new", "0"], + vec!["--backend", "cuda"], + vec!["--temperature", "NaN"], + vec!["--temperature", "inf"], + vec!["--temperature", "-1"], + vec!["--top-k", "-2"], + ] { + assert!(parse(&args).is_err(), "{args:?}"); + } + } + + #[test] + fn greedy_order_and_zero_sampling_values_are_explicit() { + for flags in [ + vec![ + "--greedy", + "--temperature", + "1", + "--top-k", + "0", + "--seed", + "0", + ], + vec![ + "--temperature", + "1", + "--greedy", + "--top-k", + "0", + "--seed", + "0", + ], + ] { + let Command::Run(args) = parse(&flags).unwrap() else { + panic!("expected run"); + }; + assert_eq!(args.temperature, 0.0); + assert_eq!(args.top_k, 0); + assert_eq!(args.seed, 0); + } + assert!(matches!(parse(&["--help"]), Ok(Command::Help))); + assert!(matches!(parse(&["--version"]), Ok(Command::Version))); + } + + #[test] + fn model_directory_override_is_preserved() { + let Command::Run(args) = parse(&["--model-dir", "/tmp/weights"]).unwrap() else { + panic!("expected run"); + }; + assert_eq!(args.model_dir, PathBuf::from("/tmp/weights")); + assert!(parse(&["--model-dir", ""]).is_err()); + } + + #[test] + fn utf8_scalar_split_over_three_tokens_streams_once() { + let mut pending = vec![0xe2]; + assert_eq!(drain_utf8(&mut pending, false), ""); + pending.push(0x82); + assert_eq!(drain_utf8(&mut pending, false), ""); + pending.push(0xac); + assert_eq!(drain_utf8(&mut pending, false), "€"); + assert!(pending.is_empty()); + } + + #[test] + fn invalid_and_unfinished_utf8_have_predictable_replacement() { + let mut pending = vec![b'a', 0xff, 0xe2]; + assert_eq!(drain_utf8(&mut pending, false), "a\u{fffd}"); + assert_eq!(drain_utf8(&mut pending, true), "\u{fffd}"); + assert!(pending.is_empty()); + assert!(output_result(Err(io::Error::from(io::ErrorKind::BrokenPipe))).is_ok()); + assert!(output_result(Err(io::Error::from(io::ErrorKind::PermissionDenied))).is_err()); + } +} diff --git a/src/engine.rs b/src/engine.rs new file mode 100644 index 0000000..3ee315b --- /dev/null +++ b/src/engine.rs @@ -0,0 +1,133 @@ +//! Asynchronous inference engine. +//! +//! A `tokio` task owns the model and backend and serves inference requests off +//! an `mpsc` queue, replying on a per-request `oneshot` channel. This is the +//! asynchronous request/response queue; forward computation is synchronous +//! inside the manager task. True continuous batching (fusing +//! queued requests into one dispatch) is future work — see the README roadmap. +//! +//! The engine is backend-agnostic (`Arc`), so it runs on the CPU +//! reference or the Metal backend without changes. + +use std::sync::Arc; + +use tokio::sync::{mpsc, oneshot}; +use tracing::{debug, info}; + +use crate::model::{Backend, Mlp, ModelError}; +use crate::tensor::Tensor; + +/// A unit of work: an input activation plus a channel to return the result on. +pub struct InferenceRequest { + pub request_id: u64, + pub input: Tensor, + pub response_tx: oneshot::Sender>, +} + +/// Owns the model + backend and serves requests from the queue. +pub struct RequestManager { + backend: Arc, + model: Arc, + request_rx: mpsc::Receiver, +} + +impl RequestManager { + pub fn new( + backend: Arc, + model: Arc, + request_rx: mpsc::Receiver, + ) -> Self { + Self { + backend, + model, + request_rx, + } + } + + /// Drains the queue until all senders are dropped, running one forward pass + /// per request on the configured backend. + pub async fn run(mut self) { + info!( + "RequestManager listening (backend: {})", + self.backend.name() + ); + let mut served = 0u64; + while let Some(req) = self.request_rx.recv().await { + if req.response_tx.is_closed() { + debug!(request_id = req.request_id, "skipped canceled request"); + continue; + } + let out = self.model.forward(&*self.backend, &req.input); + debug!(request_id = req.request_id, "served"); + // The receiver may have gone away; that's fine. + let _ = req.response_tx.send(out); + served += 1; + } + info!("RequestManager shutting down after {served} requests"); + } +} + +/// Convenience: spin up a manager on the current runtime and return a handle for +/// submitting requests. The manager stops when the returned `Submitter` (and all +/// its clones) are dropped. A zero queue depth is normalized to one slot. +pub fn spawn( + backend: Arc, + model: Arc, + queue_depth: usize, +) -> Submitter { + let (tx, rx) = mpsc::channel(queue_depth.max(1)); + let manager = RequestManager::new(backend, model, rx); + tokio::spawn(manager.run()); + Submitter { tx } +} + +/// Cloneable client handle for submitting inference requests to the engine. +#[derive(Clone)] +pub struct Submitter { + tx: mpsc::Sender, +} + +impl Submitter { + /// Submits one request and awaits its result. + pub async fn infer(&self, request_id: u64, input: Tensor) -> Result { + let (response_tx, response_rx) = oneshot::channel(); + self.tx + .send(InferenceRequest { + request_id, + input, + response_tx, + }) + .await + .map_err(|_| ModelError::EngineClosed)?; + response_rx.await.map_err(|_| ModelError::EngineClosed)? + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn closed_engine_returns_an_error() { + let (tx, rx) = mpsc::channel(1); + drop(rx); + let result = Submitter { tx }.infer(0, Tensor::zeros(vec![1, 1])).await; + assert!(matches!(result, Err(ModelError::EngineClosed))); + } + + #[tokio::test] + async fn zero_queue_depth_still_accepts_a_request() { + let model = Mlp { + layers: vec![( + Tensor::new(vec![2.0], vec![1, 1]).unwrap(), + Tensor::zeros(vec![1]), + )], + }; + let submitter = spawn(Arc::new(crate::model::CpuBackend), Arc::new(model), 0); + let out = submitter + .infer(0, Tensor::new(vec![3.0], vec![1, 1]).unwrap()) + .await + .unwrap(); + assert_eq!(out.data, vec![6.0]); + } +} diff --git a/src/gpt2.rs b/src/gpt2.rs new file mode 100644 index 0000000..cb2f743 --- /dev/null +++ b/src/gpt2.rs @@ -0,0 +1,741 @@ +//! GPT-2 inference. +//! +//! A from-scratch GPT-2 forward pass that runs on either the CPU reference or +//! the Metal backend through the [`LlmOps`] trait, so the two are checked for +//! parity exactly like the individual kernels. Heavy ops (matmul, attention, +//! layernorm, gelu) go through the backend; the cheap element-wise glue +//! (embedding gather, bias add, residual) stays in plain Rust. +//! +//! Weight layout follows HuggingFace `gpt2`: the `Conv1D` layers store weights +//! as `[in, out]` so `y = x @ W + b` is a plain matmul (no transpose), and the +//! LM head is tied to the token embedding. + +use std::collections::HashMap; + +use thiserror::Error; + +use crate::ops; +use crate::tensor::{Tensor, TensorError}; + +#[derive(Error, Debug)] +pub enum Gpt2Error { + #[error("invalid sampling input: {0}")] + InvalidSampling(&'static str), + #[error("tensor '{name}' has shape {found:?}; expected {expected:?}")] + BadShape { + name: String, + expected: Vec, + found: Vec, + }, + #[error("invalid tensor: {0}")] + Tensor(#[from] TensorError), + #[error("invalid GPT-2 configuration: {0}")] + InvalidConfig(&'static str), + #[error("missing tensor '{0}'")] + Missing(String), +} + +/// GPT-2 hyperparameters. Defaults are the 124M ("small") configuration. +#[derive(Debug, Clone, Copy)] +pub struct Config { + pub n_layer: usize, + pub n_head: usize, + pub n_embd: usize, + pub n_ctx: usize, + pub vocab_size: usize, + pub eps: f32, +} + +impl Default for Config { + fn default() -> Self { + Config { + n_layer: 12, + n_head: 12, + n_embd: 768, + n_ctx: 1024, + vocab_size: 50257, + eps: 1e-5, + } + } +} + +impl Config { + /// Validates dimensions before allocating model buffers or splitting attention heads. + pub fn validate(&self) -> Result<(), Gpt2Error> { + if self.n_layer == 0 + || self.n_head == 0 + || self.n_embd == 0 + || self.n_ctx == 0 + || self.vocab_size == 0 + { + return Err(Gpt2Error::InvalidConfig("all dimensions must be positive")); + } + if self.n_embd % self.n_head != 0 { + return Err(Gpt2Error::InvalidConfig( + "embedding width must be divisible by head count", + )); + } + if !self.eps.is_finite() || self.eps <= 0.0 { + return Err(Gpt2Error::InvalidConfig( + "normalization epsilon must be finite and positive", + )); + } + if self.n_embd.checked_mul(4).is_none() { + return Err(Gpt2Error::InvalidConfig("projection width overflows")); + } + Ok(()) + } + + pub fn head_dim(&self) -> usize { + self.n_embd / self.n_head + } +} + +/// The compute surface GPT-2 needs from a backend. Implemented by both the CPU +/// reference and the Metal backend; the model is generic over it. +pub trait LlmOps { + /// `C[m,n] = A[m,k] * B[k,n]` (row-major). Used for the Conv1D projections. + fn mm(&self, a: &[f32], b: &[f32], m: usize, k: usize, n: usize) -> Vec; + /// `y = x · Wᵀ` with `w` = `[out,in]`; used for the tied LM head. + fn lm_head(&self, x: &[f32], w: &[f32], rows: usize, in_f: usize, out_f: usize) -> Vec; + fn layernorm( + &self, + x: &[f32], + g: &[f32], + b: &[f32], + rows: usize, + d: usize, + eps: f32, + ) -> Vec; + fn gelu(&self, x: &[f32]) -> Vec; + fn mha( + &self, + q: &[f32], + k: &[f32], + v: &[f32], + seq: usize, + heads: usize, + head_dim: usize, + ) -> Vec; + fn name(&self) -> &'static str; +} + +impl LlmOps for crate::model::CpuBackend { + fn mm(&self, a: &[f32], b: &[f32], m: usize, k: usize, n: usize) -> Vec { + ops::matmul(a, b, m, k, n) + } + fn lm_head(&self, x: &[f32], w: &[f32], rows: usize, in_f: usize, out_f: usize) -> Vec { + let zeros = vec![0.0f32; out_f]; + ops::linear(x, w, &zeros, rows, in_f, out_f) + } + fn layernorm( + &self, + x: &[f32], + g: &[f32], + b: &[f32], + rows: usize, + d: usize, + eps: f32, + ) -> Vec { + ops::layernorm(x, g, b, rows, d, eps) + } + fn gelu(&self, x: &[f32]) -> Vec { + let mut y = x.to_vec(); + ops::gelu_inplace(&mut y); + y + } + fn mha( + &self, + q: &[f32], + k: &[f32], + v: &[f32], + seq: usize, + heads: usize, + head_dim: usize, + ) -> Vec { + ops::mha(q, k, v, seq, heads, head_dim) + } + fn name(&self) -> &'static str { + "cpu" + } +} + +#[cfg(target_os = "macos")] +impl LlmOps for crate::metal_backend::MetalBackend { + fn mm(&self, a: &[f32], b: &[f32], m: usize, k: usize, n: usize) -> Vec { + self.matmul_tiled(a, b, m, k, n) + } + fn lm_head(&self, x: &[f32], w: &[f32], rows: usize, in_f: usize, out_f: usize) -> Vec { + let zeros = vec![0.0f32; out_f]; + self.linear(x, w, &zeros, rows, in_f, out_f) + } + fn layernorm( + &self, + x: &[f32], + g: &[f32], + b: &[f32], + rows: usize, + d: usize, + eps: f32, + ) -> Vec { + crate::metal_backend::MetalBackend::layernorm(self, x, g, b, rows, d, eps) + } + fn gelu(&self, x: &[f32]) -> Vec { + crate::metal_backend::MetalBackend::gelu(self, x) + } + fn mha( + &self, + q: &[f32], + k: &[f32], + v: &[f32], + seq: usize, + heads: usize, + head_dim: usize, + ) -> Vec { + crate::metal_backend::MetalBackend::mha(self, q, k, v, seq, heads, head_dim) + } + fn name(&self) -> &'static str { + "metal" + } +} + +struct Layer { + ln1_w: Vec, + ln1_b: Vec, + attn_w: Vec, // c_attn.weight [n_embd, 3*n_embd] + attn_b: Vec, // [3*n_embd] + proj_w: Vec, // c_proj.weight [n_embd, n_embd] + proj_b: Vec, + ln2_w: Vec, + ln2_b: Vec, + fc_w: Vec, // mlp.c_fc.weight [n_embd, 4*n_embd] + fc_b: Vec, + fc_proj_w: Vec, // mlp.c_proj.weight [4*n_embd, n_embd] + fc_proj_b: Vec, +} + +/// A loaded GPT-2 model. All weights are owned `f32`. +pub struct Gpt2 { + pub config: Config, + wte: Vec, // [vocab, n_embd] + wpe: Vec, // [n_ctx, n_embd] + layers: Vec, + lnf_w: Vec, + lnf_b: Vec, +} + +fn take( + map: &mut HashMap, + name: &str, + shape: &[usize], +) -> Result, Gpt2Error> { + let tensor = map + .remove(name) + .ok_or_else(|| Gpt2Error::Missing(name.to_string()))?; + if tensor.shape != shape { + return Err(Gpt2Error::BadShape { + name: name.to_string(), + expected: shape.to_vec(), + found: tensor.shape, + }); + } + tensor.validate()?; + Ok(tensor.data) +} + +impl Gpt2 { + /// Builds a model from a HuggingFace `gpt2` safetensors tensor map. + pub fn from_tensors( + mut map: HashMap, + config: Config, + ) -> Result { + config.validate()?; + let d = config.n_embd; + let wte = take(&mut map, "wte.weight", &[config.vocab_size, d])?; + let wpe = take(&mut map, "wpe.weight", &[config.n_ctx, d])?; + let lnf_w = take(&mut map, "ln_f.weight", &[d])?; + let lnf_b = take(&mut map, "ln_f.bias", &[d])?; + let mut layers = Vec::with_capacity(config.n_layer); + for i in 0..config.n_layer { + let p = format!("h.{i}."); + layers.push(Layer { + ln1_w: take(&mut map, &format!("{p}ln_1.weight"), &[d])?, + ln1_b: take(&mut map, &format!("{p}ln_1.bias"), &[d])?, + attn_w: take(&mut map, &format!("{p}attn.c_attn.weight"), &[d, 3 * d])?, + attn_b: take(&mut map, &format!("{p}attn.c_attn.bias"), &[3 * d])?, + proj_w: take(&mut map, &format!("{p}attn.c_proj.weight"), &[d, d])?, + proj_b: take(&mut map, &format!("{p}attn.c_proj.bias"), &[d])?, + ln2_w: take(&mut map, &format!("{p}ln_2.weight"), &[d])?, + ln2_b: take(&mut map, &format!("{p}ln_2.bias"), &[d])?, + fc_w: take(&mut map, &format!("{p}mlp.c_fc.weight"), &[d, 4 * d])?, + fc_b: take(&mut map, &format!("{p}mlp.c_fc.bias"), &[4 * d])?, + fc_proj_w: take(&mut map, &format!("{p}mlp.c_proj.weight"), &[4 * d, d])?, + fc_proj_b: take(&mut map, &format!("{p}mlp.c_proj.bias"), &[d])?, + }); + } + Ok(Self { + config, + wte, + wpe, + layers, + lnf_w, + lnf_b, + }) + } + + /// Runs the forward pass over `tokens` and returns the logits for the final + /// position only (`[vocab_size]`), which is all generation needs. + pub fn forward(&self, backend: &B, tokens: &[usize]) -> Vec { + let cfg = self.config; + cfg.validate().expect("invalid GPT-2 configuration"); + assert!( + !tokens.is_empty(), + "GPT-2 forward requires at least one token" + ); + assert!( + tokens.len() <= cfg.n_ctx, + "token sequence exceeds GPT-2 context length" + ); + assert!( + tokens.iter().all(|&token| token < cfg.vocab_size), + "token ID exceeds GPT-2 vocabulary" + ); + let (seq, d) = (tokens.len(), cfg.n_embd); + let eps = cfg.eps; + + // Token + positional embeddings. + let mut x = vec![0.0f32; seq * d]; + for (i, &tok) in tokens.iter().enumerate() { + let wt = &self.wte[tok * d..tok * d + d]; + let wp = &self.wpe[i * d..i * d + d]; + for j in 0..d { + x[i * d + j] = wt[j] + wp[j]; + } + } + + for layer in &self.layers { + // --- attention block --- + let ln1 = backend.layernorm(&x, &layer.ln1_w, &layer.ln1_b, seq, d, eps); + let mut qkv = backend.mm(&ln1, &layer.attn_w, seq, d, 3 * d); + add_bias(&mut qkv, &layer.attn_b, seq, 3 * d); + let (q, k, v) = split_qkv(&qkv, seq, d); + let attn = backend.mha(&q, &k, &v, seq, cfg.n_head, cfg.head_dim()); + let mut proj = backend.mm(&attn, &layer.proj_w, seq, d, d); + add_bias(&mut proj, &layer.proj_b, seq, d); + residual_add(&mut x, &proj); + + // --- MLP block --- + let ln2 = backend.layernorm(&x, &layer.ln2_w, &layer.ln2_b, seq, d, eps); + let mut fc = backend.mm(&ln2, &layer.fc_w, seq, d, 4 * d); + add_bias(&mut fc, &layer.fc_b, seq, 4 * d); + let act = backend.gelu(&fc); + let mut fc2 = backend.mm(&act, &layer.fc_proj_w, seq, 4 * d, d); + add_bias(&mut fc2, &layer.fc_proj_b, seq, d); + residual_add(&mut x, &fc2); + } + + let xf = backend.layernorm(&x, &self.lnf_w, &self.lnf_b, seq, d, eps); + let last = &xf[(seq - 1) * d..seq * d]; + // Tied LM head: logits = last · wteᵀ. + backend.lm_head(last, &self.wte, 1, d, cfg.vocab_size) + } + + /// Autoregressively generates up to `max_new` tokens, returning the full + /// token sequence (prompt + generated). Stops early on the end-of-text token. + pub fn generate( + &self, + backend: &B, + prompt: &[usize], + max_new: usize, + sampler: &Sampler, + eot_token: usize, + mut on_token: impl FnMut(usize), + ) -> Vec { + self.generate_while(backend, prompt, max_new, sampler, eot_token, |token| { + on_token(token); + true + }) + } + + /// Generates until end-of-text, max_new, or the callback returns false. + pub fn generate_while( + &self, + backend: &B, + prompt: &[usize], + max_new: usize, + sampler: &Sampler, + eot_token: usize, + mut on_token: impl FnMut(usize) -> bool, + ) -> Vec { + let mut toks = prompt.to_vec(); + if max_new == 0 { + return toks; + } + self.config.validate().expect("invalid GPT-2 configuration"); + sampler.validate().expect("invalid generation sampler"); + assert!(!prompt.is_empty(), "generation requires a non-empty prompt"); + assert!( + eot_token < self.config.vocab_size, + "end-of-text ID exceeds GPT-2 vocabulary" + ); + assert!( + prompt.iter().all(|&token| token < self.config.vocab_size), + "prompt ID exceeds GPT-2 vocabulary" + ); + let n_ctx = self.config.n_ctx; + let mut rng = Rng::new(sampler.seed); + for _ in 0..max_new { + let start = toks.len().saturating_sub(n_ctx); + let logits = self.forward(backend, &toks[start..]); + let next = sampler.sample(&logits, &mut rng); + toks.push(next); + if !on_token(next) || next == eot_token { + break; + } + } + toks + } +} + +fn add_bias(x: &mut [f32], bias: &[f32], rows: usize, cols: usize) { + debug_assert_eq!(bias.len(), cols); + for r in 0..rows { + for c in 0..cols { + x[r * cols + c] += bias[c]; + } + } +} + +fn residual_add(x: &mut [f32], y: &[f32]) { + for (xi, yi) in x.iter_mut().zip(y) { + *xi += yi; + } +} + +/// Splits a `[seq, 3*d]` QKV tensor into three contiguous `[seq, d]` tensors. +fn split_qkv(qkv: &[f32], seq: usize, d: usize) -> (Vec, Vec, Vec) { + let mut q = vec![0.0f32; seq * d]; + let mut k = vec![0.0f32; seq * d]; + let mut v = vec![0.0f32; seq * d]; + for i in 0..seq { + let row = &qkv[i * 3 * d..i * 3 * d + 3 * d]; + q[i * d..i * d + d].copy_from_slice(&row[0..d]); + k[i * d..i * d + d].copy_from_slice(&row[d..2 * d]); + v[i * d..i * d + d].copy_from_slice(&row[2 * d..3 * d]); + } + (q, k, v) +} + +/// Token sampling strategy. +pub struct Sampler { + pub temperature: f32, + pub top_k: usize, + pub seed: u64, +} + +impl Sampler { + /// Greedy (argmax) sampling. + pub fn greedy() -> Self { + Sampler { + temperature: 0.0, + top_k: 0, + seed: 0, + } + } + + pub fn validate(&self) -> Result<(), Gpt2Error> { + if !self.temperature.is_finite() || self.temperature < 0.0 { + return Err(Gpt2Error::InvalidSampling( + "temperature must be finite and non-negative", + )); + } + Ok(()) + } + + /// Samples one token using this seed. Generation retains the RNG across tokens. + pub fn sample_token(&self, logits: &[f32]) -> Result { + self.sample_checked(logits, &mut Rng::new(self.seed)) + } + + fn sample(&self, logits: &[f32], rng: &mut Rng) -> usize { + self.sample_checked(logits, rng) + .expect("invalid token sampling input") + } + + fn sample_checked(&self, logits: &[f32], rng: &mut Rng) -> Result { + self.validate()?; + if logits.is_empty() { + return Err(Gpt2Error::InvalidSampling("logits must not be empty")); + } + if logits.iter().any(|logit| !logit.is_finite()) { + return Err(Gpt2Error::InvalidSampling("logits must be finite")); + } + if self.temperature == 0.0 { + return Ok(argmax(logits)); + } + let mut idx: Vec = (0..logits.len()).collect(); + if self.top_k > 0 && self.top_k < logits.len() { + idx.select_nth_unstable_by(self.top_k, |&a, &b| { + if logits[a] == logits[b] { + a.cmp(&b) + } else { + logits[b].total_cmp(&logits[a]) + } + }); + idx.truncate(self.top_k); + idx.sort_unstable(); + } + let max = idx + .iter() + .map(|&i| logits[i]) + .fold(f32::NEG_INFINITY, f32::max); + let mut probs: Vec = idx + .iter() + .map(|&i| ((f64::from(logits[i]) - f64::from(max)) / f64::from(self.temperature)).exp()) + .collect(); + let sum: f64 = probs.iter().sum(); + for p in &mut probs { + *p /= sum; + } + Ok(select_probability(&idx, &probs, f64::from(rng.next_f32()))) + } +} + +fn select_probability(indices: &[usize], probabilities: &[f64], draw: f64) -> usize { + let mut cumulative = 0.0; + for (index, &probability) in probabilities.iter().enumerate() { + cumulative += probability; + if draw < cumulative { + return indices[index]; + } + } + // Rounding near one must still select a candidate with positive weight. + let last = probabilities + .iter() + .rposition(|&p| p > 0.0) + .expect("sampling has no positive weight"); + indices[last] +} + +fn argmax(v: &[f32]) -> usize { + let mut best = 0; + for i in 1..v.len() { + if v[i] > v[best] { + best = i; + } + } + best +} + +/// Small xorshift RNG for sampling (keeps generation dependency-free). +struct Rng(u64); +impl Rng { + fn new(seed: u64) -> Self { + Rng(seed ^ 0x9E3779B97F4A7C15) + } + fn next_f32(&mut self) -> f32 { + let mut x = self.0.max(1); + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + self.0 = x; + (x >> 40) as f32 / (1u64 << 24) as f32 + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn invalid_configs_fail_before_loading_weights() { + for config in [ + Config { + n_head: 0, + ..Config::default() + }, + Config { + n_embd: 7, + ..Config::default() + }, + Config { + eps: f32::NAN, + ..Config::default() + }, + Config { + n_ctx: 0, + ..Config::default() + }, + ] { + assert!(matches!( + Gpt2::from_tensors(HashMap::new(), config), + Err(Gpt2Error::InvalidConfig(_)) + )); + } + } + + #[test] + fn wrong_embedding_shape_is_rejected_at_load_time() { + let config = Config { + n_layer: 1, + n_head: 1, + n_embd: 2, + n_ctx: 2, + vocab_size: 2, + eps: 1e-5, + }; + let mut map = HashMap::new(); + map.insert("wte.weight".into(), Tensor::zeros(vec![1, 4])); + assert!(matches!( + Gpt2::from_tensors(map, config), + Err(Gpt2Error::BadShape { .. }) + )); + } + + fn forward_test_model() -> Gpt2 { + Gpt2 { + config: Config { + n_layer: 1, + n_head: 1, + n_embd: 2, + n_ctx: 2, + vocab_size: 2, + eps: 1e-5, + }, + wte: vec![0.0; 4], + wpe: vec![0.0; 4], + layers: Vec::new(), + lnf_w: vec![1.0; 2], + lnf_b: vec![0.0; 2], + } + } + + #[test] + #[should_panic(expected = "requires at least one token")] + fn empty_forward_has_a_clear_precondition() { + forward_test_model().forward(&crate::model::CpuBackend, &[]); + } + + #[test] + #[should_panic(expected = "exceeds GPT-2 vocabulary")] + fn out_of_range_token_has_a_clear_precondition() { + forward_test_model().forward(&crate::model::CpuBackend, &[2]); + } + + #[test] + #[should_panic(expected = "exceeds GPT-2 context length")] + fn overlong_forward_has_a_clear_precondition() { + forward_test_model().forward(&crate::model::CpuBackend, &[0, 0, 0]); + } + + #[test] + fn generation_callback_can_stop_after_one_token() { + let mut calls = 0; + let tokens = forward_test_model().generate_while( + &crate::model::CpuBackend, + &[0], + 10, + &Sampler::greedy(), + 1, + |_| { + calls += 1; + false + }, + ); + assert_eq!(calls, 1); + assert_eq!(tokens.len(), 2); + } + + #[test] + fn checked_sampling_rejects_invalid_inputs() { + for temperature in [f32::NAN, f32::INFINITY, -1.0] { + assert!(Sampler { + temperature, + top_k: 1, + seed: 0 + } + .sample_token(&[0.0]) + .is_err()); + } + assert!(Sampler::greedy().sample_token(&[]).is_err()); + for value in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] { + assert!(Sampler::greedy().sample_token(&[value, 0.0]).is_err()); + } + } + + #[test] + fn extreme_finite_logits_sample_without_overflow() { + let sampler = Sampler { + temperature: f32::MIN_POSITIVE, + top_k: 0, + seed: 42, + }; + assert_eq!(sampler.sample_token(&[-f32::MAX, f32::MAX]).unwrap(), 1); + } + + #[test] + fn top_k_ties_choose_lowest_ids_and_repeat_from_a_seed() { + let sampler = Sampler { + temperature: 1.0, + top_k: 1, + seed: 42, + }; + assert_eq!(sampler.sample_token(&[0.0, 0.0, 0.0]).unwrap(), 0); + let sampler = Sampler { + top_k: 2, + ..sampler + }; + let first = sampler.sample_token(&[0.0, 0.0, 0.0]).unwrap(); + assert!(first < 2); + assert_eq!(first, sampler.sample_token(&[0.0, 0.0, 0.0]).unwrap()); + } + + #[test] + fn sampling_boundaries_never_choose_zero_weight() { + assert_eq!(select_probability(&[10, 20], &[0.0, 1.0], 0.0), 20); + assert_eq!( + select_probability( + &[10, 20, 30], + &[0.5, 0.4999999999999999, 0.0], + 0.9999999999999999 + ), + 20 + ); + } + + #[test] + fn zero_token_generation_returns_the_prompt_without_inference() { + let tokens = forward_test_model().generate( + &crate::model::CpuBackend, + &[], + 0, + &Sampler::greedy(), + 1, + |_| panic!("unexpected inference"), + ); + assert!(tokens.is_empty()); + } + + #[test] + #[should_panic(expected = "end-of-text ID exceeds")] + fn generation_rejects_unknown_stop_tokens() { + forward_test_model().generate( + &crate::model::CpuBackend, + &[0], + 1, + &Sampler::greedy(), + 999, + |_| {}, + ); + } + + #[test] + fn top_k_ties_treat_signed_zero_as_equal() { + let sampler = Sampler { + temperature: 1.0, + top_k: 1, + seed: 42, + }; + assert_eq!(sampler.sample_token(&[-0.0, 0.0]).unwrap(), 0); + } +} diff --git a/src/kv_cache/mod.rs b/src/kv_cache/mod.rs index 176dbe7..6f82bda 100644 --- a/src/kv_cache/mod.rs +++ b/src/kv_cache/mod.rs @@ -1,5 +1,22 @@ +//! Per-request Key/Value cache stored in Metal unified memory. +//! +//! Buffers are pre-allocated to `max_len` rows. Appends are bounds-checked: +//! previously `current_len` was advanced unconditionally, so a sequence longer +//! than `max_len` would write past the end of the cache buffers. + use metal::{Buffer, Device, MTLResourceOptions}; use std::collections::HashMap; +use thiserror::Error; + +#[derive(Error, Debug)] +pub enum KvCacheError { + #[error("KV cache full: have {current}/{max} rows, cannot append {requested}")] + Full { + current: usize, + max: usize, + requested: usize, + }, +} pub struct KVCache { pub k_buffer: Buffer, @@ -11,10 +28,14 @@ pub struct KVCache { impl KVCache { pub fn new(device: &Device, max_len: usize, head_dim: usize) -> Self { - let buffer_size = (max_len * head_dim * std::mem::size_of::()) as u64; - let k_buffer = device.new_buffer(buffer_size, MTLResourceOptions::StorageModeShared); - let v_buffer = device.new_buffer(buffer_size, MTLResourceOptions::StorageModeShared); - + let buffer_size = max_len + .checked_mul(head_dim) + .and_then(|v| v.checked_mul(std::mem::size_of::())) + .expect("KV Cache buffer size overflowed"); + + let k_buffer = device.new_buffer(buffer_size as u64, MTLResourceOptions::StorageModeShared); + let v_buffer = device.new_buffer(buffer_size as u64, MTLResourceOptions::StorageModeShared); + Self { k_buffer, v_buffer, @@ -23,6 +44,26 @@ impl KVCache { head_dim, } } + + /// Whether `m` more rows fit without overflowing the pre-allocated buffers. + pub fn can_fit(&self, m: usize) -> bool { + self.current_len.saturating_add(m) <= self.max_len + } + + /// Reserves room for `m` new rows and returns the offset they should be + /// written at, advancing `current_len`. Errors instead of overflowing. + pub fn advance(&mut self, m: usize) -> Result { + if !self.can_fit(m) { + return Err(KvCacheError::Full { + current: self.current_len, + max: self.max_len, + requested: m, + }); + } + let offset = self.current_len; + self.current_len += m; + Ok(offset) + } } pub struct KVStorage { @@ -43,12 +84,22 @@ impl KVStorage { } pub fn get_or_create(&mut self, request_id: u64) -> &mut KVCache { - self.caches.entry(request_id).or_insert_with(|| { - KVCache::new(&self.device, self.max_seq_len, self.head_dim) - }) + let device = &self.device; + let (max_seq_len, head_dim) = (self.max_seq_len, self.head_dim); + self.caches + .entry(request_id) + .or_insert_with(|| KVCache::new(device, max_seq_len, head_dim)) } pub fn remove(&mut self, request_id: u64) { self.caches.remove(&request_id); } + + pub fn len(&self) -> usize { + self.caches.len() + } + + pub fn is_empty(&self) -> bool { + self.caches.is_empty() + } } diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..10e67a9 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,35 @@ +//! batch_forge — a small, correctness-first inference runtime for Apple Silicon. +//! +//! The crate is split into a portable CPU reference (`ops`) that defines the +//! ground-truth numerics for every operator, and a Metal backend +//! (`metal_backend`) whose kernels are validated against that reference by the +//! parity tests in `tests/`. This mirrors how production engines (ggml, candle) +//! keep a CPU reference next to each accelerated kernel. On non-macOS targets, +//! the portable CPU modules remain available while the Metal-specific modules +//! are omitted. + +/// Asynchronous request/response inference engine built on Tokio channels. +pub mod engine; +/// From-scratch GPT-2 inference shared by the CPU and Metal backends. +pub mod gpt2; +/// Zero-copy SafeTensors loading backed by memory mapping. +pub mod loader; +/// Model definitions and backend abstractions used to run them. +pub mod model; +/// Pure-Rust reference implementations used to verify accelerated kernels. +pub mod ops; +/// Lightweight tensor types and dtype conversions shared by the backends. +pub mod tensor; +/// From-scratch byte-level BPE tokenization for GPT-2-compatible vocabularies. +pub mod tokenizer; + +#[cfg(target_os = "macos")] +pub mod kv_cache; + +/// Metal compute backend, available when building on macOS. +#[cfg(target_os = "macos")] +pub mod metal_backend; + +/// Default Metal shader source, embedded at compile time. +#[cfg(target_os = "macos")] +pub const SHADER_SOURCE: &str = include_str!("shaders/compute.metal"); diff --git a/src/loader.rs b/src/loader.rs index 88930a6..05fa543 100644 --- a/src/loader.rs +++ b/src/loader.rs @@ -1,11 +1,19 @@ +//! Safetensors loading over `mmap`. +//! +//! [`SafeModel`] owns the memory map for the lifetime of the handle, so tensor +//! views borrow from a stable backing store. This replaces the previous +//! `transmute` + `mem::forget` approach, which laundered a borrowed slice into +//! `&'static` and leaked the mapping on every load. + use std::collections::HashMap; use std::fs::File; use std::path::Path; -use memmap2::MmapOptions; + +use memmap2::Mmap; use safetensors::SafeTensors; use thiserror::Error; -use crate::tensor::{DataType, TensorView}; +use crate::tensor::{DataType, Tensor, TensorError, TensorView}; #[derive(Error, Debug)] pub enum LoaderError { @@ -13,29 +21,60 @@ pub enum LoaderError { Io(#[from] std::io::Error), #[error("Safetensors error: {0}")] SafeTensors(#[from] safetensors::SafeTensorError), + #[error("Tensor error: {0}")] + Tensor(#[from] TensorError), } -/// Loads a Safetensors file via memory mapping, returning zero-copy tensor views. -pub fn load_safetensors<'a>(path: &Path) -> Result>, LoaderError> { - let file = File::open(path)?; - let mmap = unsafe { MmapOptions::new().map(&file)? }; - - // We use Box::leak to keep the mmap object alive for the lifetime of the program, - // which allows us to have a &'static [u8] view of the memory-mapped file. - let mmap_ref: &'static [u8] = unsafe { std::mem::transmute(&mmap[..]) }; - std::mem::forget(mmap); // Prevent mmap from being dropped +/// A memory-mapped safetensors checkpoint. Holding this handle keeps the mapping +/// alive; tensor views and owned tensors are produced on demand from it. +pub struct SafeModel { + mmap: Mmap, +} - let st = SafeTensors::deserialize(mmap_ref)?; - let mut tensors = HashMap::new(); +impl SafeModel { + /// Opens and memory-maps a safetensors file. The header is not parsed until + /// [`SafeModel::with_tensors`] or [`SafeModel::load_f32`] is called. + pub fn open(path: &Path) -> Result { + let file = File::open(path)?; + // SAFETY: the file is not mutated for the lifetime of the mapping; the + // mapping is owned by `self` and dropped (unmapped) with it. + let mmap = unsafe { Mmap::map(&file)? }; + Ok(Self { mmap }) + } - for name in st.names() { - let view = st.tensor(name)?; - let dtype = DataType::from(view.dtype()); - let shape = view.shape().to_vec(); - let data = view.data(); + /// Parses the checkpoint and invokes `f` with zero-copy views that borrow + /// from the mapping. The scoped-closure shape keeps the views' lifetime tied + /// to the borrow of `self`, which is what makes this sound *and* copy-free + /// (e.g. uploading weights straight to a GPU buffer without a heap copy). + pub fn with_tensors( + &self, + f: impl FnOnce(&HashMap>) -> R, + ) -> Result { + let st = SafeTensors::deserialize(&self.mmap)?; + let mut views = HashMap::new(); + for name in st.names() { + let raw = st.tensor(name)?; + let dtype = DataType::try_from(raw.dtype())?; + let view = TensorView::new(raw.shape(), dtype, raw.data())?; + views.insert(name.to_string(), view); + } + Ok(f(&views)) + } - tensors.insert(name.to_string(), TensorView::new(shape, dtype, data)); + /// Loads every F32 tensor into owned [`Tensor`]s. Non-F32 tensors cause a + /// [`TensorError::NotF32`]; the demo model is exported entirely in F32. + pub fn load_f32(&self) -> Result, LoaderError> { + self.with_tensors(|views| { + let mut out = HashMap::with_capacity(views.len()); + for (name, view) in views { + out.insert(name.clone(), view.to_tensor_f32()?); + } + Ok(out) + })? } +} - Ok(tensors) +/// Convenience: open `path` and load all F32 tensors as owned [`Tensor`]s. +pub fn load_safetensors(path: &Path) -> Result, LoaderError> { + SafeModel::open(path)?.load_f32() } diff --git a/src/main.rs b/src/main.rs index 9d87331..141896e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,144 +1,335 @@ -use tracing::{info, error}; -use tokio::sync::{mpsc, oneshot}; +//! batch_forge CLI. +//! +//! Loads an exported model, runs its forward pass on the CPU reference and (on +//! Apple Silicon) the Metal backend, cross-checks the two, and optionally +//! verifies against a saved JAX/NumPy reference output. `--requests N` exercises +//! the async engine with N concurrent in-flight requests. + +use std::path::{Path, PathBuf}; +use std::process::ExitCode; use std::sync::Arc; -use std::path::PathBuf; -mod tensor; -mod loader; +use tracing::{error, info, warn}; -#[cfg(target_os = "macos")] -mod metal_backend; +use batch_forge::loader; +use batch_forge::model::{Backend, CpuBackend, Mlp}; +use batch_forge::tensor::Tensor; -#[cfg(target_os = "macos")] -mod kv_cache; +/// Tolerance for the reference-verification pass/fail check. +const VERIFY_TOL: f32 = 1e-3; -/// Represents a request for model inference. #[derive(Debug)] -struct InferenceRequest { - request_id: u64, - input_data: Vec, - response_tx: oneshot::Sender>, +struct Args { + model: PathBuf, + verify: Option, + backend: BackendChoice, + requests: usize, } -/// Manages the asynchronous request queue and dynamic batching. -struct RequestManager { - backend: Arc, - kv_storage: kv_cache::KVStorage, - request_rx: mpsc::Receiver, +#[derive(Debug, Clone, Copy, PartialEq)] +enum BackendChoice { + Cpu, + Metal, + Both, } -impl RequestManager { - pub fn new(backend: Arc, rx: mpsc::Receiver) -> Self { - let kv_storage = kv_cache::KVStorage::new(backend.device.clone(), 1024, 64); - Self { backend, kv_storage, request_rx: rx } +fn parse_args() -> Result { + let mut model = PathBuf::from("model.safetensors"); + let mut verify = None; + let default_backend = if cfg!(target_os = "macos") { + BackendChoice::Both + } else { + BackendChoice::Cpu + }; + let mut backend = default_backend; + let mut requests = 0usize; + + let mut args = std::env::args().skip(1); + while let Some(arg) = args.next() { + match arg.as_str() { + "--model" | "-m" => model = args.next().ok_or("--model needs a path")?.into(), + "--verify" | "-v" => verify = Some(args.next().ok_or("--verify needs a path")?.into()), + "--backend" | "-b" => { + backend = match args.next().as_deref() { + Some("cpu") => BackendChoice::Cpu, + Some("metal") => BackendChoice::Metal, + Some("both") => BackendChoice::Both, + other => return Err(format!("unknown backend {other:?}")), + } + } + "--requests" | "-r" => { + requests = args + .next() + .ok_or("--requests needs a number")? + .parse() + .map_err(|_| "--requests must be an integer")?; + } + "--help" | "-h" => return Err("help".into()), + "--prompt" | "-p" => return Err("text prompts require the generate binary: cargo run --bin generate -- --prompt TEXT".into()), + other => return Err(format!("unknown argument: {other}")), + } } + Ok(Args { + model, + verify, + backend, + requests, + }) +} + +fn print_help() { + println!( + "batch_forge {} — verified Metal inference for JAX/Equinox models\n\ +\n\ +USAGE:\n\ + batch_forge [--model PATH] [--backend cpu|metal|both] [--verify PATH] [--requests N]\n\ +\n\ +OPTIONS:\n\ + -m, --model PATH Safetensors checkpoint (default: model.safetensors)\n\ + -b, --backend WHICH Which backend(s) to run (default: both on macOS, cpu elsewhere)\n\ + -v, --verify PATH Reference safetensors with `input`/`output`; checks numerical parity\n\ + -r, --requests N Run the async engine with N concurrent requests (default: 0)\n\ + Text generation: cargo run --bin generate -- --help\n\ + -h, --help Show this help\n\ +\n\ +Generate a demo model with: python python/make_demo_model.py", + env!("CARGO_PKG_VERSION") + ); +} - pub async fn run(mut self) { - info!("RequestManager: Listening for incoming requests..."); - - while let Some(req) = self.request_rx.recv().await { - let backend = Arc::clone(&self.backend); - let kv_cache = self.kv_storage.get_or_create(req.request_id); - - // For testing: dummy KV update and attention - let m = 1; // 1 token at a time for generation - let d = 64; - - let new_k_data = vec![0.1f32; (m * d) as usize]; - let new_v_data = vec![0.2f32; (m * d) as usize]; - let q_data = req.input_data; // Assume input is Q for this test - - let buf_new_k = backend.create_buffer(&new_k_data); - let buf_new_v = backend.create_buffer(&new_v_data); - let buf_q = backend.create_buffer(&q_data); - let buf_o = backend.create_buffer_uninitialized::((m * d) as usize); - - // 1. Update KV Cache - backend.update_kv_cache( - &buf_new_k, - &buf_new_v, - &kv_cache.k_buffer, - &kv_cache.v_buffer, - m as u32, - kv_cache.current_len as u32, - d as u32 - ); - kv_cache.current_len += m; - - // 2. Perform KV Attention - backend.kv_attention( - &buf_q, - &kv_cache.k_buffer, - &kv_cache.v_buffer, - &buf_o, - m as u32, - kv_cache.current_len as u32, - d as u32 - ); - - let ptr = buf_o.contents() as *const f32; - let mut o_data = vec![0.0f32; (m * d) as usize]; - unsafe { std::ptr::copy_nonoverlapping(ptr, o_data.as_mut_ptr(), o_data.len()); } - - let _ = req.response_tx.send(o_data); +/// Short summary of a tensor: shape, first values, and L2 norm. +fn summarize(t: &Tensor) -> String { + let l2 = t.data.iter().map(|v| v * v).sum::().sqrt(); + let head: Vec = t.data.iter().take(4).map(|v| format!("{v:+.4}")).collect(); + format!( + "shape {:?}, ‖·‖₂={l2:.4}, head=[{}, …]", + t.shape, + head.join(", ") + ) +} + +/// Builds the input: from the reference file if present, else a deterministic vector. +fn build_input(verify: &Option, in_features: usize) -> Result { + if let Some(path) = verify { + let map = loader::load_safetensors(path).map_err(|e| format!("load --verify: {e}"))?; + if let Some(input) = map.get("input") { + return Ok(input.clone()); + } + return Err("--verify file has no `input` tensor; reference verification requires the original input".into()); + } + let data = (0..in_features).map(|i| (i as f32 * 0.01).sin()).collect(); + Tensor::new(data, vec![1, in_features]).map_err(|e| e.to_string()) +} + +#[cfg(target_os = "macos")] +fn make_metal() -> Option> { + match batch_forge::metal_backend::MetalBackend::new(batch_forge::SHADER_SOURCE) { + Ok(b) => Some(Arc::new(b)), + Err(e) => { + warn!("Metal unavailable ({e}); falling back to CPU"); + None } } } #[tokio::main] -async fn main() { - tracing_subscriber::fmt::init(); - info!("Starting batch_forge Async Engine"); - - // Phase 1: Load SafeTensors (if available) - let model_path = PathBuf::from("model.safetensors"); - if model_path.exists() { - match loader::load_safetensors(&model_path) { - Ok(tensors) => { - info!("Successfully loaded {} tensors from Safetensors.", tensors.len()); +async fn main() -> ExitCode { + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into()), + ) + .init(); + + let args = match parse_args() { + Ok(a) => a, + Err(e) => { + if e != "help" { + eprintln!("error: {e}\n"); + } + print_help(); + return if e == "help" { + ExitCode::SUCCESS + } else { + ExitCode::FAILURE + }; + } + }; + + #[cfg(not(target_os = "macos"))] + if args.backend == BackendChoice::Metal { + eprintln!("error: Metal requires macOS; use --backend cpu"); + return ExitCode::FAILURE; + } + + info!("batch_forge {} starting", env!("CARGO_PKG_VERSION")); + + if !args.model.exists() { + warn!("model not found at {}", args.model.display()); + println!( + "\nNo checkpoint at `{}`. Generate the demo model first:\n python python/make_demo_model.py\n\ +or export your own Equinox model:\n python python/export_eqx.py --out model.safetensors", + args.model.display() + ); + return ExitCode::FAILURE; + } + + let model = match load_model(&args.model) { + Ok(m) => Arc::new(m), + Err(e) => { + error!("{e}"); + return ExitCode::FAILURE; + } + }; + info!( + "loaded MLP: {} layers, in={}, out={}", + model.layers.len(), + model.in_features(), + model.out_features() + ); + + let input = match build_input(&args.verify, model.in_features()) { + Ok(t) => t, + Err(e) => { + error!("{e}"); + return ExitCode::FAILURE; + } + }; + + // --- CPU reference forward (always available) --- + let cpu_out = match model.forward(&CpuBackend, &input) { + Ok(output) => output, + Err(error) => { + error!("CPU forward failed: {error}"); + return ExitCode::FAILURE; + } + }; + let mut exit = ExitCode::SUCCESS; + if args.backend != BackendChoice::Metal { + info!("[cpu] output: {}", summarize(&cpu_out)); + } + + // --- Metal forward + cross-backend check --- + // `mut` is only used on macOS, where the Metal result replaces the CPU one. + #[cfg_attr(not(target_os = "macos"), allow(unused_mut))] + let mut production = cpu_out.clone(); + #[cfg(target_os = "macos")] + if args.backend != BackendChoice::Cpu { + if let Some(metal) = make_metal() { + let metal_out = match model.forward(metal.as_ref(), &input) { + Ok(output) => output, + Err(error) => { + error!("Metal forward failed: {error}"); + return ExitCode::FAILURE; + } + }; + info!("[metal] output: {}", summarize(&metal_out)); + let diff = cpu_out.max_abs_diff(&metal_out); + info!("[check] CPU vs Metal max|Δ| = {diff:.3e}"); + if !diff.is_finite() || diff > VERIFY_TOL { + error!("CPU/Metal parity failed (tolerance {VERIFY_TOL})"); + exit = ExitCode::FAILURE; + } + production = metal_out; + } else if args.backend == BackendChoice::Metal { + error!("requested Metal backend is unavailable"); + return ExitCode::FAILURE; + } + } + + // --- Reference verification --- + if let Some(path) = &args.verify { + match verify_against_reference(path, &production) { + Ok(diff) => { + if diff <= VERIFY_TOL { + info!( + "[verify] PASS — max|Δ| vs reference = {diff:.3e} (tol {VERIFY_TOL:.0e})" + ); + } else { + error!( + "[verify] FAIL — max|Δ| vs reference = {diff:.3e} (tol {VERIFY_TOL:.0e})" + ); + exit = ExitCode::FAILURE; + } } Err(e) => { - error!("Failed to load model: {}", e); + error!("[verify] {e}"); + exit = ExitCode::FAILURE; } } - } else { - info!("No model.safetensors found. Run python export_eqx.py to generate it."); } - // Initialize backend - let shader_source = include_str!("shaders/compute.metal"); - let backend = Arc::new(metal_backend::MetalBackend::new(shader_source).expect("Failed to init Metal")); - - // Set up request queue - let (tx, rx) = mpsc::channel(100); - let manager = RequestManager::new(Arc::clone(&backend), rx); - - // Spawn the manager in its own task - tokio::spawn(async move { - manager.run().await; - }); - - // Simulate an autoregressive generation loop for a single request - let tx_clone = tx.clone(); - let handle = tokio::spawn(async move { - let request_id = 42; - info!("Starting Autoregressive Generation for Request {}", request_id); - - for step in 0..5 { - let (resp_tx, resp_rx) = oneshot::channel(); - // In a real model, Q would be derived from the previous step's output - let input_q = vec![0.5f32; 64]; - - tx_clone.send(InferenceRequest { - request_id, - input_data: input_q, - response_tx: resp_tx, - }).await.unwrap(); - - let result = resp_rx.await.unwrap(); - info!("Step {}: Generation output (first 5 elements): {:?}", step, &result[0..5]); + // --- Async engine demo --- + if args.requests > 0 + && !run_async_demo(Arc::clone(&model), &input, args.requests, args.backend).await + { + exit = ExitCode::FAILURE; + } + + exit +} + +fn load_model(path: &Path) -> Result { + let tensors = loader::load_safetensors(path).map_err(|e| format!("load model: {e}"))?; + Mlp::from_tensors(&tensors).map_err(|e| format!("build model: {e}")) +} + +fn verify_against_reference(path: &Path, produced: &Tensor) -> Result { + let map = loader::load_safetensors(path).map_err(|e| format!("load reference: {e}"))?; + let reference = map + .get("output") + .ok_or("reference file has no `output` tensor")?; + if reference.shape != produced.shape { + return Err(format!( + "shape mismatch: produced {:?}, reference {:?}", + produced.shape, reference.shape + )); + } + Ok(produced.max_abs_diff(reference)) +} + +async fn run_async_demo( + model: Arc, + input: &Tensor, + requests: usize, + choice: BackendChoice, +) -> bool { + let backend: Arc = pick_async_backend(choice); + info!( + "[engine] dispatching {requests} concurrent requests on `{}`", + backend.name() + ); + let submitter = batch_forge::engine::spawn(backend, model, requests.max(1)); + + let start = std::time::Instant::now(); + let mut handles = Vec::with_capacity(requests); + for id in 0..requests as u64 { + let s = submitter.clone(); + let inp = input.clone(); + handles.push(tokio::spawn(async move { s.infer(id, inp).await })); + } + let mut ok = 0usize; + for h in handles { + if matches!(h.await, Ok(Ok(_))) { + ok += 1; } - }); + } + drop(submitter); + let elapsed = start.elapsed(); + let rps = ok as f64 / elapsed.as_secs_f64(); + info!( + "[engine] {ok}/{requests} succeeded in {:.2?} ({rps:.0} req/s)", + elapsed + ); + ok == requests +} - handle.await.unwrap(); +fn pick_async_backend(choice: BackendChoice) -> Arc { + #[cfg(target_os = "macos")] + if choice != BackendChoice::Cpu { + if let Some(metal) = make_metal() { + return metal; + } + } + let _ = choice; + Arc::new(CpuBackend) } diff --git a/src/metal_backend.rs b/src/metal_backend.rs index c40366c..db62b03 100644 --- a/src/metal_backend.rs +++ b/src/metal_backend.rs @@ -1,175 +1,523 @@ -use metal::{Buffer, CommandQueue, CompileOptions, ComputePipelineState, Device, Library, MTLResourceOptions, MTLSize}; -use std::error::Error; -use tracing::{info}; +//! Metal compute backend. +//! +//! Each method mirrors a function in [`crate::ops`] and is validated against it +//! by `tests/parity.rs`. Buffers use Metal's unified (shared) memory, so on +//! Apple Silicon there is no discrete host↔device copy — `create_buffer` and +//! `read_buffer` read/write the same physical pages the GPU sees. + +use bytemuck::Pod; +use metal::{ + Buffer, CommandQueue, ComputeCommandEncoderRef, ComputePipelineState, Device, Library, + MTLResourceOptions, MTLSize, +}; +use thiserror::Error; +use tracing::info; + +use crate::model::Backend; +use crate::tensor::Tensor; + +#[derive(Error, Debug)] +pub enum BackendError { + #[error("Buffer size computation overflowed")] + BufferOverflow, + #[error("Initialization error: {0}")] + Init(String), +} pub struct MetalBackend { pub device: Device, command_queue: CommandQueue, + #[allow(dead_code)] // kept alive so pipeline states remain valid library: Library, matmul_pipeline: ComputePipelineState, + linear_pipeline: ComputePipelineState, quant_matmul_pipeline: ComputePipelineState, + gelu_pipeline: ComputePipelineState, kv_attention_pipeline: ComputePipelineState, update_kv_cache_pipeline: ComputePipelineState, + layernorm_pipeline: ComputePipelineState, + rmsnorm_pipeline: ComputePipelineState, + rope_pipeline: ComputePipelineState, + matmul_tiled_pipeline: ComputePipelineState, + mha_pipeline: ComputePipelineState, } impl MetalBackend { - pub fn new(shader_source: &str) -> Result> { - let device = Device::system_default().ok_or("No Metal device found. Are you on a Mac?")?; + pub fn new(shader_source: &str) -> Result { + let device = Device::system_default() + .ok_or_else(|| BackendError::Init("No Metal device found. Are you on a Mac?".into()))?; info!("Initialized Metal device: {}", device.name()); let command_queue = device.new_command_queue(); - - let options = CompileOptions::new(); - let library = device.new_library_with_source(shader_source, &options) - .map_err(|e| format!("Failed to compile shader: {}", e))?; - - let matmul_func = library.get_function("matmul", None) - .map_err(|e| format!("Failed to find function 'matmul': {}", e))?; - let matmul_pipeline = device.new_compute_pipeline_state_with_function(&matmul_func) - .map_err(|e| format!("Failed to create compute pipeline: {}", e))?; - - let quant_matmul_func = library.get_function("quant_matmul", None) - .map_err(|e| format!("Failed to find function 'quant_matmul': {}", e))?; - let quant_matmul_pipeline = device.new_compute_pipeline_state_with_function(&quant_matmul_func) - .map_err(|e| format!("Failed to create compute pipeline: {}", e))?; - - let kv_attention_func = library.get_function("kv_attention", None) - .map_err(|e| format!("Failed to find function 'kv_attention': {}", e))?; - let kv_attention_pipeline = device.new_compute_pipeline_state_with_function(&kv_attention_func) - .map_err(|e| format!("Failed to create compute pipeline: {}", e))?; - - let update_kv_cache_func = library.get_function("update_kv_cache", None) - .map_err(|e| format!("Failed to find function 'update_kv_cache': {}", e))?; - let update_kv_cache_pipeline = device.new_compute_pipeline_state_with_function(&update_kv_cache_func) - .map_err(|e| format!("Failed to create compute pipeline: {}", e))?; + let library = device + .new_library_with_source(shader_source, &metal::CompileOptions::new()) + .map_err(|e| BackendError::Init(format!("Failed to compile shaders: {e}")))?; + + let pso = |name: &str| -> Result { + let func = library + .get_function(name, None) + .map_err(|e| BackendError::Init(format!("missing kernel '{name}': {e}")))?; + device + .new_compute_pipeline_state_with_function(&func) + .map_err(|e| BackendError::Init(format!("pipeline '{name}': {e}"))) + }; Ok(Self { - device, + matmul_pipeline: pso("matmul")?, + linear_pipeline: pso("linear")?, + quant_matmul_pipeline: pso("quant_matmul")?, + gelu_pipeline: pso("gelu_forward")?, + kv_attention_pipeline: pso("kv_attention")?, + update_kv_cache_pipeline: pso("update_kv_cache")?, + layernorm_pipeline: pso("layernorm")?, + rmsnorm_pipeline: pso("rmsnorm")?, + rope_pipeline: pso("rope")?, + matmul_tiled_pipeline: pso("matmul_tiled")?, + mha_pipeline: pso("mha")?, command_queue, library, - matmul_pipeline, - quant_matmul_pipeline, - kv_attention_pipeline, - update_kv_cache_pipeline, + device, }) } - pub fn create_buffer(&self, data: &[T]) -> Buffer { - let length = (data.len() * std::mem::size_of::()) as u64; - self.device.new_buffer_with_data( + // --- buffer helpers ---------------------------------------------------- + + pub fn create_buffer(&self, data: &[T]) -> Result { + let length = std::mem::size_of_val(data); + Ok(self.device.new_buffer_with_data( data.as_ptr() as *const _, - length, + length as u64, MTLResourceOptions::StorageModeShared, - ) + )) } - pub fn create_buffer_uninitialized(&self, len: usize) -> Buffer { - let length = (len * std::mem::size_of::()) as u64; - self.device.new_buffer( - length, - MTLResourceOptions::StorageModeShared, - ) + pub fn create_buffer_uninitialized(&self, len: usize) -> Result { + let length = len + .checked_mul(std::mem::size_of::()) + .ok_or(BackendError::BufferOverflow)?; + Ok(self + .device + .new_buffer(length.max(1) as u64, MTLResourceOptions::StorageModeShared)) } - pub fn matmul(&self, a: &Buffer, b: &Buffer, c: &Buffer, m: u32, n: u32, k: u32) { - let command_buffer = self.command_queue.new_command_buffer(); - let encoder = command_buffer.new_compute_command_encoder(); - - encoder.set_compute_pipeline_state(&self.matmul_pipeline); - encoder.set_buffer(0, Some(a), 0); - encoder.set_buffer(1, Some(b), 0); - encoder.set_buffer(2, Some(c), 0); - - let m_buf = self.create_buffer(&[m]); - let n_buf = self.create_buffer(&[n]); - let k_buf = self.create_buffer(&[k]); - - encoder.set_buffer(3, Some(&m_buf), 0); - encoder.set_buffer(4, Some(&n_buf), 0); - encoder.set_buffer(5, Some(&k_buf), 0); - - let w = self.matmul_pipeline.thread_execution_width(); - let h = self.matmul_pipeline.max_total_threads_per_threadgroup() / w; - - let threads_per_threadgroup = MTLSize::new(w, h, 1); - let threadgroups_per_grid = MTLSize::new( - (n as u64 + w - 1) / w, - (m as u64 + h - 1) / h, - 1, + /// Copies `len` elements of type `T` out of a shared buffer. + pub fn read_buffer(&self, buffer: &Buffer, len: usize) -> Vec { + let mut out = vec![T::zeroed(); len]; + let ptr = buffer.contents() as *const T; + // SAFETY: shared-storage buffer, `len` elements were allocated/written. + unsafe { std::ptr::copy_nonoverlapping(ptr, out.as_mut_ptr(), len) }; + out + } + + // --- dispatch helpers -------------------------------------------------- + + fn run_1d( + &self, + pso: &ComputePipelineState, + set: impl FnOnce(&ComputeCommandEncoderRef), + n: u64, + ) { + let cb = self.command_queue.new_command_buffer(); + let enc = cb.new_compute_command_encoder(); + enc.set_compute_pipeline_state(pso); + set(enc); + let max = pso.max_total_threads_per_threadgroup(); + let tgw = n.clamp(1, max); + enc.dispatch_thread_groups(MTLSize::new(n.div_ceil(tgw), 1, 1), MTLSize::new(tgw, 1, 1)); + enc.end_encoding(); + cb.commit(); + cb.wait_until_completed(); + } + + fn run_2d( + &self, + pso: &ComputePipelineState, + set: impl FnOnce(&ComputeCommandEncoderRef), + gx: u64, + gy: u64, + ) { + let cb = self.command_queue.new_command_buffer(); + let enc = cb.new_compute_command_encoder(); + enc.set_compute_pipeline_state(pso); + set(enc); + let w = pso.thread_execution_width(); + let h = (pso.max_total_threads_per_threadgroup() / w).max(1); + enc.dispatch_thread_groups( + MTLSize::new(gx.div_ceil(w), gy.div_ceil(h), 1), + MTLSize::new(w, h, 1), + ); + enc.end_encoding(); + cb.commit(); + cb.wait_until_completed(); + } + + // --- high-level ops (Vec in / Vec out, used by parity tests) ----------- + + pub fn matmul(&self, a: &[f32], b: &[f32], m: usize, k: usize, n: usize) -> Vec { + let ba = self.create_buffer(a).unwrap(); + let bb = self.create_buffer(b).unwrap(); + let bc = self.create_buffer_uninitialized::(m * n).unwrap(); + let (bm, bn, bk) = self.dims3(m, n, k); + self.run_2d( + &self.matmul_pipeline, + |e| { + e.set_buffer(0, Some(&ba), 0); + e.set_buffer(1, Some(&bb), 0); + e.set_buffer(2, Some(&bc), 0); + e.set_buffer(3, Some(&bm), 0); + e.set_buffer(4, Some(&bn), 0); + e.set_buffer(5, Some(&bk), 0); + }, + n as u64, + m as u64, + ); + self.read_buffer(&bc, m * n) + } + + /// Shared-memory tiled matmul: `C[m,n] = A[m,k] * B[k,n]`. Same result as + /// [`MetalBackend::matmul`] but uses full 16x16 threadgroups with on-chip tiles. + pub fn matmul_tiled(&self, a: &[f32], b: &[f32], m: usize, k: usize, n: usize) -> Vec { + const TILE: u64 = 16; + let ba = self.create_buffer(a).unwrap(); + let bb = self.create_buffer(b).unwrap(); + let bc = self.create_buffer_uninitialized::(m * n).unwrap(); + let (bm, bn, bk) = self.dims3(m, n, k); + let cb = self.command_queue.new_command_buffer(); + let enc = cb.new_compute_command_encoder(); + enc.set_compute_pipeline_state(&self.matmul_tiled_pipeline); + enc.set_buffer(0, Some(&ba), 0); + enc.set_buffer(1, Some(&bb), 0); + enc.set_buffer(2, Some(&bc), 0); + enc.set_buffer(3, Some(&bm), 0); + enc.set_buffer(4, Some(&bn), 0); + enc.set_buffer(5, Some(&bk), 0); + // Full TILExTILE threadgroups so boundary threads still load zeros and + // participate in the barriers. + enc.dispatch_thread_groups( + MTLSize::new((n as u64).div_ceil(TILE), (m as u64).div_ceil(TILE), 1), + MTLSize::new(TILE, TILE, 1), ); + enc.end_encoding(); + cb.commit(); + cb.wait_until_completed(); + self.read_buffer(&bc, m * n) + } - encoder.dispatch_thread_groups(threadgroups_per_grid, threads_per_threadgroup); - encoder.end_encoding(); - - command_buffer.commit(); - command_buffer.wait_until_completed(); - } - - pub fn kv_attention(&self, q: &Buffer, k_cache: &Buffer, v_cache: &Buffer, o: &Buffer, m: u32, cur_seq_len: u32, d: u32) { - let command_buffer = self.command_queue.new_command_buffer(); - let encoder = command_buffer.new_compute_command_encoder(); - - encoder.set_compute_pipeline_state(&self.kv_attention_pipeline); - encoder.set_buffer(0, Some(q), 0); - encoder.set_buffer(1, Some(k_cache), 0); - encoder.set_buffer(2, Some(v_cache), 0); - encoder.set_buffer(3, Some(o), 0); - - let m_buf = self.create_buffer(&[m]); - let cur_seq_len_buf = self.create_buffer(&[cur_seq_len]); - let d_buf = self.create_buffer(&[d]); - - encoder.set_buffer(4, Some(&m_buf), 0); - encoder.set_buffer(5, Some(&cur_seq_len_buf), 0); - encoder.set_buffer(6, Some(&d_buf), 0); - - let max_threads = self.kv_attention_pipeline.max_total_threads_per_threadgroup(); - let threads_per_threadgroup = MTLSize::new(std::cmp::min(m as u64, max_threads), 1, 1); - let threadgroups_per_grid = MTLSize::new( - (m as u64 + threads_per_threadgroup.width - 1) / threads_per_threadgroup.width, - 1, - 1, + /// Multi-head causal self-attention. Q/K/V are `[seq, heads*head_dim]`. + pub fn mha( + &self, + q: &[f32], + k: &[f32], + v: &[f32], + seq: usize, + heads: usize, + head_dim: usize, + ) -> Vec { + let bq = self.create_buffer(q).unwrap(); + let bk = self.create_buffer(k).unwrap(); + let bv = self.create_buffer(v).unwrap(); + let bo = self + .create_buffer_uninitialized::(seq * heads * head_dim) + .unwrap(); + let (bs, bh, bd) = self.dims3(seq, heads, head_dim); + let cb = self.command_queue.new_command_buffer(); + let enc = cb.new_compute_command_encoder(); + enc.set_compute_pipeline_state(&self.mha_pipeline); + enc.set_buffer(0, Some(&bq), 0); + enc.set_buffer(1, Some(&bk), 0); + enc.set_buffer(2, Some(&bv), 0); + enc.set_buffer(3, Some(&bo), 0); + enc.set_buffer(4, Some(&bs), 0); + enc.set_buffer(5, Some(&bh), 0); + enc.set_buffer(6, Some(&bd), 0); + // One thread per (query, head); non-uniform dispatch handles the edges. + let th = heads.min(1024) as u64; + let tw = (1024 / th).min(seq as u64).max(1); + enc.dispatch_threads( + MTLSize::new(seq as u64, heads as u64, 1), + MTLSize::new(tw, th, 1), ); + enc.end_encoding(); + cb.commit(); + cb.wait_until_completed(); + self.read_buffer(&bo, seq * heads * head_dim) + } - encoder.dispatch_thread_groups(threadgroups_per_grid, threads_per_threadgroup); - encoder.end_encoding(); - - command_buffer.commit(); - command_buffer.wait_until_completed(); - } - - pub fn update_kv_cache(&self, new_k: &Buffer, new_v: &Buffer, k_cache: &Buffer, v_cache: &Buffer, m: u32, offset: u32, d: u32) { - let command_buffer = self.command_queue.new_command_buffer(); - let encoder = command_buffer.new_compute_command_encoder(); - - encoder.set_compute_pipeline_state(&self.update_kv_cache_pipeline); - encoder.set_buffer(0, Some(new_k), 0); - encoder.set_buffer(1, Some(new_v), 0); - encoder.set_buffer(2, Some(k_cache), 0); - encoder.set_buffer(3, Some(v_cache), 0); - - let m_buf = self.create_buffer(&[m]); - let offset_buf = self.create_buffer(&[offset]); - let d_buf = self.create_buffer(&[d]); - - encoder.set_buffer(4, Some(&m_buf), 0); - encoder.set_buffer(5, Some(&offset_buf), 0); - encoder.set_buffer(6, Some(&d_buf), 0); - - let w = self.update_kv_cache_pipeline.thread_execution_width(); - let h = self.update_kv_cache_pipeline.max_total_threads_per_threadgroup() / w; - - let threads_per_threadgroup = MTLSize::new(w, h, 1); - let threadgroups_per_grid = MTLSize::new( - (m as u64 + w - 1) / w, - (d as u64 + h - 1) / h, - 1, + pub fn linear( + &self, + x: &[f32], + w: &[f32], + bias: &[f32], + rows: usize, + in_f: usize, + out_f: usize, + ) -> Vec { + let bx = self.create_buffer(x).unwrap(); + let bw = self.create_buffer(w).unwrap(); + let bb = self.create_buffer(bias).unwrap(); + let by = self + .create_buffer_uninitialized::(rows * out_f) + .unwrap(); + let (br, bi, bo) = self.dims3(rows, in_f, out_f); + self.run_2d( + &self.linear_pipeline, + |e| { + e.set_buffer(0, Some(&bx), 0); + e.set_buffer(1, Some(&bw), 0); + e.set_buffer(2, Some(&bb), 0); + e.set_buffer(3, Some(&by), 0); + e.set_buffer(4, Some(&br), 0); + e.set_buffer(5, Some(&bi), 0); + e.set_buffer(6, Some(&bo), 0); + }, + out_f as u64, + rows as u64, ); + self.read_buffer(&by, rows * out_f) + } + + pub fn quant_matmul( + &self, + a_int8: &[i8], + scales: &[f32], + b: &[f32], + m: usize, + k: usize, + n: usize, + ) -> Vec { + let ba = self.create_buffer(a_int8).unwrap(); + let bb = self.create_buffer(b).unwrap(); + let bc = self.create_buffer_uninitialized::(m * n).unwrap(); + let bs = self.create_buffer(scales).unwrap(); + let (bm, bn, bk) = self.dims3(m, n, k); + self.run_2d( + &self.quant_matmul_pipeline, + |e| { + e.set_buffer(0, Some(&ba), 0); + e.set_buffer(1, Some(&bb), 0); + e.set_buffer(2, Some(&bc), 0); + e.set_buffer(3, Some(&bs), 0); + e.set_buffer(4, Some(&bm), 0); + e.set_buffer(5, Some(&bn), 0); + e.set_buffer(6, Some(&bk), 0); + }, + n as u64, + m as u64, + ); + self.read_buffer(&bc, m * n) + } - encoder.dispatch_thread_groups(threadgroups_per_grid, threads_per_threadgroup); - encoder.end_encoding(); + pub fn gelu(&self, x: &[f32]) -> Vec { + let bx = self.create_buffer(x).unwrap(); + let by = self.create_buffer_uninitialized::(x.len()).unwrap(); + let bn = self.create_buffer(&[x.len() as u32]).unwrap(); + self.run_1d( + &self.gelu_pipeline, + |e| { + e.set_buffer(0, Some(&bx), 0); + e.set_buffer(1, Some(&by), 0); + e.set_buffer(2, Some(&bn), 0); + }, + x.len() as u64, + ); + self.read_buffer(&by, x.len()) + } + + pub fn layernorm( + &self, + x: &[f32], + gamma: &[f32], + beta: &[f32], + rows: usize, + d: usize, + eps: f32, + ) -> Vec { + let bx = self.create_buffer(x).unwrap(); + let bg = self.create_buffer(gamma).unwrap(); + let bb = self.create_buffer(beta).unwrap(); + let by = self.create_buffer_uninitialized::(rows * d).unwrap(); + let br = self.create_buffer(&[rows as u32]).unwrap(); + let bd = self.create_buffer(&[d as u32]).unwrap(); + let be = self.create_buffer(&[eps]).unwrap(); + self.run_1d( + &self.layernorm_pipeline, + |e| { + e.set_buffer(0, Some(&bx), 0); + e.set_buffer(1, Some(&bg), 0); + e.set_buffer(2, Some(&bb), 0); + e.set_buffer(3, Some(&by), 0); + e.set_buffer(4, Some(&br), 0); + e.set_buffer(5, Some(&bd), 0); + e.set_buffer(6, Some(&be), 0); + }, + rows as u64, + ); + self.read_buffer(&by, rows * d) + } + + pub fn rmsnorm(&self, x: &[f32], gamma: &[f32], rows: usize, d: usize, eps: f32) -> Vec { + let bx = self.create_buffer(x).unwrap(); + let bg = self.create_buffer(gamma).unwrap(); + let by = self.create_buffer_uninitialized::(rows * d).unwrap(); + let br = self.create_buffer(&[rows as u32]).unwrap(); + let bd = self.create_buffer(&[d as u32]).unwrap(); + let be = self.create_buffer(&[eps]).unwrap(); + self.run_1d( + &self.rmsnorm_pipeline, + |e| { + e.set_buffer(0, Some(&bx), 0); + e.set_buffer(1, Some(&bg), 0); + e.set_buffer(2, Some(&by), 0); + e.set_buffer(3, Some(&br), 0); + e.set_buffer(4, Some(&bd), 0); + e.set_buffer(5, Some(&be), 0); + }, + rows as u64, + ); + self.read_buffer(&by, rows * d) + } + + pub fn rope( + &self, + x: &[f32], + positions: &[u32], + rows: usize, + d: usize, + theta: f32, + ) -> Vec { + let bx = self.create_buffer(x).unwrap(); + let bp = self.create_buffer(positions).unwrap(); + let br = self.create_buffer(&[rows as u32]).unwrap(); + let bd = self.create_buffer(&[d as u32]).unwrap(); + let bt = self.create_buffer(&[theta]).unwrap(); + self.run_1d( + &self.rope_pipeline, + |e| { + e.set_buffer(0, Some(&bx), 0); + e.set_buffer(1, Some(&bp), 0); + e.set_buffer(2, Some(&br), 0); + e.set_buffer(3, Some(&bd), 0); + e.set_buffer(4, Some(&bt), 0); + }, + rows as u64, + ); + self.read_buffer(&bx, rows * d) + } + + /// Single-head attention over Vec inputs (used by parity tests). + #[allow(clippy::too_many_arguments)] + pub fn attention( + &self, + q: &[f32], + k: &[f32], + v: &[f32], + m: usize, + seq: usize, + d: usize, + causal: bool, + q_offset: usize, + ) -> Vec { + let bq = self.create_buffer(q).unwrap(); + let bk = self.create_buffer(k).unwrap(); + let bv = self.create_buffer(v).unwrap(); + let bo = self.create_buffer_uninitialized::(m * d).unwrap(); + self.kv_attention(&bq, &bk, &bv, &bo, m, seq, d, causal, q_offset); + self.read_buffer(&bo, m * d) + } + + // --- buffer-level ops (used by the cached generation loop) ------------- + + /// Attention reading K/V straight from persistent cache buffers. + #[allow(clippy::too_many_arguments)] + pub fn kv_attention( + &self, + q: &Buffer, + k_cache: &Buffer, + v_cache: &Buffer, + o: &Buffer, + m: usize, + cur_seq_len: usize, + d: usize, + causal: bool, + q_offset: usize, + ) { + let bm = self.create_buffer(&[m as u32]).unwrap(); + let bs = self.create_buffer(&[cur_seq_len as u32]).unwrap(); + let bd = self.create_buffer(&[d as u32]).unwrap(); + let bc = self.create_buffer(&[causal as u32]).unwrap(); + let bo = self.create_buffer(&[q_offset as u32]).unwrap(); + self.run_1d( + &self.kv_attention_pipeline, + |e| { + e.set_buffer(0, Some(q), 0); + e.set_buffer(1, Some(k_cache), 0); + e.set_buffer(2, Some(v_cache), 0); + e.set_buffer(3, Some(o), 0); + e.set_buffer(4, Some(&bm), 0); + e.set_buffer(5, Some(&bs), 0); + e.set_buffer(6, Some(&bd), 0); + e.set_buffer(7, Some(&bc), 0); + e.set_buffer(8, Some(&bo), 0); + }, + m as u64, + ); + } + + /// Writes `m` new K/V rows into the cache buffers at `offset`. + #[allow(clippy::too_many_arguments)] + pub fn update_kv_cache( + &self, + new_k: &Buffer, + new_v: &Buffer, + k_cache: &Buffer, + v_cache: &Buffer, + m: usize, + offset: usize, + d: usize, + ) { + let bm = self.create_buffer(&[m as u32]).unwrap(); + let boff = self.create_buffer(&[offset as u32]).unwrap(); + let bd = self.create_buffer(&[d as u32]).unwrap(); + self.run_2d( + &self.update_kv_cache_pipeline, + |e| { + e.set_buffer(0, Some(new_k), 0); + e.set_buffer(1, Some(new_v), 0); + e.set_buffer(2, Some(k_cache), 0); + e.set_buffer(3, Some(v_cache), 0); + e.set_buffer(4, Some(&bm), 0); + e.set_buffer(5, Some(&boff), 0); + e.set_buffer(6, Some(&bd), 0); + }, + m as u64, + d as u64, + ); + } + + fn dims3(&self, a: usize, b: usize, c: usize) -> (Buffer, Buffer, Buffer) { + ( + self.create_buffer(&[a as u32]).unwrap(), + self.create_buffer(&[b as u32]).unwrap(), + self.create_buffer(&[c as u32]).unwrap(), + ) + } +} + +impl Backend for MetalBackend { + fn linear(&self, x: &Tensor, w: &Tensor, b: &Tensor) -> Tensor { + let (n, in_f) = x.dims2().expect("linear input must be 2-D"); + let (out_f, in_w) = w.dims2().expect("weight must be 2-D"); + assert_eq!(in_f, in_w, "linear in-features mismatch"); + let data = MetalBackend::linear(self, &x.data, &w.data, &b.data, n, in_f, out_f); + Tensor::new(data, vec![n, out_f]).expect("linear output shape is consistent") + } + + fn gelu(&self, x: &Tensor) -> Tensor { + let data = MetalBackend::gelu(self, &x.data); + Tensor::new(data, x.shape.clone()).expect("gelu preserves shape") + } - command_buffer.commit(); - command_buffer.wait_until_completed(); + fn name(&self) -> &'static str { + "metal" } } diff --git a/src/model.rs b/src/model.rs new file mode 100644 index 0000000..de24a4a --- /dev/null +++ b/src/model.rs @@ -0,0 +1,327 @@ +//! Model definitions and the backend abstraction used to run them. +//! +//! [`Mlp`] mirrors `python/export_eqx.py::SimpleMLP`. It is generic over a +//! [`Backend`], so the exact same forward pass runs on the portable +//! [`CpuBackend`] reference and on the Metal backend, which is how the +//! end-to-end test asserts CPU/GPU agreement. + +use std::collections::{BTreeSet, HashMap}; +use thiserror::Error; + +use crate::ops; +use crate::tensor::{Tensor, TensorError}; + +#[derive(Error, Debug)] +pub enum ModelError { + #[error("inference engine closed before returning a result")] + EngineClosed, + #[error("invalid tensor: {0}")] + Tensor(#[from] TensorError), + #[error("missing tensor '{0}' in checkpoint")] + MissingTensor(String), + #[error("tensor '{name}' has unexpected shape {shape:?}")] + BadShape { name: String, shape: Vec }, + #[error("no layers found in checkpoint")] + NoLayers, + #[error("non-contiguous layer indices: expected layers.{expected}, found layers.{found}")] + NonContiguousLayers { expected: usize, found: usize }, + #[error("dimension mismatch: layer expects input width {expected}, got {got}")] + WidthMismatch { expected: usize, got: usize }, +} + +/// The minimal compute surface a feed-forward model needs from a backend. +/// +/// Keeping this trait tiny is deliberate: a new backend (CPU, Metal, and later +/// Vulkan/WebGPU) only has to implement these two ops to run the whole model. +pub trait Backend { + /// `y = x · Wᵀ + b`, with `x`[n×in], `w`[out×in], `b`[out] → `[n×out]`. + fn linear(&self, x: &Tensor, w: &Tensor, b: &Tensor) -> Tensor; + /// Element-wise GELU (tanh approximation). + fn gelu(&self, x: &Tensor) -> Tensor; + /// Human-readable backend name, used in logs and the demo output. + fn name(&self) -> &'static str; +} + +/// Portable, always-available reference backend. Defines correct behavior. +pub struct CpuBackend; + +impl Backend for CpuBackend { + fn linear(&self, x: &Tensor, w: &Tensor, b: &Tensor) -> Tensor { + let (n, in_f) = x.dims2().expect("linear input must be 2-D"); + let (out_f, in_w) = w.dims2().expect("weight must be 2-D"); + assert_eq!(in_f, in_w, "linear in-features mismatch"); + let data = ops::linear(&x.data, &w.data, &b.data, n, in_f, out_f); + Tensor::new(data, vec![n, out_f]).expect("linear output shape is consistent") + } + + fn gelu(&self, x: &Tensor) -> Tensor { + let mut out = x.clone(); + ops::gelu_inplace(&mut out.data); + out + } + + fn name(&self) -> &'static str { + "cpu" + } +} + +/// A stack of `Linear` layers with GELU between them (no activation after the +/// final layer) — the structure exported by `SimpleMLP`. +pub struct Mlp { + /// `(weight[out×in], bias[out])` per layer, in execution order. + pub layers: Vec<(Tensor, Tensor)>, +} + +impl Mlp { + /// Builds an MLP from a name→tensor map using the `layers.{i}.weight` / + /// `layers.{i}.bias` naming emitted by the exporter. + pub fn from_tensors(map: &HashMap) -> Result { + let layer_indices: BTreeSet = map + .keys() + .filter_map(|name| { + name.strip_prefix("layers.")? + .strip_suffix(".weight")? + .parse() + .ok() + }) + .collect(); + for (expected, &found) in layer_indices.iter().enumerate() { + if expected != found { + return Err(ModelError::NonContiguousLayers { expected, found }); + } + } + + let mut layers = Vec::new(); + let mut i = 0; + loop { + let w_name = format!("layers.{i}.weight"); + let b_name = format!("layers.{i}.bias"); + let Some(w) = map.get(&w_name) else { break }; + let b = map + .get(&b_name) + .ok_or_else(|| ModelError::MissingTensor(b_name.clone()))?; + w.validate()?; + b.validate()?; + if w.shape.len() != 2 || w.shape.contains(&0) { + return Err(ModelError::BadShape { + name: w_name, + shape: w.shape.clone(), + }); + } + if b.shape.len() != 1 || b.shape[0] != w.shape[0] { + return Err(ModelError::BadShape { + name: b_name, + shape: b.shape.clone(), + }); + } + if let Some((previous, _)) = layers.last() { + let previous: &Tensor = previous; + if previous.shape[0] != w.shape[1] { + return Err(ModelError::WidthMismatch { + expected: previous.shape[0], + got: w.shape[1], + }); + } + } + layers.push((w.clone(), b.clone())); + i += 1; + } + if layers.is_empty() { + return Err(ModelError::NoLayers); + } + Ok(Self { layers }) + } + + /// Revalidates public layers before inference. + pub fn validate(&self) -> Result<(), ModelError> { + if self.layers.is_empty() { + return Err(ModelError::NoLayers); + } + let mut previous_width = None; + for (index, (weight, bias)) in self.layers.iter().enumerate() { + weight.validate()?; + bias.validate()?; + if weight.shape.len() != 2 || weight.shape.contains(&0) { + return Err(ModelError::BadShape { + name: format!("layers.{index}.weight"), + shape: weight.shape.clone(), + }); + } + if bias.shape.as_slice() != [weight.shape[0]] { + return Err(ModelError::BadShape { + name: format!("layers.{index}.bias"), + shape: bias.shape.clone(), + }); + } + if let Some(expected) = previous_width { + if weight.shape[1] != expected { + return Err(ModelError::WidthMismatch { + expected, + got: weight.shape[1], + }); + } + } + previous_width = Some(weight.shape[0]); + } + Ok(()) + } + + /// Input feature width expected by the first layer. + pub fn in_features(&self) -> usize { + self.layers[0].0.shape[1] + } + + /// Output feature width produced by the last layer. + pub fn out_features(&self) -> usize { + self.layers.last().unwrap().0.shape[0] + } + + /// Runs the forward pass on the given backend. `x` is `[n × in_features]`. + /// + /// `?Sized` so it accepts both concrete backends and `&dyn Backend` (used by + /// the async engine, which holds an `Arc`). + pub fn forward( + &self, + backend: &B, + x: &Tensor, + ) -> Result { + self.validate()?; + x.validate()?; + let (_, in_f) = x.dims2().map_err(|_| ModelError::WidthMismatch { + expected: self.in_features(), + got: 0, + })?; + if in_f != self.in_features() { + return Err(ModelError::WidthMismatch { + expected: self.in_features(), + got: in_f, + }); + } + let mut h = x.clone(); + let last = self.layers.len() - 1; + for (idx, (w, b)) in self.layers.iter().enumerate() { + h = backend.linear(&h, w, b); + if idx != last { + h = backend.gelu(&h); + } + } + Ok(h) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn tiny_model() -> Mlp { + // 2 -> 2 -> 2, identity-ish weights. + let mut map = HashMap::new(); + map.insert( + "layers.0.weight".into(), + Tensor::new(vec![1.0, 0.0, 0.0, 1.0], vec![2, 2]).unwrap(), + ); + map.insert( + "layers.0.bias".into(), + Tensor::new(vec![0.0, 0.0], vec![2]).unwrap(), + ); + map.insert( + "layers.1.weight".into(), + Tensor::new(vec![2.0, 0.0, 0.0, 2.0], vec![2, 2]).unwrap(), + ); + map.insert( + "layers.1.bias".into(), + Tensor::new(vec![1.0, 1.0], vec![2]).unwrap(), + ); + Mlp::from_tensors(&map).unwrap() + } + + #[test] + fn builds_layers_in_order() { + let m = tiny_model(); + assert_eq!(m.layers.len(), 2); + assert_eq!(m.in_features(), 2); + assert_eq!(m.out_features(), 2); + } + + #[test] + fn cpu_forward_is_correct() { + let m = tiny_model(); + let x = Tensor::new(vec![1.0, -1.0], vec![1, 2]).unwrap(); + // layer0 (identity) -> [1,-1]; gelu -> [gelu(1), gelu(-1)]; + // layer1 (2x + 1) -> [2*gelu(1)+1, 2*gelu(-1)+1] + let y = m.forward(&CpuBackend, &x).unwrap(); + let expected0 = 2.0 * ops::gelu(1.0) + 1.0; + let expected1 = 2.0 * ops::gelu(-1.0) + 1.0; + assert!((y.data[0] - expected0).abs() < 1e-5); + assert!((y.data[1] - expected1).abs() < 1e-5); + } + + #[test] + fn rejects_wrong_input_width() { + let m = tiny_model(); + let x = Tensor::new(vec![1.0, 2.0, 3.0], vec![1, 3]).unwrap(); + assert!(matches!( + m.forward(&CpuBackend, &x), + Err(ModelError::WidthMismatch { .. }) + )); + } + + #[test] + fn rejects_disconnected_checkpoint_layers() { + let m = tiny_model(); + let mut map = HashMap::new(); + for (i, (w, b)) in m.layers.into_iter().enumerate() { + map.insert(format!("layers.{i}.weight"), w); + map.insert(format!("layers.{i}.bias"), b); + } + map.insert("layers.1.weight".into(), Tensor::zeros(vec![2, 3])); + assert!(matches!( + Mlp::from_tensors(&map), + Err(ModelError::WidthMismatch { .. }) + )); + } + + #[test] + fn forward_rejects_corrupt_inputs_and_empty_models() { + let corrupt = Tensor { + data: vec![1.0], + shape: vec![1, 2], + }; + assert!(matches!( + tiny_model().forward(&CpuBackend, &corrupt), + Err(ModelError::Tensor(_)) + )); + assert!(matches!( + Mlp { layers: Vec::new() }.forward(&CpuBackend, &corrupt), + Err(ModelError::NoLayers) + )); + } + + #[test] + fn checkpoint_rejects_zero_feature_width() { + let mut map = HashMap::new(); + map.insert("layers.0.weight".into(), Tensor::zeros(vec![2, 0])); + map.insert("layers.0.bias".into(), Tensor::zeros(vec![2])); + assert!(matches!( + Mlp::from_tensors(&map), + Err(ModelError::BadShape { .. }) + )); + } + + #[test] + fn forward_revalidates_mutated_public_weights_and_biases() { + let mut model = tiny_model(); + let input = Tensor::zeros(vec![1, 2]); + model.layers[0].0.data.pop(); + assert!(matches!( + model.forward(&CpuBackend, &input), + Err(ModelError::Tensor(_)) + )); + let mut model = tiny_model(); + model.layers[0].1 = Tensor::zeros(vec![1]); + assert!(matches!( + model.forward(&CpuBackend, &input), + Err(ModelError::BadShape { .. }) + )); + } +} diff --git a/src/ops.rs b/src/ops.rs new file mode 100644 index 0000000..cd4907c --- /dev/null +++ b/src/ops.rs @@ -0,0 +1,463 @@ +//! Portable, dependency-free CPU reference implementations of every operator. +//! +//! These functions define the *ground-truth* numerics for the engine. The Metal +//! kernels in `metal_backend` are validated against them by the parity tests in +//! `tests/parity.rs`, and the documented tolerance bounds in `docs/correctness.md` +//! refer to the maximum observed deviation between these and the GPU path. +//! +//! Conventions: +//! * All matrices are row-major. +//! * `gelu` uses the tanh approximation, matching `jax.nn.gelu(approximate=True)`. +//! * `rope` uses the rotate-half (GPT-NeoX / HF) convention. + +fn checked_elements(rows: usize, cols: usize) -> usize { + rows.checked_mul(cols) + .expect("operator dimensions overflow") +} + +/// Standard matrix multiply: `A`[m×k] · `B`[k×n] → `C`[m×n] (row-major). +pub fn matmul(a: &[f32], b: &[f32], m: usize, k: usize, n: usize) -> Vec { + assert_eq!(a.len(), checked_elements(m, k), "A has wrong length"); + assert_eq!(b.len(), checked_elements(k, n), "B has wrong length"); + let mut c = vec![0.0f32; checked_elements(m, n)]; + if c.is_empty() { + return c; + } + for row in 0..m { + for i in 0..k { + let a_ik = a[row * k + i]; + // Hoisting a_ik and walking B contiguously keeps the reference + // cache-friendly without changing the (well-defined) summation order. + let b_row = &b[i * n..i * n + n]; + let c_row = &mut c[row * n..row * n + n]; + for col in 0..n { + c_row[col] += a_ik * b_row[col]; + } + } + } + c +} + +/// Affine layer matching `equinox.nn.Linear`: `y = x · Wᵀ + b`. +/// +/// `x`[n×in], `w`[out×in], `b`[out] → `[n×out]`. +pub fn linear(x: &[f32], w: &[f32], b: &[f32], n: usize, in_f: usize, out_f: usize) -> Vec { + assert_eq!(x.len(), checked_elements(n, in_f)); + assert_eq!(w.len(), checked_elements(out_f, in_f)); + assert_eq!(b.len(), out_f); + let mut y = vec![0.0f32; checked_elements(n, out_f)]; + if y.is_empty() { + return y; + } + for row in 0..n { + for o in 0..out_f { + let mut acc = b[o]; + let x_row = &x[row * in_f..row * in_f + in_f]; + let w_row = &w[o * in_f..o * in_f + in_f]; + for i in 0..in_f { + acc += x_row[i] * w_row[i]; + } + y[row * out_f + o] = acc; + } + } + y +} + +const SQRT_2_OVER_PI: f32 = 0.7978846; + +/// GELU activation (tanh approximation), matching `jax.nn.gelu(approximate=True)`. +pub fn gelu(x: f32) -> f32 { + 0.5 * x * (1.0 + (SQRT_2_OVER_PI * (x + 0.044715 * x * x * x)).tanh()) +} + +pub fn gelu_inplace(x: &mut [f32]) { + for v in x.iter_mut() { + *v = gelu(*v); + } +} + +/// SiLU / swish activation: `x · sigmoid(x)`. +pub fn silu(x: f32) -> f32 { + x / (1.0 + (-x).exp()) +} + +/// Numerically-stable softmax over a single vector (in place). +pub fn softmax_inplace(x: &mut [f32]) { + let max = x.iter().copied().fold(f32::NEG_INFINITY, f32::max); + let mut sum = 0.0f32; + for v in x.iter_mut() { + *v = (*v - max).exp(); + sum += *v; + } + if sum > 0.0 { + for v in x.iter_mut() { + *v /= sum; + } + } +} + +/// Row-wise LayerNorm over the last dimension of size `d`. +/// `y = (x - mean) / sqrt(var + eps) * gamma + beta`, population variance. +pub fn layernorm( + x: &[f32], + gamma: &[f32], + beta: &[f32], + rows: usize, + d: usize, + eps: f32, +) -> Vec { + assert!(d > 0, "normalization width must be positive"); + assert!( + eps.is_finite() && eps >= 0.0, + "normalization epsilon must be finite and non-negative" + ); + assert_eq!(x.len(), checked_elements(rows, d)); + assert_eq!(gamma.len(), d); + assert_eq!(beta.len(), d); + let mut out = vec![0.0f32; checked_elements(rows, d)]; + for r in 0..rows { + let row = &x[r * d..r * d + d]; + let mean = row.iter().sum::() / d as f32; + let var = row.iter().map(|v| (v - mean) * (v - mean)).sum::() / d as f32; + let inv_std = 1.0 / (var + eps).sqrt(); + for c in 0..d { + out[r * d + c] = (row[c] - mean) * inv_std * gamma[c] + beta[c]; + } + } + out +} + +/// Row-wise RMSNorm over the last dimension of size `d`: `y = x / sqrt(mean(x²) + eps) * gamma`. +pub fn rmsnorm(x: &[f32], gamma: &[f32], rows: usize, d: usize, eps: f32) -> Vec { + assert!(d > 0, "normalization width must be positive"); + assert!( + eps.is_finite() && eps >= 0.0, + "normalization epsilon must be finite and non-negative" + ); + assert_eq!(x.len(), checked_elements(rows, d)); + assert_eq!(gamma.len(), d); + let mut out = vec![0.0f32; checked_elements(rows, d)]; + for r in 0..rows { + let row = &x[r * d..r * d + d]; + let ms = row.iter().map(|v| v * v).sum::() / d as f32; + let inv = 1.0 / (ms + eps).sqrt(); + for c in 0..d { + out[r * d + c] = row[c] * inv * gamma[c]; + } + } + out +} + +/// Applies rotary position embeddings (rotate-half / GPT-NeoX convention) to +/// `x`[rows×d] in place. `positions[r]` gives the absolute position of row `r`. +/// `d` must be even. +pub fn rope_inplace(x: &mut [f32], positions: &[usize], rows: usize, d: usize, theta: f32) { + assert!(d > 0, "rope head dim must be positive"); + assert!( + theta.is_finite() && theta > 0.0, + "rope theta must be finite and positive" + ); + assert_eq!(x.len(), checked_elements(rows, d)); + assert_eq!(positions.len(), rows); + assert_eq!(d % 2, 0, "rope head dim must be even"); + let half = d / 2; + for r in 0..rows { + let pos = positions[r] as f32; + let row = &mut x[r * d..r * d + d]; + for i in 0..half { + let inv_freq = theta.powf(-2.0 * i as f32 / d as f32); + let angle = pos * inv_freq; + let (sin, cos) = angle.sin_cos(); + let x1 = row[i]; + let x2 = row[i + half]; + row[i] = x1 * cos - x2 * sin; + row[i + half] = x2 * cos + x1 * sin; + } + } +} + +/// Single-head scaled dot-product attention. +/// +/// `q`[m×d] attends over `k`/`v`[seq×d] → `out`[m×d]. When `causal` is set, query +/// row `i` may only attend to keys `j ≤ i + q_offset` (matching cached generation, +/// where the `m` new queries start at absolute position `q_offset`). +#[allow(clippy::too_many_arguments)] +pub fn attention( + q: &[f32], + k: &[f32], + v: &[f32], + m: usize, + seq: usize, + d: usize, + causal: bool, + q_offset: usize, +) -> Vec { + assert!(d > 0, "attention width must be positive"); + assert!( + seq > 0 || m == 0, + "non-empty attention queries require keys" + ); + assert_eq!(q.len(), checked_elements(m, d)); + assert_eq!(k.len(), checked_elements(seq, d)); + assert_eq!(v.len(), checked_elements(seq, d)); + let scale = 1.0 / (d as f32).sqrt(); + let mut out = vec![0.0f32; checked_elements(m, d)]; + let mut scores = vec![0.0f32; seq]; + for qi in 0..m { + let limit = if causal { + qi.saturating_add(q_offset).saturating_add(1).min(seq) + } else { + seq + }; + for (kj, score) in scores.iter_mut().enumerate().take(limit) { + let mut dot = 0.0f32; + for di in 0..d { + dot += q[qi * d + di] * k[kj * d + di]; + } + *score = dot * scale; + } + softmax_inplace(&mut scores[..limit]); + for di in 0..d { + let mut acc = 0.0f32; + for (kj, &w) in scores[..limit].iter().enumerate() { + acc += w * v[kj * d + di]; + } + out[qi * d + di] = acc; + } + } + out +} + +/// Multi-head causal self-attention. Q/K/V are `[seq, heads*head_dim]` with the +/// heads laid out contiguously per row; output has the same shape. Query `i` +/// attends causally over keys `j <= i`. This is the GPT-2 attention reference. +pub fn mha(q: &[f32], k: &[f32], v: &[f32], seq: usize, heads: usize, head_dim: usize) -> Vec { + assert!( + heads > 0 && head_dim > 0, + "attention heads and head width must be positive" + ); + let hd = checked_elements(heads, head_dim); + assert_eq!(q.len(), checked_elements(seq, hd)); + assert_eq!(k.len(), checked_elements(seq, hd)); + assert_eq!(v.len(), checked_elements(seq, hd)); + let scale = 1.0 / (head_dim as f32).sqrt(); + let mut out = vec![0.0f32; checked_elements(seq, hd)]; + let mut scores = vec![0.0f32; seq]; + for h in 0..heads { + let base = h * head_dim; + for qi in 0..seq { + let limit = qi + 1; // causal + for (kj, score) in scores.iter_mut().enumerate().take(limit) { + let mut dot = 0.0f32; + for d in 0..head_dim { + dot += q[qi * hd + base + d] * k[kj * hd + base + d]; + } + *score = dot * scale; + } + softmax_inplace(&mut scores[..limit]); + for d in 0..head_dim { + let mut acc = 0.0f32; + for (kj, &w) in scores[..limit].iter().enumerate() { + acc += w * v[kj * hd + base + d]; + } + out[qi * hd + base + d] = acc; + } + } + } + out +} + +/// Dequantizes a per-row INT8 weight matrix: `out[r,c] = q[r,c] · scale[r]`. +pub fn dequantize_int8(q: &[i8], scales: &[f32], rows: usize, cols: usize) -> Vec { + assert_eq!(q.len(), checked_elements(rows, cols)); + assert_eq!(scales.len(), rows); + let mut out = vec![0.0f32; checked_elements(rows, cols)]; + if out.is_empty() { + return out; + } + for r in 0..rows { + let s = scales[r]; + for c in 0..cols { + out[r * cols + c] = q[r * cols + c] as f32 * s; + } + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + fn approx(a: &[f32], b: &[f32], tol: f32) { + assert_eq!(a.len(), b.len()); + for (x, y) in a.iter().zip(b) { + assert!((x - y).abs() <= tol, "expected {y}, got {x}"); + } + } + + #[test] + fn matmul_identity() { + let a = [1.0, 2.0, 3.0, 4.0]; // 2x2 + let id = [1.0, 0.0, 0.0, 1.0]; + assert_eq!(matmul(&a, &id, 2, 2, 2), a.to_vec()); + } + + #[test] + fn matmul_known() { + // [1 2 3] · [[1],[0],[-1]] = [1*1 + 2*0 + 3*-1] = [-2] + let a = [1.0, 2.0, 3.0]; + let b = [1.0, 0.0, -1.0]; + assert_eq!(matmul(&a, &b, 1, 3, 1), vec![-2.0]); + } + + #[test] + fn linear_matches_manual() { + // x = [1,2], W = [[1,0],[0,1],[1,1]] (out=3,in=2), b=[1,2,3] + let x = [1.0, 2.0]; + let w = [1.0, 0.0, 0.0, 1.0, 1.0, 1.0]; + let b = [1.0, 2.0, 3.0]; + let y = linear(&x, &w, &b, 1, 2, 3); + assert_eq!(y, vec![1.0 + 1.0, 2.0 + 2.0, 3.0 + 3.0]); + } + + #[test] + fn gelu_reference_points() { + assert!((gelu(0.0)).abs() < 1e-7); + // gelu(1) ≈ 0.8412 with tanh approximation + assert!((gelu(1.0) - 0.841_192).abs() < 1e-4); + assert!((gelu(-1.0) - -0.158_808).abs() < 1e-4); + } + + #[test] + fn softmax_sums_to_one() { + let mut x = [1.0, 2.0, 3.0, 4.0]; + softmax_inplace(&mut x); + let s: f32 = x.iter().sum(); + assert!((s - 1.0).abs() < 1e-6); + // monotonic: larger logit -> larger prob + assert!(x[3] > x[0]); + } + + #[test] + fn layernorm_zero_mean_unit_var() { + let x = [1.0, 2.0, 3.0, 4.0]; + let gamma = [1.0; 4]; + let beta = [0.0; 4]; + let y = layernorm(&x, &gamma, &beta, 1, 4, 1e-5); + let mean: f32 = y.iter().sum::() / 4.0; + assert!(mean.abs() < 1e-4); + let var: f32 = y.iter().map(|v| v * v).sum::() / 4.0; + assert!((var - 1.0).abs() < 1e-2); + } + + #[test] + fn rmsnorm_scales_correctly() { + let x = [3.0, 4.0]; // rms = sqrt((9+16)/2) = sqrt(12.5) + let gamma = [1.0, 1.0]; + let y = rmsnorm(&x, &gamma, 1, 2, 0.0); + let rms = (12.5f32).sqrt(); + approx(&y, &[3.0 / rms, 4.0 / rms], 1e-6); + } + + #[test] + fn rope_preserves_norm() { + // Rotation is norm-preserving per (i, i+half) pair. + let mut x = [0.3, 0.7, -0.2, 0.5]; + let before: f32 = x.iter().map(|v| v * v).sum(); + rope_inplace(&mut x, &[5], 1, 4, 10000.0); + let after: f32 = x.iter().map(|v| v * v).sum(); + assert!((before - after).abs() < 1e-5); + } + + #[test] + fn rope_position_zero_is_identity() { + let mut x = [0.3, 0.7, -0.2, 0.5]; + let orig = x; + rope_inplace(&mut x, &[0], 1, 4, 10000.0); + approx(&x, &orig, 1e-6); + } + + #[test] + fn attention_uniform_when_keys_equal() { + // All keys identical -> uniform weights -> output = mean of V rows. + let d = 2; + let seq = 3; + let q = [1.0, 1.0]; + let k = [0.5, 0.5, 0.5, 0.5, 0.5, 0.5]; + let v = [1.0, 0.0, 2.0, 0.0, 3.0, 0.0]; + let out = attention(&q, &k, &v, 1, seq, d, false, 0); + approx(&out, &[2.0, 0.0], 1e-5); // mean of [1,2,3] = 2 + } + + #[test] + fn attention_causal_first_query_sees_only_first_key() { + let d = 2; + let seq = 2; + let q = [1.0, 0.0]; // query at offset 0 -> only key 0 visible + let k = [10.0, 0.0, 0.0, 10.0]; + let v = [5.0, 0.0, 0.0, 9.0]; + let out = attention(&q, &k, &v, 1, seq, d, true, 0); + approx(&out, &[5.0, 0.0], 1e-5); + } + + #[test] + fn dequant_int8_roundtrip() { + let q = [10i8, -10, 100]; + let scales = [0.1f32]; + let out = dequantize_int8(&q, &scales, 1, 3); + approx(&out, &[1.0, -1.0, 10.0], 1e-6); + } + + #[test] + #[should_panic(expected = "operator dimensions overflow")] + fn matmul_rejects_overflow_before_allocation() { + matmul(&[], &[], usize::MAX, 2, 1); + } + + #[test] + fn normalization_rejects_invalid_epsilon() { + for eps in [f32::NAN, f32::INFINITY, -1.0] { + assert!( + std::panic::catch_unwind(|| layernorm(&[1.0], &[1.0], &[0.0], 1, 1, eps)).is_err() + ); + assert!(std::panic::catch_unwind(|| rmsnorm(&[1.0], &[1.0], 1, 1, eps)).is_err()); + } + } + + #[test] + fn normalization_distinguishes_empty_batches_from_zero_width() { + assert!(layernorm(&[], &[1.0], &[0.0], 0, 1, 1e-5).is_empty()); + assert!(rmsnorm(&[], &[1.0], 0, 1, 0.0).is_empty()); + assert!(std::panic::catch_unwind(|| layernorm(&[], &[], &[], 1, 0, 0.0)).is_err()); + } + + #[test] + fn rope_rejects_invalid_frequency_bases() { + for theta in [0.0, -1.0, f32::NAN, f32::INFINITY] { + assert!( + std::panic::catch_unwind(|| rope_inplace(&mut [1.0, 2.0], &[0], 1, 2, theta)) + .is_err() + ); + } + } + + #[test] + fn attention_large_offset_attends_to_all_available_keys() { + let out = attention(&[0.0], &[0.0, 0.0], &[2.0, 4.0], 1, 2, 1, true, usize::MAX); + assert_eq!(out, vec![3.0]); + } + + #[test] + fn mha_rejects_zero_and_overflowing_head_dimensions() { + assert!(std::panic::catch_unwind(|| mha(&[], &[], &[], 0, 0, 1)).is_err()); + assert!(std::panic::catch_unwind(|| mha(&[], &[], &[], 0, usize::MAX, 2)).is_err()); + assert!(mha(&[], &[], &[], 0, 1, 1).is_empty()); + } + + #[test] + fn empty_matrix_outputs_skip_unused_rows() { + assert!(matmul(&[], &[], usize::MAX, 0, 0).is_empty()); + assert!(linear(&[], &[], &[], usize::MAX, 0, 0).is_empty()); + assert!(dequantize_int8(&[], &[1.0; 3], 3, 0).is_empty()); + } +} diff --git a/src/shaders/compute.metal b/src/shaders/compute.metal index 7f6d486..4f9eddb 100644 --- a/src/shaders/compute.metal +++ b/src/shaders/compute.metal @@ -1,6 +1,21 @@ #include using namespace metal; +// GELU (tanh approximation), matching ops::gelu / jax.nn.gelu(approximate=True). +// +// The tanh argument is clamped: Metal's fast-math `tanh` evaluates exp(2·arg), +// which overflows to inf (→ NaN) once the argument is large (GPT-2 activations +// can drive it past ~70). tanh has already saturated to ±1 by |arg|=15, so the +// clamp is numerically exact in f32 while avoiding the overflow. +inline float gelu_approx(float x) { + float inner = 0.7978845608f * (x + 0.044715f * x * x * x); + inner = clamp(inner, -15.0f, 15.0f); + return 0.5f * x * (1.0f + tanh(inner)); +} + +// --------------------------------------------------------------------------- +// Dense matmul: C[M,N] = A[M,K] * B[K,N] (row-major, naive one-thread-per-output) +// --------------------------------------------------------------------------- kernel void matmul( device const float* A [[buffer(0)]], device const float* B [[buffer(1)]], @@ -12,9 +27,8 @@ kernel void matmul( ) { uint row = gid.y; uint col = gid.x; - if (row < M && col < N) { - float sum = 0.0; + float sum = 0.0f; for (uint i = 0; i < K; ++i) { sum += A[row * K + i] * B[i * N + col]; } @@ -22,6 +36,34 @@ kernel void matmul( } } +// --------------------------------------------------------------------------- +// Linear (equinox.nn.Linear): Y[N,Out] = X[N,In] * W[Out,In]^T + B[Out] +// Avoids an explicit weight transpose by indexing W row-major as [Out,In]. +// --------------------------------------------------------------------------- +kernel void linear( + device const float* X [[buffer(0)]], + device const float* W [[buffer(1)]], + device const float* Bias [[buffer(2)]], + device float* Y [[buffer(3)]], + constant uint& Rows [[buffer(4)]], + constant uint& In [[buffer(5)]], + constant uint& Out [[buffer(6)]], + uint2 gid [[thread_position_in_grid]] +) { + uint row = gid.y; + uint o = gid.x; + if (row < Rows && o < Out) { + float acc = Bias[o]; + for (uint i = 0; i < In; ++i) { + acc += X[row * In + i] * W[o * In + i]; + } + Y[row * Out + o] = acc; + } +} + +// --------------------------------------------------------------------------- +// Weight-only INT8 matmul: C = (A_int8 * per-row scale) * B +// --------------------------------------------------------------------------- kernel void quant_matmul( device const char* A_int8 [[buffer(0)]], device const float* B_f32 [[buffer(1)]], @@ -34,9 +76,8 @@ kernel void quant_matmul( ) { uint row = gid.y; uint col = gid.x; - if (row < M && col < N) { - float sum = 0.0; + float sum = 0.0f; float scale = A_scales[row]; for (uint i = 0; i < K; ++i) { float a_val = (float)A_int8[row * K + i] * scale; @@ -46,71 +87,95 @@ kernel void quant_matmul( } } -// Fused Attention with KV-Cache support +// Elementwise GELU. +kernel void gelu_forward( + device const float* X [[buffer(0)]], + device float* Y [[buffer(1)]], + constant uint& N [[buffer(2)]], + uint gid [[thread_position_in_grid]] +) { + if (gid < N) { + Y[gid] = gelu_approx(X[gid]); + } +} + +// --------------------------------------------------------------------------- +// Single-head attention with KV-cache. +// +// One thread per query row. The previous version recomputed the full Q·K dot +// product inside the output-dimension loop, making it O(M·S·D^2). This computes +// each score once per key, so the cost is O(M·S·D). Causal masking restricts +// query row i to keys j <= i + QOffset (QOffset = absolute position of the first +// new query, i.e. the cache length before this step). +// --------------------------------------------------------------------------- kernel void kv_attention( - device const float* Q [[buffer(0)]], // M x D (Queries for current step) - device const float* K_cache [[buffer(1)]], // MaxSeq x D (Keys Cache) - device const float* V_cache [[buffer(2)]], // MaxSeq x D (Values Cache) - device float* O [[buffer(3)]], // M x D (Output) - constant uint& M [[buffer(4)]], // Current query length (usually 1 for generation) - constant uint& CurSeqLen [[buffer(5)]], // Current total sequence length including new tokens - constant uint& D [[buffer(6)]], // Head Dimension + device const float* Q [[buffer(0)]], // M x D + device const float* K_cache [[buffer(1)]], // CurSeqLen x D + device const float* V_cache [[buffer(2)]], // CurSeqLen x D + device float* O [[buffer(3)]], // M x D + constant uint& M [[buffer(4)]], + constant uint& CurSeqLen [[buffer(5)]], + constant uint& D [[buffer(6)]], + constant uint& Causal [[buffer(7)]], + constant uint& QOffset [[buffer(8)]], uint gid [[thread_position_in_grid]] ) { uint q_idx = gid; if (q_idx >= M) return; - float max_score = -1e9; - - // Iterate over the cached keys up to the current sequence length - for (uint k_idx = 0; k_idx < CurSeqLen; ++k_idx) { - float score = 0.0; - for (uint d = 0; d < D; ++d) { - score += Q[q_idx * D + d] * K_cache[k_idx * D + d]; - } - score /= sqrt((float)D); - max_score = max(max_score, score); - } - - float sum_exp = 0.0; - for (uint k_idx = 0; k_idx < CurSeqLen; ++k_idx) { - float score = 0.0; - for (uint d = 0; d < D; ++d) { - score += Q[q_idx * D + d] * K_cache[k_idx * D + d]; - } - score /= sqrt((float)D); - sum_exp += exp(score - max_score); - } - - for (uint d = 0; d < D; ++d) { - float out_val = 0.0; - for (uint k_idx = 0; k_idx < CurSeqLen; ++k_idx) { - float score = 0.0; - for (uint d_inner = 0; d_inner < D; ++d_inner) { - score += Q[q_idx * D + d_inner] * K_cache[k_idx * D + d_inner]; - } - score /= sqrt((float)D); - float weight = exp(score - max_score) / sum_exp; - out_val += weight * V_cache[k_idx * D + d]; - } - O[q_idx * D + d] = out_val; + uint limit = CurSeqLen; + if (Causal != 0u) { + uint c = q_idx + QOffset + 1u; + limit = (c < CurSeqLen) ? c : CurSeqLen; + } + + // Output starts at zero; if nothing is visible it stays zero. + for (uint d = 0; d < D; ++d) O[q_idx * D + d] = 0.0f; + if (limit == 0u) return; + + float inv_sqrt_d = rsqrt((float)D); + + // Pass 1: running max for numerical stability. + float max_score = -INFINITY; + for (uint k = 0; k < limit; ++k) { + float s = 0.0f; + for (uint d = 0; d < D; ++d) s += Q[q_idx * D + d] * K_cache[k * D + d]; + s *= inv_sqrt_d; + max_score = max(max_score, s); + } + + // Pass 2: denominator. + float sum_exp = 0.0f; + for (uint k = 0; k < limit; ++k) { + float s = 0.0f; + for (uint d = 0; d < D; ++d) s += Q[q_idx * D + d] * K_cache[k * D + d]; + s *= inv_sqrt_d; + sum_exp += exp(s - max_score); + } + + // Pass 3: weighted sum of V (score computed once per key, not per dim). + for (uint k = 0; k < limit; ++k) { + float s = 0.0f; + for (uint d = 0; d < D; ++d) s += Q[q_idx * D + d] * K_cache[k * D + d]; + s *= inv_sqrt_d; + float w = exp(s - max_score) / sum_exp; + for (uint d = 0; d < D; ++d) O[q_idx * D + d] += w * V_cache[k * D + d]; } } -// Helper kernel to update KV-Cache with new tokens +// Append M new tokens to the KV cache at row `Offset`. kernel void update_kv_cache( device const float* NewK [[buffer(0)]], // M x D device const float* NewV [[buffer(1)]], // M x D device float* K_cache [[buffer(2)]], // MaxSeq x D device float* V_cache [[buffer(3)]], // MaxSeq x D - constant uint& M [[buffer(4)]], // New tokens length - constant uint& Offset [[buffer(5)]], // Starting position in cache - constant uint& D [[buffer(6)]], // Head Dimension + constant uint& M [[buffer(4)]], + constant uint& Offset [[buffer(5)]], + constant uint& D [[buffer(6)]], uint2 gid [[thread_position_in_grid]] ) { uint tok_idx = gid.x; uint d_idx = gid.y; - if (tok_idx < M && d_idx < D) { uint cache_pos = (tok_idx + Offset) * D + d_idx; uint input_pos = tok_idx * D + d_idx; @@ -118,3 +183,172 @@ kernel void update_kv_cache( V_cache[cache_pos] = NewV[input_pos]; } } + +// Row-wise LayerNorm over the last dimension. +kernel void layernorm( + device const float* X [[buffer(0)]], + device const float* Gamma [[buffer(1)]], + device const float* Beta [[buffer(2)]], + device float* Y [[buffer(3)]], + constant uint& Rows [[buffer(4)]], + constant uint& D [[buffer(5)]], + constant float& Eps [[buffer(6)]], + uint gid [[thread_position_in_grid]] +) { + uint r = gid; + if (r >= Rows) return; + float mean = 0.0f; + for (uint c = 0; c < D; ++c) mean += X[r * D + c]; + mean /= (float)D; + float var = 0.0f; + for (uint c = 0; c < D; ++c) { + float diff = X[r * D + c] - mean; + var += diff * diff; + } + var /= (float)D; + float inv_std = rsqrt(var + Eps); + for (uint c = 0; c < D; ++c) { + Y[r * D + c] = (X[r * D + c] - mean) * inv_std * Gamma[c] + Beta[c]; + } +} + +// Row-wise RMSNorm over the last dimension. +kernel void rmsnorm( + device const float* X [[buffer(0)]], + device const float* Gamma [[buffer(1)]], + device float* Y [[buffer(2)]], + constant uint& Rows [[buffer(3)]], + constant uint& D [[buffer(4)]], + constant float& Eps [[buffer(5)]], + uint gid [[thread_position_in_grid]] +) { + uint r = gid; + if (r >= Rows) return; + float ms = 0.0f; + for (uint c = 0; c < D; ++c) ms += X[r * D + c] * X[r * D + c]; + ms /= (float)D; + float inv = rsqrt(ms + Eps); + for (uint c = 0; c < D; ++c) { + Y[r * D + c] = X[r * D + c] * inv * Gamma[c]; + } +} + +// Rotary position embedding (rotate-half / GPT-NeoX convention), applied in place. +kernel void rope( + device float* X [[buffer(0)]], + device const uint* Positions [[buffer(1)]], + constant uint& Rows [[buffer(2)]], + constant uint& D [[buffer(3)]], + constant float& Theta [[buffer(4)]], + uint gid [[thread_position_in_grid]] +) { + uint r = gid; + if (r >= Rows) return; + uint half_d = D / 2; + float pos = (float)Positions[r]; + for (uint i = 0; i < half_d; ++i) { + float inv_freq = pow(Theta, -2.0f * (float)i / (float)D); + float angle = pos * inv_freq; + float s = sin(angle); + float c = cos(angle); + float x1 = X[r * D + i]; + float x2 = X[r * D + i + half_d]; + X[r * D + i] = x1 * c - x2 * s; + X[r * D + i + half_d] = x2 * c + x1 * s; + } +} + +// --------------------------------------------------------------------------- +// Tiled matmul: C[M,N] = A[M,K] * B[K,N] using threadgroup shared memory. +// +// Each 16x16 threadgroup cooperatively stages tiles of A and B into fast +// on-chip memory, so each global element is read once per tile instead of +// once per output. This is the standard shared-memory GEMM and is many times +// faster than the naive one-thread-per-output `matmul` above. +// --------------------------------------------------------------------------- +#define TILE 16 +kernel void matmul_tiled( + device const float* A [[buffer(0)]], + device const float* B [[buffer(1)]], + device float* C [[buffer(2)]], + constant uint& M [[buffer(3)]], + constant uint& N [[buffer(4)]], + constant uint& K [[buffer(5)]], + uint2 tid [[thread_position_in_threadgroup]], + uint2 gid [[thread_position_in_grid]] +) { + threadgroup float As[TILE][TILE]; + threadgroup float Bs[TILE][TILE]; + + uint row = gid.y; + uint col = gid.x; + float acc = 0.0f; + + uint n_tiles = (K + TILE - 1) / TILE; + for (uint t = 0; t < n_tiles; ++t) { + uint a_col = t * TILE + tid.x; + uint b_row = t * TILE + tid.y; + As[tid.y][tid.x] = (row < M && a_col < K) ? A[row * K + a_col] : 0.0f; + Bs[tid.y][tid.x] = (b_row < K && col < N) ? B[b_row * N + col] : 0.0f; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint k = 0; k < TILE; ++k) { + acc += As[tid.y][k] * Bs[k][tid.x]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (row < M && col < N) { + C[row * N + col] = acc; + } +} + +// --------------------------------------------------------------------------- +// Multi-head causal self-attention. +// +// Q/K/V are [S, H*Dh] (heads laid out contiguously per row). One thread handles +// one (head, query) pair, attending causally over keys j <= query. Output is +// [S, H*Dh]. This is what powers the GPT-2 attention blocks. +// --------------------------------------------------------------------------- +kernel void mha( + device const float* Q [[buffer(0)]], + device const float* K [[buffer(1)]], + device const float* V [[buffer(2)]], + device float* O [[buffer(3)]], + constant uint& S [[buffer(4)]], // sequence length + constant uint& H [[buffer(5)]], // number of heads + constant uint& Dh [[buffer(6)]], // head dimension + uint2 gid [[thread_position_in_grid]] +) { + uint qpos = gid.x; + uint head = gid.y; + if (qpos >= S || head >= H) return; + + uint HD = H * Dh; + uint base = head * Dh; + float scale = rsqrt((float)Dh); + + // Pass 1: max logit over visible keys (causal: k <= qpos). + float max_s = -INFINITY; + for (uint k = 0; k <= qpos; ++k) { + float s = 0.0f; + for (uint d = 0; d < Dh; ++d) s += Q[qpos * HD + base + d] * K[k * HD + base + d]; + s *= scale; + max_s = max(max_s, s); + } + // Pass 2: softmax denominator. + float sum = 0.0f; + for (uint k = 0; k <= qpos; ++k) { + float s = 0.0f; + for (uint d = 0; d < Dh; ++d) s += Q[qpos * HD + base + d] * K[k * HD + base + d]; + sum += exp(s * scale - max_s); + } + // Pass 3: weighted sum of V. + for (uint d = 0; d < Dh; ++d) O[qpos * HD + base + d] = 0.0f; + for (uint k = 0; k <= qpos; ++k) { + float s = 0.0f; + for (uint d = 0; d < Dh; ++d) s += Q[qpos * HD + base + d] * K[k * HD + base + d]; + float w = exp(s * scale - max_s) / sum; + for (uint d = 0; d < Dh; ++d) O[qpos * HD + base + d] += w * V[k * HD + base + d]; + } +} diff --git a/src/tensor.rs b/src/tensor.rs index 7096e9d..ef65e18 100644 --- a/src/tensor.rs +++ b/src/tensor.rs @@ -1,7 +1,26 @@ use bytemuck::Pod; use safetensors::tensor::Dtype; +use thiserror::Error; -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Error, Debug)] +pub enum TensorError { + #[error("Unsupported dtype mapping: {0:?}")] + UnsupportedDtype(Dtype), + #[error("Shape mismatch: expected {expected} bytes, found {found}")] + ShapeMismatch { expected: usize, found: usize }, + #[error("Shape mismatch: expected {expected} elements, found {found}")] + ElementCountMismatch { expected: usize, found: usize }, + #[error("Buffer overflow detected when computing tensor size")] + BufferOverflow, + #[error("Failed to allocate tensor buffer for {elements} elements")] + AllocationFailed { elements: usize }, + #[error("Expected a {expected}-D tensor, found shape {shape:?}")] + RankMismatch { expected: usize, shape: Vec }, + #[error("Tensor is dtype {0:?}, expected F32")] + NotF32(DataType), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum DataType { F32, F16, @@ -12,23 +31,38 @@ pub enum DataType { I64, } -impl From for DataType { - fn from(dt: Dtype) -> Self { +impl DataType { + pub fn size_in_bytes(&self) -> usize { + match self { + DataType::F32 | DataType::I32 => 4, + DataType::F16 | DataType::BF16 => 2, + DataType::I64 => 8, + DataType::I8 | DataType::U8 => 1, + } + } +} + +impl TryFrom for DataType { + type Error = TensorError; + + fn try_from(dt: Dtype) -> Result { match dt { - Dtype::F32 => DataType::F32, - Dtype::F16 => DataType::F16, - Dtype::BF16 => DataType::BF16, - Dtype::I8 => DataType::I8, - Dtype::U8 => DataType::U8, - Dtype::I32 => DataType::I32, - Dtype::I64 => DataType::I64, - _ => unimplemented!("Unsupported dtype mapping: {:?}", dt), + Dtype::F32 => Ok(DataType::F32), + Dtype::F16 => Ok(DataType::F16), + Dtype::BF16 => Ok(DataType::BF16), + Dtype::I8 => Ok(DataType::I8), + Dtype::U8 => Ok(DataType::U8), + Dtype::I32 => Ok(DataType::I32), + Dtype::I64 => Ok(DataType::I64), + _ => Err(TensorError::UnsupportedDtype(dt)), } } } -/// A zero-copy view into a memory-mapped tensor buffer. -#[derive(Debug)] +/// A view into a memory-mapped tensor buffer: the bytes are borrowed zero-copy +/// from the mapping, while the (tiny) shape vector is owned so the view can +/// outlive the transient parser handle it came from. +#[derive(Debug, Clone)] pub struct TensorView<'data> { pub shape: Vec, pub dtype: DataType, @@ -36,13 +70,296 @@ pub struct TensorView<'data> { } impl<'data> TensorView<'data> { - pub fn new(shape: Vec, dtype: DataType, data: &'data [u8]) -> Self { - Self { shape, dtype, data } + pub fn new(shape: &[usize], dtype: DataType, data: &'data [u8]) -> Result { + let expected_bytes = num_bytes(shape, dtype)?; + if data.len() != expected_bytes { + return Err(TensorError::ShapeMismatch { + expected: expected_bytes, + found: data.len(), + }); + } + Ok(Self { + shape: shape.to_vec(), + dtype, + data, + }) + } + + /// Number of elements described by the shape. + pub fn numel(&self) -> usize { + self.shape.iter().product() } - /// Safely casts the underlying byte buffer to a typed slice if the dtype matches. + /// Attempts to reinterpret bytes as a Pod slice, checking alignment and length. + /// This does not check semantic dtype or convert byte order. Use + /// [`TensorView::to_tensor_f32`] for portable, owned F32 materialization. pub fn as_slice(&self) -> Option<&[T]> { - // In a full implementation, we would verify `T` matches `self.dtype`. bytemuck::try_cast_slice(self.data).ok() } + + /// Materializes an owned f32 [`Tensor`], copying out of the mapped buffer. + /// + /// Reads little-endian f32 directly from the bytes, so it works even when the + /// tensor's offset in the mmap is not 4-byte aligned (which a zero-copy + /// `bytemuck` cast would reject). Only F32 source data is supported. + pub fn to_tensor_f32(&self) -> Result { + if self.dtype != DataType::F32 { + return Err(TensorError::NotF32(self.dtype)); + } + let expected = num_bytes(&self.shape, self.dtype)?; + if self.data.len() != expected { + return Err(TensorError::ShapeMismatch { + expected, + found: self.data.len(), + }); + } + let data = self + .data + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(); + Ok(Tensor { + data, + shape: self.shape.clone(), + }) + } +} + +fn checked_numel(shape: &[usize]) -> Result { + if shape.contains(&0) { + return Ok(0); + } + shape.iter().try_fold(1usize, |elements, &dim| { + elements.checked_mul(dim).ok_or(TensorError::BufferOverflow) + }) +} + +fn num_bytes(shape: &[usize], dtype: DataType) -> Result { + checked_numel(shape)? + .checked_mul(dtype.size_in_bytes()) + .ok_or(TensorError::BufferOverflow) +} + +/// A simple owned, row-major f32 tensor used by the CPU reference ops and as the +/// common currency between backends. +#[derive(Debug, Clone, PartialEq)] +pub struct Tensor { + pub data: Vec, + pub shape: Vec, +} + +impl Tensor { + /// Creates a tensor from data, validating that the element count matches the shape. + /// Returns [`TensorError::BufferOverflow`] if the shape's element count overflows. + pub fn new(data: Vec, shape: Vec) -> Result { + let expected = checked_numel(&shape)?; + if data.len() != expected { + return Err(TensorError::ElementCountMismatch { + expected, + found: data.len(), + }); + } + Ok(Self { data, shape }) + } + + /// A zero-filled tensor of the given shape. + /// + /// # Panics + /// + /// Panics if the shape overflows or the backing buffer cannot be allocated. + /// Use [`Tensor::try_zeros`] when those failures need to be handled. + pub fn zeros(shape: Vec) -> Self { + Self::try_zeros(shape).expect("failed to allocate zero-filled tensor") + } + + /// Creates a zero-filled tensor, returning shape-overflow and allocation errors. + pub fn try_zeros(shape: Vec) -> Result { + let elements = checked_numel(&shape)?; + let mut data = Vec::new(); + data.try_reserve_exact(elements) + .map_err(|_| TensorError::AllocationFailed { elements })?; + data.resize(elements, 0.0); + Ok(Self { data, shape }) + } + + /// Checks that the public shape and data fields still describe the same tensor. + pub fn validate(&self) -> Result<(), TensorError> { + let expected = checked_numel(&self.shape)?; + if self.data.len() != expected { + return Err(TensorError::ElementCountMismatch { + expected, + found: self.data.len(), + }); + } + Ok(()) + } + + pub fn numel(&self) -> usize { + self.data.len() + } + + /// Interprets the tensor as 2-D, returning `(rows, cols)`. + pub fn dims2(&self) -> Result<(usize, usize), TensorError> { + match self.shape.as_slice() { + [r, c] => Ok((*r, *c)), + _ => Err(TensorError::RankMismatch { + expected: 2, + shape: self.shape.clone(), + }), + } + } + + /// Maximum absolute element-wise difference against another tensor of equal shape. + /// Useful for parity / tolerance assertions. Returns `f32::INFINITY` on shape or + /// data-length mismatch, or when an element-wise difference is not finite. + pub fn max_abs_diff(&self, other: &Tensor) -> f32 { + if self.shape != other.shape + || self.data.len() != other.data.len() + || self.validate().is_err() + || other.validate().is_err() + { + return f32::INFINITY; + } + let mut max_diff = 0.0_f32; + for (a, b) in self.data.iter().zip(&other.data) { + let diff = (a - b).abs(); + if !diff.is_finite() { + return f32::INFINITY; + } + max_diff = max_diff.max(diff); + } + max_diff + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_valid_tensor_view() { + let data = vec![0u8; 8]; + let view = TensorView::new(&[2, 1], DataType::F32, &data); + assert!(view.is_ok()); + } + + #[test] + fn test_shape_mismatch() { + let data = vec![0u8; 7]; // F32 requires multiple of 4 + let view = TensorView::new(&[2, 1], DataType::F32, &data); + assert!(matches!(view, Err(TensorError::ShapeMismatch { .. }))); + } + + #[test] + fn test_buffer_overflow() { + let data = vec![0u8; 8]; + let view = TensorView::new(&[usize::MAX, 2], DataType::F32, &data); + assert!(matches!(view, Err(TensorError::BufferOverflow))); + } + + #[test] + fn test_view_to_tensor_roundtrip() { + let floats = [1.0f32, 2.0, 3.0, 4.0]; + let bytes: &[u8] = bytemuck::cast_slice(&floats); + let view = TensorView::new(&[2, 2], DataType::F32, bytes).unwrap(); + let t = view.to_tensor_f32().unwrap(); + assert_eq!(t.shape, vec![2, 2]); + assert_eq!(t.data, vec![1.0, 2.0, 3.0, 4.0]); + } + + #[test] + fn test_max_abs_diff() { + let a = Tensor::new(vec![1.0, 2.0, 3.0], vec![3]).unwrap(); + let b = Tensor::new(vec![1.0, 2.5, 3.0], vec![3]).unwrap(); + assert!((a.max_abs_diff(&b) - 0.5).abs() < 1e-9); + } + + #[test] + fn test_owned_tensor_rejects_shape_overflow() { + assert!(matches!( + Tensor::new(Vec::new(), vec![usize::MAX, 2]), + Err(TensorError::BufferOverflow) + )); + } + + #[test] + fn test_max_abs_diff_rejects_non_finite_values() { + let finite = Tensor::new(vec![0.0], vec![1]).unwrap(); + for value in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] { + let invalid = Tensor::new(vec![value], vec![1]).unwrap(); + assert_eq!(invalid.max_abs_diff(&finite), f32::INFINITY); + assert_eq!(finite.max_abs_diff(&invalid), f32::INFINITY); + assert_eq!(invalid.max_abs_diff(&invalid), f32::INFINITY); + } + } + + #[test] + fn test_max_abs_diff_rejects_mismatched_data_lengths() { + let valid = Tensor::new(vec![1.0, 2.0], vec![2]).unwrap(); + let truncated = Tensor { + data: vec![1.0], + shape: vec![2], + }; + assert_eq!(valid.max_abs_diff(&truncated), f32::INFINITY); + assert_eq!(truncated.max_abs_diff(&valid), f32::INFINITY); + } + + #[test] + fn test_fallible_zeros_rejects_shape_overflow() { + assert!(matches!( + Tensor::try_zeros(vec![usize::MAX, 2]), + Err(TensorError::BufferOverflow) + )); + } + + #[test] + fn test_fallible_zeros_preserves_empty_shape() { + let tensor = Tensor::try_zeros(vec![2, 0, 3]).unwrap(); + assert_eq!(tensor.shape, vec![2, 0, 3]); + assert!(tensor.data.is_empty()); + } + + #[test] + fn test_fallible_zeros_initializes_values() { + let tensor = Tensor::try_zeros(vec![2, 2]).unwrap(); + assert_eq!(tensor.data, vec![0.0; 4]); + } + + #[test] + fn zero_dimensions_short_circuit_overflow_in_any_order() { + for shape in [vec![usize::MAX, 2, 0], vec![0, usize::MAX, 2]] { + assert!(Tensor::try_zeros(shape).unwrap().data.is_empty()); + } + } + + #[test] + fn conversion_revalidates_public_view_fields() { + let bytes = [0u8; 4]; + let view = TensorView { + shape: vec![2], + dtype: DataType::F32, + data: &bytes, + }; + assert!(matches!( + view.to_tensor_f32(), + Err(TensorError::ShapeMismatch { .. }) + )); + let overflowing = TensorView { + shape: vec![usize::MAX, 2], + ..view + }; + assert!(matches!( + overflowing.to_tensor_f32(), + Err(TensorError::BufferOverflow) + )); + } + + #[test] + fn matching_corrupt_buffers_do_not_pass_parity() { + let corrupt = Tensor { + data: vec![1.0], + shape: vec![2], + }; + assert_eq!(corrupt.max_abs_diff(&corrupt), f32::INFINITY); + } } diff --git a/src/tokenizer.rs b/src/tokenizer.rs new file mode 100644 index 0000000..56e28c0 --- /dev/null +++ b/src/tokenizer.rs @@ -0,0 +1,471 @@ +//! GPT-2 byte-level BPE tokenizer, implemented from scratch. +//! +//! Mirrors the original GPT-2 encoder: bytes are mapped into a printable +//! unicode alphabet, text is pre-tokenized into word-like chunks, and byte-pair +//! merges are applied in rank order from `merges.txt`. Loads the same +//! `vocab.json` / `merges.txt` HuggingFace ships for `gpt2`. + +use std::collections::HashMap; +use std::fs; +use std::path::Path; + +use thiserror::Error; + +#[derive(Error, Debug)] +pub enum TokenizerError { + #[error("unknown tokenizer ID: {0}")] + UnknownTokenId(usize), + #[error("duplicate BPE merge pair: {0:?}")] + DuplicateMerge(String), + #[error("invalid BPE merge line: {0:?}")] + InvalidMergeLine(String), + #[error("IO error: {0}")] + Io(#[from] std::io::Error), + #[error("failed to parse vocab.json: {0}")] + Json(#[from] serde_json::Error), + #[error("vocab.json assigns token ID {0} to more than one token")] + DuplicateTokenId(usize), + #[error("vocab.json is missing token ID {0} from its dense range")] + MissingTokenId(usize), + #[error("vocab.json is missing tokenizer symbol {0:?}")] + MissingVocabSymbol(String), +} + +pub struct Tokenizer { + encoder: HashMap, + decoder: HashMap, + bpe_ranks: HashMap<(String, String), usize>, + byte_encoder: [char; 256], + byte_decoder: HashMap, +} + +/// GPT-2's reversible bytes→unicode mapping (every byte gets a printable char). +fn bytes_to_unicode() -> [char; 256] { + let mut in_set = [false; 256]; + let mut cp = [0u32; 256]; + let push_range = |a: u32, b: u32, in_set: &mut [bool; 256], cp: &mut [u32; 256]| { + for c in a..=b { + in_set[c as usize] = true; + cp[c as usize] = c; + } + }; + push_range(b'!' as u32, b'~' as u32, &mut in_set, &mut cp); + push_range(0xA1, 0xAC, &mut in_set, &mut cp); + push_range(0xAE, 0xFF, &mut in_set, &mut cp); + let mut n = 0u32; + for b in 0..256usize { + if !in_set[b] { + cp[b] = 256 + n; + n += 1; + } + } + let mut arr = ['\0'; 256]; + for b in 0..256usize { + arr[b] = char::from_u32(cp[b]).unwrap(); + } + arr +} + +impl Tokenizer { + /// Loads a tokenizer from vocab.json and merges.txt. + pub fn from_files(vocab_path: &Path, merges_path: &Path) -> Result { + let vocab_raw = fs::read_to_string(vocab_path)?; + let vocab: HashMap = serde_json::from_str(&vocab_raw)?; + let merges_raw = fs::read_to_string(merges_path)?; + Self::from_assets(vocab, &merges_raw) + } + + fn from_assets( + vocab: HashMap, + merges_raw: &str, + ) -> Result { + let mut decoder = HashMap::with_capacity(vocab.len()); + for (token, &id) in &vocab { + if decoder.insert(id, token.clone()).is_some() { + return Err(TokenizerError::DuplicateTokenId(id)); + } + } + for id in 0..vocab.len() { + if !decoder.contains_key(&id) { + return Err(TokenizerError::MissingTokenId(id)); + } + } + + let byte_encoder = bytes_to_unicode(); + for &symbol in &byte_encoder { + let symbol = symbol.to_string(); + if !vocab.contains_key(&symbol) { + return Err(TokenizerError::MissingVocabSymbol(symbol)); + } + } + + let mut bpe_ranks = HashMap::new(); + for (rank, line) in merges_raw + .lines() + .map(str::trim) + .filter(|l| !l.is_empty() && !l.starts_with("#version:")) + .enumerate() + { + let mut it = line.split_whitespace(); + if let (Some(a), Some(b), None) = (it.next(), it.next(), it.next()) { + for operand in [a, b] { + if !vocab.contains_key(operand) { + return Err(TokenizerError::MissingVocabSymbol(operand.to_string())); + } + } + let merged = format!("{a}{b}"); + if !vocab.contains_key(&merged) { + return Err(TokenizerError::MissingVocabSymbol(merged)); + } + if bpe_ranks + .insert((a.to_string(), b.to_string()), rank) + .is_some() + { + return Err(TokenizerError::DuplicateMerge(line.to_string())); + } + } else { + return Err(TokenizerError::InvalidMergeLine(line.to_string())); + } + } + + let byte_decoder = byte_encoder + .iter() + .enumerate() + .map(|(b, &c)| (c, b as u8)) + .collect(); + + Ok(Self { + encoder: vocab, + decoder, + bpe_ranks, + byte_encoder, + byte_decoder, + }) + } + + pub fn vocab_size(&self) -> usize { + self.encoder.len() + } + + /// Encodes text into token ids. + pub fn encode(&self, text: &str) -> Vec { + let mut ids = Vec::new(); + for chunk in pre_tokenize(text) { + // Map each UTF-8 byte of the chunk into the unicode alphabet. + let mapped: String = chunk + .bytes() + .map(|b| self.byte_encoder[b as usize]) + .collect(); + for sym in self.bpe(&mapped) { + if let Some(&id) = self.encoder.get(&sym) { + ids.push(id); + } + } + } + ids + } + + /// Decodes token IDs, returning an error instead of skipping unknown IDs. + pub fn try_decode(&self, ids: &[usize]) -> Result { + Ok(String::from_utf8_lossy(&self.try_decode_bytes(ids)?).into_owned()) + } + + /// Decodes raw bytes while checking every token ID. + pub fn try_decode_bytes(&self, ids: &[usize]) -> Result, TokenizerError> { + let mut bytes = Vec::new(); + for id in ids { + let token = self + .decoder + .get(id) + .ok_or(TokenizerError::UnknownTokenId(*id))?; + bytes.extend( + token + .chars() + .filter_map(|c| self.byte_decoder.get(&c).copied()), + ); + } + Ok(bytes) + } + + /// Decodes token ids back into text. + pub fn decode(&self, ids: &[usize]) -> String { + String::from_utf8_lossy(&self.decode_bytes(ids)).into_owned() + } + + /// Decodes raw bytes, allowing streaming callers to retain partial UTF-8 sequences. + pub fn decode_bytes(&self, ids: &[usize]) -> Vec { + let mapped: String = ids + .iter() + .filter_map(|id| self.decoder.get(id)) + .flat_map(|s| s.chars()) + .collect(); + mapped + .chars() + .filter_map(|c| self.byte_decoder.get(&c).copied()) + .collect() + } + + /// Applies BPE merges to one pre-tokenized, byte-mapped chunk. + fn bpe(&self, token: &str) -> Vec { + let mut word: Vec = token.chars().map(|c| c.to_string()).collect(); + if word.len() < 2 { + return word; + } + loop { + // Find the adjacent pair with the lowest merge rank. + let mut best: Option<(usize, (String, String))> = None; + for i in 0..word.len() - 1 { + let pair = (word[i].clone(), word[i + 1].clone()); + if let Some(&rank) = self.bpe_ranks.get(&pair) { + if best.as_ref().map_or(true, |(r, _)| rank < *r) { + best = Some((rank, pair)); + } + } + } + let Some((_, pair)) = best else { break }; + // Merge every non-overlapping occurrence of that pair. + let merged = format!("{}{}", pair.0, pair.1); + let mut next = Vec::with_capacity(word.len()); + let mut i = 0; + while i < word.len() { + if i + 1 < word.len() && word[i] == pair.0 && word[i + 1] == pair.1 { + next.push(merged.clone()); + i += 2; + } else { + next.push(word[i].clone()); + i += 1; + } + } + word = next; + } + word + } +} + +/// Splits text into GPT-2-style chunks (contractions, words with an optional +/// leading space, number runs, punctuation runs, and whitespace runs). This is +/// a hand-rolled equivalent of GPT-2's regex that covers ordinary prose. +fn pre_tokenize(text: &str) -> Vec { + let chars: Vec = text.chars().collect(); + let n = chars.len(); + let mut out = Vec::new(); + let mut i = 0; + while i < n { + let c = chars[i]; + + // Contractions: 's 't 're 've 'm 'll 'd + if c == '\'' && i + 1 < n { + let two: String = chars[i + 1..n.min(i + 3)].iter().collect(); + if ["re", "ve", "ll"].contains(&two.as_str()) { + out.push(format!("'{two}")); + i += 3; + continue; + } + let one = chars[i + 1]; + if matches!(one, 's' | 't' | 'm' | 'd') { + out.push(format!("'{one}")); + i += 2; + continue; + } + } + + // Optional single leading space attached to the following word/number/punct. + let lead_space = c == ' ' && i + 1 < n && !chars[i + 1].is_whitespace(); + let start = i; + let k = if lead_space { i + 1 } else { i }; + if k < n && !chars[k].is_whitespace() { + let cat = category(chars[k]); + let mut e = k + 1; + while e < n && category(chars[e]) == cat && !chars[e].is_whitespace() { + e += 1; + } + out.push(chars[start..e].iter().collect()); + i = e; + continue; + } + + // Whitespace run. + if c.is_whitespace() { + let mut e = i; + while e < n && chars[e].is_whitespace() { + e += 1; + } + if e < n && e - i > 1 { + e -= 1; + } + out.push(chars[i..e].iter().collect()); + i = e; + continue; + } + + out.push(c.to_string()); + i += 1; + } + out +} + +#[derive(PartialEq, Eq)] +enum Cat { + Letter, + Digit, + Other, +} + +fn category(c: char) -> Cat { + if c.is_alphabetic() { + Cat::Letter + } else if c.is_numeric() { + Cat::Digit + } else { + Cat::Other + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn maybe_tokenizer() -> Option { + let v = Path::new("models/gpt2/vocab.json"); + let m = Path::new("models/gpt2/merges.txt"); + if v.exists() && m.exists() { + Tokenizer::from_files(v, m).ok() + } else { + None + } + } + + #[test] + fn encodes_known_example() { + let Some(tok) = maybe_tokenizer() else { + eprintln!("skipping: GPT-2 tokenizer files not present"); + return; + }; + // Canonical GPT-2 encoding. + assert_eq!(tok.encode("Hello world"), vec![15496, 995]); + } + + #[test] + fn roundtrips() { + let Some(tok) = maybe_tokenizer() else { return }; + for s in [ + "The quick brown fox.", + "batch_forge runs GPT-2!", + "I'm here", + ] { + assert_eq!(tok.decode(&tok.encode(s)), s); + } + } + + fn byte_vocab() -> HashMap { + bytes_to_unicode() + .iter() + .enumerate() + .map(|(id, symbol)| (symbol.to_string(), id)) + .collect() + } + + #[test] + fn rejects_duplicate_token_ids() { + let vocab = HashMap::from([("first".to_string(), 0), ("second".to_string(), 0)]); + assert!(matches!( + Tokenizer::from_assets(vocab, ""), + Err(TokenizerError::DuplicateTokenId(0)) + )); + } + + #[test] + fn rejects_sparse_token_ids() { + let vocab = HashMap::from([("first".to_string(), 0), ("third".to_string(), 2)]); + assert!(matches!( + Tokenizer::from_assets(vocab, ""), + Err(TokenizerError::MissingTokenId(1)) + )); + } + + #[test] + fn rejects_vocab_missing_a_byte_symbol() { + let vocab = HashMap::from([("a".to_string(), 0)]); + assert!(matches!( + Tokenizer::from_assets(vocab, ""), + Err(TokenizerError::MissingVocabSymbol(_)) + )); + } + + #[test] + fn rejects_merge_results_missing_from_vocab() { + let vocab = byte_vocab(); + assert!(matches!( + Tokenizer::from_assets(vocab, "#version: 0.2\na b\n"), + Err(TokenizerError::MissingVocabSymbol(symbol)) if symbol == "ab" + )); + } + + #[test] + fn valid_byte_vocab_encodes_and_decodes_a_merge() { + let mut vocab = byte_vocab(); + vocab.insert("ab".to_string(), 256); + let tokenizer = Tokenizer::from_assets(vocab, "#version: 0.2\na b\n").unwrap(); + assert_eq!(tokenizer.encode("ab"), vec![256]); + assert_eq!(tokenizer.decode(&[256]), "ab"); + } + + #[test] + fn whitespace_runs_match_gpt2_lookahead() { + assert_eq!(pre_tokenize("a b"), vec!["a", " ", " b"]); + assert_eq!(pre_tokenize("a b"), vec!["a", " ", " b"]); + assert_eq!(pre_tokenize("a "), vec!["a", " "]); + assert_eq!(pre_tokenize("a\n b"), vec!["a", "\n", " b"]); + } + + #[test] + fn hash_prefix_merges_are_not_comments() { + let mut vocab = byte_vocab(); + vocab.insert("##".into(), 256); + let tokenizer = Tokenizer::from_assets(vocab, "#version: 0.2\n\n# #\n").unwrap(); + assert_eq!(tokenizer.encode("##"), vec![256]); + } + + #[test] + fn malformed_merge_lines_are_reported() { + for line in ["a", "a b c"] { + assert!(matches!( + Tokenizer::from_assets(byte_vocab(), line), + Err(TokenizerError::InvalidMergeLine(_)) + )); + } + } + + #[test] + fn duplicate_merges_are_rejected() { + let mut vocab = byte_vocab(); + vocab.insert("ab".into(), 256); + assert!(matches!( + Tokenizer::from_assets(vocab, "a b\na b\n"), + Err(TokenizerError::DuplicateMerge(_)) + )); + } + + #[test] + fn missing_merge_operands_are_reported() { + let mut vocab = byte_vocab(); + vocab.insert("not_presenta".into(), 256); + assert!( + matches!(Tokenizer::from_assets(vocab, "not_present a"), Err(TokenizerError::MissingVocabSymbol(symbol)) if symbol == "not_present") + ); + } + + #[test] + fn checked_decode_reports_unknown_ids_and_preserves_bytes() { + let tokenizer = Tokenizer::from_assets(byte_vocab(), "").unwrap(); + assert!(matches!( + tokenizer.try_decode(&[999]), + Err(TokenizerError::UnknownTokenId(999)) + )); + assert_eq!( + tokenizer.try_decode_bytes(&[0xe2, 0x82, 0xac]).unwrap(), + vec![0xe2, 0x82, 0xac] + ); + assert_eq!(tokenizer.try_decode(&[0xe2, 0x82, 0xac]).unwrap(), "€"); + } +} diff --git a/tests/gpt2_e2e.rs b/tests/gpt2_e2e.rs new file mode 100644 index 0000000..563b703 --- /dev/null +++ b/tests/gpt2_e2e.rs @@ -0,0 +1,59 @@ +//! End-to-end GPT-2 verification. +//! +//! Requires the real `gpt2` weights under `models/gpt2/` (gitignored), so these +//! tests skip cleanly when the weights are absent (e.g. in CI). Locally they +//! assert the Rust forward reproduces HuggingFace GPT-2's prediction and that +//! the CPU and Metal backends agree. + +use std::path::Path; + +use batch_forge::gpt2::{Config, Gpt2}; +use batch_forge::loader; +use batch_forge::model::CpuBackend; + +const MODEL: &str = "models/gpt2/model.safetensors"; +// "The meaning of life is" +const TOKENS: [usize; 5] = [464, 3616, 286, 1204, 318]; +// Greedy next-token ranking from HuggingFace GPT-2 (verified by python/gpt2_reference.py). +const EXPECTED_TOP5: [usize; 5] = [407, 284, 262, 326, 257]; + +fn load() -> Option { + if !Path::new(MODEL).exists() { + eprintln!("skipping: {MODEL} not present"); + return None; + } + let tensors = loader::load_safetensors(Path::new(MODEL)).expect("load weights"); + Some(Gpt2::from_tensors(tensors, Config::default()).expect("build model")) +} + +fn top5(logits: &[f32]) -> Vec { + let mut idx: Vec = (0..logits.len()).collect(); + idx.sort_by(|&a, &b| logits[b].partial_cmp(&logits[a]).unwrap()); + idx.truncate(5); + idx +} + +#[test] +fn cpu_forward_matches_huggingface() { + let Some(model) = load() else { return }; + let logits = model.forward(&CpuBackend, &TOKENS); + assert_eq!(top5(&logits), EXPECTED_TOP5, "GPT-2 prediction diverged"); +} + +#[cfg(target_os = "macos")] +#[test] +fn cpu_metal_logits_agree() { + let Some(model) = load() else { return }; + let metal = batch_forge::metal_backend::MetalBackend::new(batch_forge::SHADER_SOURCE) + .expect("metal init"); + let cpu = model.forward(&CpuBackend, &TOKENS); + let gpu = model.forward(&metal, &TOKENS); + let max_diff = cpu + .iter() + .zip(&gpu) + .map(|(a, b)| (a - b).abs()) + .fold(0.0f32, f32::max); + eprintln!("[gpt2] CPU vs Metal logits max|Δ| = {max_diff:.3e}"); + assert!(max_diff < 1e-2, "CPU/Metal logits diverged: {max_diff}"); + assert_eq!(top5(&gpu), EXPECTED_TOP5); +} diff --git a/tests/parity.rs b/tests/parity.rs new file mode 100644 index 0000000..6ccc312 --- /dev/null +++ b/tests/parity.rs @@ -0,0 +1,296 @@ +//! CPU-reference vs Metal parity tests. +//! +//! Every Metal kernel is checked against the pure-Rust reference in +//! `batch_forge::ops` on randomized inputs. These run on Apple Silicon only +//! (the kernels need a Metal device) and back the tolerance bounds documented +//! in `docs/correctness.md`. Run with `cargo test --test parity -- --nocapture` +//! to print the observed maximum absolute deviation for each op. +#![cfg(target_os = "macos")] + +use batch_forge::metal_backend::MetalBackend; +use batch_forge::{ops, SHADER_SOURCE}; + +/// Tiny deterministic xorshift RNG so tests are reproducible without a dep. +struct Rng(u64); +impl Rng { + fn new(seed: u64) -> Self { + Rng(seed | 1) + } + fn next_u64(&mut self) -> u64 { + let mut x = self.0; + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + self.0 = x; + x + } + /// Uniform f32 in [-1, 1). + fn f32(&mut self) -> f32 { + ((self.next_u64() >> 40) as f32 / (1u64 << 24) as f32) * 2.0 - 1.0 + } + fn i8(&mut self) -> i8 { + ((self.next_u64() % 255) as i32 - 127) as i8 + } +} + +fn vecf(rng: &mut Rng, n: usize) -> Vec { + (0..n).map(|_| rng.f32()).collect() +} + +fn max_diff(a: &[f32], b: &[f32]) -> f32 { + assert_eq!(a.len(), b.len(), "length mismatch"); + a.iter() + .zip(b) + .map(|(x, y)| { + let diff = (x - y).abs(); + if diff.is_finite() { + diff + } else { + f32::INFINITY + } + }) + .fold(0.0f32, f32::max) +} + +#[test] +fn max_diff_rejects_non_finite_values() { + for value in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] { + assert_eq!(max_diff(&[value], &[0.0]), f32::INFINITY); + assert_eq!(max_diff(&[value], &[value]), f32::INFINITY); + } +} + +fn backend() -> MetalBackend { + MetalBackend::new(SHADER_SOURCE).expect("Metal init failed") +} + +/// Asserts parity and prints the observed deviation (visible with --nocapture). +fn check(name: &str, cpu: &[f32], gpu: &[f32], tol: f32) { + let d = max_diff(cpu, gpu); + eprintln!("[parity] {name:<22} max|Δ| = {d:.3e} (tol {tol:.0e})"); + assert!(d <= tol, "{name}: max diff {d:.3e} exceeds tol {tol:.0e}"); +} + +#[test] +fn matmul_parity() { + let (m, k, n) = (32, 64, 48); + let mut rng = Rng::new(1); + let a = vecf(&mut rng, m * k); + let b = vecf(&mut rng, k * n); + let cpu = ops::matmul(&a, &b, m, k, n); + let gpu = backend().matmul(&a, &b, m, k, n); + check("matmul", &cpu, &gpu, 1e-3); +} + +#[test] +fn linear_parity() { + let (rows, in_f, out_f) = (16, 64, 40); + let mut rng = Rng::new(2); + let x = vecf(&mut rng, rows * in_f); + let w = vecf(&mut rng, out_f * in_f); + let b = vecf(&mut rng, out_f); + let cpu = ops::linear(&x, &w, &b, rows, in_f, out_f); + let gpu = backend().linear(&x, &w, &b, rows, in_f, out_f); + check("linear", &cpu, &gpu, 1e-3); +} + +#[test] +fn gelu_parity() { + let mut rng = Rng::new(3); + let x = vecf(&mut rng, 4096); + let cpu: Vec = x.iter().map(|&v| ops::gelu(v)).collect(); + let gpu = backend().gelu(&x); + check("gelu", &cpu, &gpu, 1e-5); +} + +#[test] +fn layernorm_parity() { + let (rows, d) = (16, 64); + let mut rng = Rng::new(4); + let x = vecf(&mut rng, rows * d); + let gamma = vecf(&mut rng, d); + let beta = vecf(&mut rng, d); + let cpu = ops::layernorm(&x, &gamma, &beta, rows, d, 1e-5); + let gpu = backend().layernorm(&x, &gamma, &beta, rows, d, 1e-5); + check("layernorm", &cpu, &gpu, 1e-4); +} + +#[test] +fn rmsnorm_parity() { + let (rows, d) = (16, 64); + let mut rng = Rng::new(5); + let x = vecf(&mut rng, rows * d); + let gamma = vecf(&mut rng, d); + let cpu = ops::rmsnorm(&x, &gamma, rows, d, 1e-5); + let gpu = backend().rmsnorm(&x, &gamma, rows, d, 1e-5); + check("rmsnorm", &cpu, &gpu, 1e-4); +} + +#[test] +fn rope_parity() { + let (rows, d) = (8, 64); + let mut rng = Rng::new(6); + let mut cpu = vecf(&mut rng, rows * d); + let gpu_in = cpu.clone(); + let positions: Vec = (0..rows).collect(); + let positions_u32: Vec = positions.iter().map(|&p| p as u32).collect(); + ops::rope_inplace(&mut cpu, &positions, rows, d, 10000.0); + let gpu = backend().rope(&gpu_in, &positions_u32, rows, d, 10000.0); + check("rope", &cpu, &gpu, 1e-3); +} + +#[test] +fn attention_parity_noncausal() { + let (m, seq, d) = (4, 12, 32); + let mut rng = Rng::new(7); + let q = vecf(&mut rng, m * d); + let k = vecf(&mut rng, seq * d); + let v = vecf(&mut rng, seq * d); + let cpu = ops::attention(&q, &k, &v, m, seq, d, false, 0); + let gpu = backend().attention(&q, &k, &v, m, seq, d, false, 0); + check("attention(full)", &cpu, &gpu, 1e-3); +} + +#[test] +fn attention_parity_causal() { + let (m, seq, d, q_offset) = (4, 10, 32, 3); + let mut rng = Rng::new(8); + let q = vecf(&mut rng, m * d); + let k = vecf(&mut rng, seq * d); + let v = vecf(&mut rng, seq * d); + let cpu = ops::attention(&q, &k, &v, m, seq, d, true, q_offset); + let gpu = backend().attention(&q, &k, &v, m, seq, d, true, q_offset); + check("attention(causal)", &cpu, &gpu, 1e-3); +} + +#[test] +fn quant_matmul_parity() { + let (m, k, n) = (24, 48, 32); + let mut rng = Rng::new(9); + let a_i8: Vec = (0..m * k).map(|_| rng.i8()).collect(); + let scales: Vec = (0..m).map(|_| 0.01 + rng.f32().abs() * 0.05).collect(); + let b = vecf(&mut rng, k * n); + // Reference: dequantize then dense matmul. + let deq = ops::dequantize_int8(&a_i8, &scales, m, k); + let cpu = ops::matmul(&deq, &b, m, k, n); + let gpu = backend().quant_matmul(&a_i8, &scales, &b, m, k, n); + check("quant_matmul", &cpu, &gpu, 1e-3); +} + +#[test] +fn update_kv_cache_writes_correct_rows() { + let be = backend(); + let (d, max_len) = (4usize, 8usize); + let k_cache = be.create_buffer(&vec![0.0f32; max_len * d]).unwrap(); + let v_cache = be.create_buffer(&vec![0.0f32; max_len * d]).unwrap(); + + let k1: Vec = (0..2 * d).map(|i| i as f32).collect(); + let v1: Vec = (0..2 * d).map(|i| (i + 100) as f32).collect(); + let bk1 = be.create_buffer(&k1).unwrap(); + let bv1 = be.create_buffer(&v1).unwrap(); + be.update_kv_cache(&bk1, &bv1, &k_cache, &v_cache, 2, 0, d); + + let k2: Vec = (0..3 * d).map(|i| (i + 1000) as f32).collect(); + let v2: Vec = (0..3 * d).map(|i| (i + 2000) as f32).collect(); + let bk2 = be.create_buffer(&k2).unwrap(); + let bv2 = be.create_buffer(&v2).unwrap(); + be.update_kv_cache(&bk2, &bv2, &k_cache, &v_cache, 3, 2, d); + + let k_read: Vec = be.read_buffer(&k_cache, 5 * d); + let v_read: Vec = be.read_buffer(&v_cache, 5 * d); + assert_eq!(&k_read[..2 * d], &k1[..]); + assert_eq!(&k_read[2 * d..5 * d], &k2[..]); + assert_eq!(&v_read[..2 * d], &v1[..]); + assert_eq!(&v_read[2 * d..5 * d], &v2[..]); +} + +/// End-to-end check of the cached generation path: build the cache with +/// `update_kv_cache`, attend with `kv_attention` reading straight from cache +/// buffers, and compare to the reference attention over the assembled K/V. +#[test] +fn cached_attention_matches_reference() { + let be = backend(); + let (d, seq) = (16usize, 6usize); + let mut rng = Rng::new(10); + let k_full = vecf(&mut rng, seq * d); + let v_full = vecf(&mut rng, seq * d); + let q = vecf(&mut rng, d); // single query (m = 1) + + let k_cache = be.create_buffer(&vec![0.0f32; seq * d]).unwrap(); + let v_cache = be.create_buffer(&vec![0.0f32; seq * d]).unwrap(); + let bk = be.create_buffer(&k_full).unwrap(); + let bv = be.create_buffer(&v_full).unwrap(); + be.update_kv_cache(&bk, &bv, &k_cache, &v_cache, seq, 0, d); + + let bq = be.create_buffer(&q).unwrap(); + let bo = be.create_buffer_uninitialized::(d).unwrap(); + be.kv_attention(&bq, &k_cache, &v_cache, &bo, 1, seq, d, false, 0); + let gpu: Vec = be.read_buffer(&bo, d); + + let cpu = ops::attention(&q, &k_full, &v_full, 1, seq, d, false, 0); + check("cached_attention", &cpu, &gpu, 1e-3); +} + +#[test] +fn matmul_tiled_parity() { + let (m, k, n) = (40, 72, 56); // deliberately non-multiples of the 16 tile + let mut rng = Rng::new(11); + let a = vecf(&mut rng, m * k); + let b = vecf(&mut rng, k * n); + let cpu = ops::matmul(&a, &b, m, k, n); + let gpu = backend().matmul_tiled(&a, &b, m, k, n); + check("matmul_tiled", &cpu, &gpu, 1e-3); +} + +#[test] +fn gelu_scaled_parity() { + let be = backend(); + let mut rng = Rng::new(123); + // GPT-2 MLP hidden size (5 x 3072) with realistic magnitudes. + let x: Vec = (0..5 * 3072).map(|_| rng.f32() * 12.0).collect(); + let cpu: Vec = x.iter().map(|&v| ops::gelu(v)).collect(); + let gpu = be.gelu(&x); + let nan = gpu.iter().filter(|v| !v.is_finite()).count(); + eprintln!("[parity] gelu_scaled nan count = {nan}"); + check("gelu_scaled", &cpu, &gpu, 1e-2); +} + +#[test] +fn gpt2_sizes_parity() { + let be = backend(); + let mut rng = Rng::new(99); + // c_attn-shaped matmul + let a = vecf(&mut rng, 5 * 768); + let b = vecf(&mut rng, 768 * 2304); + let cpu = ops::matmul(&a, &b, 5, 768, 2304); + let gpu = be.matmul_tiled(&a, &b, 5, 768, 2304); + check("mm 5x768x2304", &cpu, &gpu, 3e-3); + // GPT-2 MHA shape + let hd = 12 * 64; + let q = vecf(&mut rng, 5 * hd); + let k = vecf(&mut rng, 5 * hd); + let v = vecf(&mut rng, 5 * hd); + let cpu = ops::mha(&q, &k, &v, 5, 12, 64); + let gpu = be.mha(&q, &k, &v, 5, 12, 64); + check("mha 5x12x64", &cpu, &gpu, 2e-3); + // tied LM head shape + let x = vecf(&mut rng, 768); + let w = vecf(&mut rng, 50257 * 768); + let bz = vec![0.0f32; 50257]; + let cpu = ops::linear(&x, &w, &bz, 1, 768, 50257); + let gpu = be.linear(&x, &w, &bz, 1, 768, 50257); + check("linear lm_head", &cpu, &gpu, 4e-3); +} + +#[test] +fn mha_parity() { + let (seq, heads, head_dim) = (7, 3, 8); + let hd = heads * head_dim; + let mut rng = Rng::new(12); + let q = vecf(&mut rng, seq * hd); + let k = vecf(&mut rng, seq * hd); + let v = vecf(&mut rng, seq * hd); + let cpu = ops::mha(&q, &k, &v, seq, heads, head_dim); + let gpu = backend().mha(&q, &k, &v, seq, heads, head_dim); + check("mha", &cpu, &gpu, 1e-3); +}