feat: rust port

[starpit]

This includes, under rust/

• a port to Rust
• cublas acceleration
• metal acceleration
• an KTIR-to-KTIR optimizer that covers a few important cases
• 4 test fixtures that have KTIR/MLIR for full forwards for the following scenarios:
  • smollm-135m decode
  • smollm-135m prefill m=8
  • llama-3.2-1b decode
  • llama-3.2-1b prefill m=32
• goldens verified vs transformers on these fixtures
• a direct Python↔Rust differential conformance harness: the example programs run
  through both the Python reference interpreter and the Rust port and are diffed
  head-to-head — bit-exact (max-abs 0) on the CPU/AMX interpreter across 19 programs,
  and tolerance-banded on the Metal fast path (NAX/simdgroup GEMM + fused map, each
  with a mandatory offload proof so a silent CPU fallback fails). Gated per-push in CI.
• unit and also e2e tests (the latter run benches and goldens validation)

The rust/PERFORMANCE.md file describes the massive gains possible on an M5, compared to the current Python interpreter. on llama 3.2 1b, 500x on decode, 1000x on prefill.

Seperately, I can describe how one can generate other test fixtures.

[lchu6]

Here is my view of this effort:

1. For ktir-cpu, correctness/reference behavior matters more than raw speed (if raw speed even matter), so this repo benefit less/much-less than some other torch-spyre repos.
2. It is still a good addition - it is 100% isolated so I am not scared to adopt it. and it can be useful for big validation workloads and future cross-repo reuse.
3. My real concern is maintenance and semantic drift. How can we make sure this is guarded by Python parity, and who is the owner for that.

@fabianlim wdyt?

[starpit]

Fwiw, all of the Python tests are ported.

Performance comparo — whole-model forward, ms/pass (Apple M5, f16; one warm-up excluded). Rust = the resident GPU executor (production default: prefill computes only the last token's logits, as generation does); Python = the per-node reference interpreter, with the #124 perf fixes now merged.

config	Rust (ms/pass)	Python (ms/pass)	speedup
smollm2-135m decode	17.6	2,536.3	144×
smollm2-135m prefill (m=8)	45.6	60,397.4	1324×
llama-3.2-1b decode	71.6	31,777.6	444×
llama-3.2-1b prefill (m=32)	169.1	977,688.3	5782×

Python is the per-node reference interpreter after the now-merged #124 (vectorized ktdp.load offsets, ravel-not-flatten allocation reads, O(log n) allocation lookup) — so these are the de-gratuitized numbers; the gap is still 144–5782×. llama-3.2-1b prefill needs a raised LX cap (KTIR_LX_MB) to run under the faithful 2 MB LX per-node reference.

(Prefill numbers are the production last-token default — only the final position's logits are computed, the same token generation samples, validated exact against the real model. The all-rows prefill path (every position, used by the comprehensive golden) is llama 186.6 / smollm2 45.8 ms. Python numbers unchanged.)

Update (fused m=1 attention): decode attention had been running as a ~1500-op interpreter storm per layer (the H heads unrolled in the node body — the BLAS GEMVs are near-free; the cost was the const/splat/access_tile plumbing around them). A structural m=1 attention recognizer + a fused CPU dispatch — per head, QKᵀ GEMV → softmax → scores·V GEMV (f32 accumulate, GQA + per-position context mask + 1/√d scale + context/diagonal split) — collapses it to ~3·H BLAS+softmax primitives, golden-faithful (max-abs identical to the decomposed oracle). Resident decode: llama 97.4→71.6 (1.35×), smollm2 33.9→17.6 (2.05×) on the M5. Because it removes ~20k op-dispatches/token of pure CPU plumbing, it wins more on slower-CPU devices: real vLLM llama-3.2-1b decode on an M1 Max went 5.4→11 tok/s — clearing the CPU bottleneck lets the M1 Max's ~2.5× memory bandwidth win the (now-dominant) weight-streaming GEMMs and pull ahead of the M5. KTIR_NO_FUSE_ATTN restores the decomposed path. Prefill is unchanged (the fused path is m=1-only).

Update (NAX matmul kernel): the prefill numbers reflect a vectorized matmul2d threadgroup loader — wide 4-element coalesced device loads + threadgroup stores replacing the per-element div/mod + per-element bounds checks in the staging path (the GEMMs were loader-bound, streaming weights at only ~12 GB/s). llama, with the largest GEMMs, gains most (prefill 196.6→169.1, −14%); smollm2's tiny m=8 GEMMs are in the noise. Bit-identical output (golden unchanged). The NAX kernels are also now AOT-precompiled (embedded metallibs, with the mandatory -mmacosx-version-min=26.2 workaround for the SDK-26.5 matmul2d half-K miscompile, JIT fallback retained) — but that's startup-only and does not affect these steady-state numbers.

Note: this single-pass ms/pass bench marshals sources once, so it does not capture the biggest serving-loop fix — a resident-weight-caching bug where set_sources re-decoded + re-uploaded the whole ~2 GB weight set every decode token in the real serving loop (vLLM calls set_sources per token to thread the KV cache). Keeping only the provably-immutable weights resident took real vLLM decode from 1.7 → 4.8 tok/s (~88% of weights now stay resident across steps), and the fused m=1 attention above took it further. The table numbers are steady-state with weights already resident.

[fabianlim]

@starpit — before we proceed with merging, I'd like to suggest we align on a landing plan first. I've opened issue #128 which documents the proposed commit sequence and flags a few items that need to be resolved (CI fixes, the examples/ compile-time dependency, sparse checkout ergonomics for Python contributors, and some bench scripts with hardcoded local paths).

Can you take a look at #128 and let us know if the proposed approach works for you, or if there's a different split you had in mind?

cc @kiszk @lchu6 @nwang-ibm

[fabianlim]

@starpit as written in #128, I will need to insist to not have the tests triggered only on rust paths. My reason is very strong, I simply do not want to put burden on ktir-cpu developers, to fix conformance discrepancies between ktir-cpu and rust. If we had this, then everytime we update ktir-cpu such that it goes out of sync with rust, then the CI will break, and we cannot merge our PRs until we reconcile the differences

My suggestion was that we have a workflow-dispatch that we can trigger an adhoc CI-run on main or a particular commit to test for breakages.

cc: @mudhakar @raghukiran1224 @lchu6

[fabianlim]

torch-spyre is a spyre project. im not sure we want to run GPU Ci tests in our organization. Im concerned about billing, im not sure how GPU runners are charged to our organization.

That being said, from the above comment, im confused with the above comment, this runs on MacOS, so its running on an Apple silicon GPU?

Regardless, macos runners I believe are 10X more expensive than linux boxes

[fabianlim]

We need to resolve python developer burden and CI concerns before this can be merged. This projects main goal is to focus on correctness and latency modeling of ktir-cpu, which is evolving at a very fast pace. We are very appreciative for this contribution, but we cannot have the rust port in a critical CI path, which will impede the original goal of this repo. I hope you understand and try to accomodate all the concerns in #128 . I noticed alot of requests like i) staging the merge, and ii) introducing the sparse checkout (or figuring out an alternative if that does not work) is rejected.

Pls understand that since this project is high-velocity, im very cautious in introducing such a big change. This repo is used in evaluating our upstream triton repos and many others, the blast radius will be very big if something goes wrong.