Compare commits

...

68 Commits

Author SHA1 Message Date
42da25a37c feat(bench): version-aware benchmark harness + neuron build metadata
All checks were successful
CI / CUDA type-check (push) Successful in 1m36s
CI / Format (push) Successful in 31s
CI / Clippy (push) Successful in 2m47s
CI / Test (push) Successful in 4m33s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
Adds automated, longitudinal performance tracking across neuron builds,
replacing manual script/bench.py runs and hand edits to benchmarks.md.

neuron build metadata + GET /version:
- cortex-core: shared BuildInfo type (build_info.rs).
- neuron build.rs captures git SHA (preferring injected HELEXA_BUILD_SHA,
  else git, else "unknown"), dirty flag, build timestamp, rustc version,
  profile, target, enabled cargo features, and best-effort candle-core
  version from Cargo.lock.
- New GET /version endpoint (version.rs) + clap --version long form.
- SHA injected in CI (build-neuron step) and helexa-neuron.spec
  (%{?helexa_commit}) so tarball RPMs report the real SHA. /version is
  now the canonical "which build is live" probe.

helexa-bench crate:
- Continuous daemon: hits each neuron directly on :13131, exercises each
  warm (status==loaded) model, records every run into a SQLite
  system-of-record stamped with the neuron's full BuildInfo.
- Version-aware: skips any (target, build SHA, model, scenario) cell
  already at samples_per_version, so a steady fleet costs only cheap
  /version + /models polls until a new SHA ships.
- Extensible Scenario trait; phase-1 chat-latency family ported verbatim
  from bench.py (synthetic 128/4096-tok prompts, /no_think, streamed
  TTFT + decode-window tok/s). `report` regenerates the benchmarks table.
- kind="openai" comparison targets scaffolded, not yet wired.

Packaging: data/helexa-bench.service (+ sysusers), prebuilt-binary RPM
spec (outbound-only, no firewalld), and build/package/publish wiring in
build-prerelease.yml with change detection.

Tests: cortex-core BuildInfo round-trip, neuron GET /version integration,
helexa-bench unit (prompt/SSE/config/store) + end-to-end sweep
(record -> skip -> resume on new SHA). Docs updated (benchmarks.md,
CLAUDE.md addendum).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-13 15:26:02 +03:00
30d50d6215 Merge pull request 'fix(ci): drop the unused flash-attn feature from neuron builds (#42)' (#46) from fix/42-drop-flash-attn into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 33s
build-prerelease / Build neuron-blackwell (push) Successful in 1m28s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m43s
build-prerelease / Build cortex binary (push) Successful in 2m45s
build-prerelease / Test (push) Successful in 4m32s
build-prerelease / Package cortex RPM (push) Successful in 1m23s
build-prerelease / Build neuron-ampere (push) Successful in 2m5s
build-prerelease / Build neuron-ada (push) Successful in 3m29s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 1m39s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 1m40s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 1m41s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 53s
2026-06-13 07:15:43 +00:00
9a312098dd fix(ci): drop the unused flash-attn feature from neuron builds (#42)
All checks were successful
CI / Format (push) Successful in 37s
CI / CUDA type-check (push) Successful in 2m9s
CI / Clippy (push) Successful in 2m17s
CI / Test (push) Successful in 4m56s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Format (pull_request) Successful in 34s
CI / CUDA type-check (pull_request) Successful in 1m37s
CI / Clippy (pull_request) Successful in 2m31s
CI / Test (pull_request) Successful in 4m45s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
The neuron fleet builds with `cuda cudnn flash-attn`, but nothing in
neuron uses flash-attn: the qwen3_5 (27B) arch is hand-rolled, the
candle-transformers qwen3 model has no flash path, llama is built with
use_flash_attn=false, and `grep flash crates/neuron/src` is empty. The
feature only pulls in candle-flash-attn's sm_80/sm_86 CUDA kernel
sweep — which is exactly where ptxas SIGSEGVs/hangs in #42 (3 hits in
one day, the last a ~4-hour hang that stalled the whole deploy behind
the ampere job).

Dropping the feature removes the #42 failure surface at the root (not
a mitigation) and cuts the longest, most fragile part of each flavour
build. No runtime change — nothing called those kernels. Removed from
all three flavour builds in build-prerelease.yml and from deploy-dev.yml;
ci.yml's cuda-check already used `--features cuda` only.

Closes #42

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-13 09:43:14 +03:00
98e9749f22 Merge pull request 'feat(neuron): speculative decoding — acceptance core + config (#25, phase 1)' (#45) from feat/25-speculative-decoding into main
Some checks failed
build-prerelease / Build neuron-ampere (push) Blocked by required conditions
build-prerelease / Resolve version stamps + change detection (push) Successful in 32s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m11s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Test (push) Successful in 5m22s
build-prerelease / Build neuron-blackwell (push) Successful in 10m12s
build-prerelease / Build neuron-ada (push) Successful in 14m21s
build-prerelease / Package helexa-neuron-ada RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been cancelled
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been cancelled
2026-06-13 06:39:49 +00:00
ec764a2cac feat(neuron): speculative decoding — acceptance core + config (#25, phase 1)
All checks were successful
CI / Format (push) Successful in 36s
CI / Format (pull_request) Successful in 31s
CI / CUDA type-check (push) Successful in 1m50s
CI / CUDA type-check (pull_request) Successful in 1m44s
CI / Clippy (push) Successful in 2m38s
CI / Test (push) Successful in 4m21s
CI / Clippy (pull_request) Successful in 2m29s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Test (pull_request) Successful in 4m37s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
First phase of speculative decoding: the pure, state-free acceptance
logic and per-target config, unit-tested in isolation before the
draft/verify loop and GDN-state rollback wire it into the generation
path.

greedy_accept walks the drafter's K proposed tokens against the
target's greedy token at each of the K+1 positions, accepting the
longest matching prefix and always committing one bonus token on top
(the target's correction at the first mismatch, or a free extra token
when the whole draft matched). So a round commits 1..=K+1 tokens —
never zero, guaranteeing forward progress even with a useless drafter.
Greedy is exact for temperature-0 (the fleet probe + #22 bench
regime); stochastic acceptance is a later phase.

SpeculativeConfig carries the drafter id (must share the target's
tokenizer — Qwen3.5-0.8B for the Qwen3.6-27B target, both qwen3_5,
byte-identical tokenizer, confirmed on beast) and the draft length K.

6 unit tests: full accept, partial accept, zero accept (progress
guarantee), last-position mismatch, single-token draft, config
gating. Not yet wired into the decode path — phase 2 (single-GPU
draft/verify) follows. Design + phasing on the issue.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-13 08:30:21 +03:00
4c1bdba31d Merge pull request 'feat(neuron): chunk the single-GPU vision prefill (parity with TP) (#18)' (#44) from feat/18-single-gpu-vision-chunked into main
Some checks failed
build-prerelease / Resolve version stamps + change detection (push) Successful in 32s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m11s
build-prerelease / Test (push) Successful in 4m7s
build-prerelease / Build neuron-blackwell (push) Successful in 9m55s
build-prerelease / Build neuron-ada (push) Successful in 13m10s
build-prerelease / Build neuron-ampere (push) Has started running
build-prerelease / Package helexa-neuron-ada RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been cancelled
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been cancelled
2026-06-13 05:23:07 +00:00
988ef5afc2 feat(neuron): chunk the single-GPU vision prefill (parity with TP) (#18)
All checks were successful
CI / Format (push) Successful in 31s
CI / Format (pull_request) Successful in 41s
CI / CUDA type-check (pull_request) Successful in 1m27s
CI / CUDA type-check (push) Successful in 2m11s
CI / Clippy (push) Successful in 2m38s
CI / Clippy (pull_request) Successful in 2m31s
CI / Test (pull_request) Successful in 4m13s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Test (push) Successful in 4m39s
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
The single-GPU vision path was still single-shot: a long vision-bearing
prompt to a single-GPU-loaded qwen3_5 had the OOM exposure the TP path
shed in fa01350 (it was only guard-rejected, never served).

Mirror TpQwen3_5ForCausalLM::prefill_with_images_chunked onto the
single-GPU Qwen3_5ForCausalLM: encode the image(s) once, walk the
pre-expanded prompt in prefill_chunk_tokens() windows splicing the
per-chunk <|image_pad|> rows, accumulate KV + GDN state across chunks
via the growing offset, keep the last chunk's logits. Interleaved
M-RoPE positions are computed once over the whole prompt and sliced
per chunk (an image compresses the position space, so per-chunk offset
arithmetic would be wrong) — so Qwen3_5Model::forward_inner gains an
explicit position_ids path alongside the internal-from-grids
(single-shot) and plain (text/decode) paths, plus a forward_with_positions
entry point. The device-worker ForwardLogitsWithImages handler now
calls the chunked method; chunk size comes from prefill_chunk_tokens()
on the worker thread, so the Job/handle surface and the callers are
unchanged.

The shared validate_vision_prefill VRAM/KV backstop stays (TP keeps it
too) — chunking bounds activation memory, not the accumulating KV
cache, so the guard still does useful work.

Verified on real weights (Qwen3.5-0.8B): extended the #15 vision
reference test to also run the chunked path with chunk_size=64 over the
217-token prompt (4 chunks; the ~196-token image-pad run spans them).
Chunked vs single-shot logits: cosine 1.000000, max_abs 0.0001;
argmax matches the HF reference. The test covers all three
forward_inner branches (text plain / single-shot vision / chunked
vision) on a real single-GPU qwen3_5 load.

Closes #18

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-13 08:17:11 +03:00
a1450789d2 Merge pull request 'docs(learnings): source-control P1 + P2 sprint learnings' (#43) from docs/learnings-p1-p2 into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 35s
build-prerelease / Lint (fmt + clippy) (push) Has been skipped
build-prerelease / Build neuron-blackwell (push) Has been skipped
build-prerelease / Build neuron-ampere (push) Has been skipped
build-prerelease / Build neuron-ada (push) Has been skipped
build-prerelease / Package helexa-neuron-ada RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been skipped
build-prerelease / Test (push) Has been skipped
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been skipped
2026-06-13 04:21:11 +00:00
2eaa776d85 docs(learnings): source-control P1 + P2 sprint learnings
All checks were successful
CI / Format (push) Successful in 35s
CI / Format (pull_request) Successful in 35s
CI / CUDA type-check (push) Successful in 1m32s
CI / CUDA type-check (pull_request) Successful in 1m37s
CI / Clippy (push) Successful in 2m30s
CI / Clippy (pull_request) Successful in 2m31s
CI / Test (push) Successful in 4m24s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Test (pull_request) Successful in 4m27s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
doc/plan/* is gitignored, so the P1 learnings briefing could never be
committed. Move it to doc/learnings/p1.md (verbatim) and add
doc/learnings/p2.md capturing the P2 sprint (#11/#23/#1/#15).

The P2 doc's headline: CI green != correct. Four correctness bugs
passed every CI gate and surfaced only on the live fleet (post-gen
snapshots never re-match reasoning models; full-prompt snapshots
break on BPE retokenization; the chunked delta-rule's nilpotent-
squaring shortcut NaNs on correlated keys; the 0.8B masked two of
these by luck). Plus the device-worker/TP state patterns, the
deploy-dev + systemd-drop-in A/B loop, the per-package change-
detection fleet-split failure mode (#42), and the f32-fixture
numerical-validation rig (#15).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-13 07:13:36 +03:00
7918995e5a chore(ci): retrigger build-prerelease — ampere ptxas segfault (flash-attn sm_86, runner-side) on 538cc87
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 30s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m12s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Test (push) Successful in 3m57s
build-prerelease / Build neuron-blackwell (push) Successful in 9m5s
build-prerelease / Build neuron-ampere (push) Successful in 15m3s
build-prerelease / Build neuron-ada (push) Successful in 19m4s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 3m9s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 3m13s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 3m50s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 1m12s
2026-06-13 00:12:24 +03:00
538cc87572 Merge pull request 'feat(neuron): numerical validation against the transformers reference (#15)' (#41) from feat/15-numerical-reference into main
Some checks failed
build-prerelease / Resolve version stamps + change detection (push) Successful in 29s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m13s
build-prerelease / Build cortex binary (push) Successful in 2m43s
build-prerelease / Test (push) Successful in 4m25s
build-prerelease / Package cortex RPM (push) Successful in 1m39s
build-prerelease / Build neuron-ampere (push) Failing after 8m54s
build-prerelease / Build neuron-blackwell (push) Successful in 10m11s
build-prerelease / Build neuron-ada (push) Successful in 14m9s
build-prerelease / Package helexa-neuron-ada RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been skipped
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 56s
2026-06-12 20:43:37 +00:00
1c4b53cbf1 feat(neuron): numerical validation against the transformers reference (#15)
All checks were successful
CI / Format (push) Successful in 44s
CI / Format (pull_request) Successful in 37s
CI / CUDA type-check (push) Successful in 1m39s
CI / CUDA type-check (pull_request) Successful in 2m6s
CI / Clippy (push) Successful in 2m23s
CI / Clippy (pull_request) Successful in 2m24s
CI / Test (push) Successful in 4m21s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Test (pull_request) Successful in 4m1s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
script/dump_reference.py captures fixtures from the HF qwen3_5
implementation (token ids + reference tensors, f32 by default so the
comparison pins math rather than dtype noise);
tests/numerical_reference.rs replays them through our arch and
asserts argmax equality, cosine similarity, and max-abs ceilings. The
tests self-skip without NEURON_REF_MODEL_PATH so CI stays green
without weights.

Measured on beast (f32-vs-f32): text logits max_abs 0.000 / cosine
1.000000 (the >64-token prompt routes through the chunked GDN
prefill, so the production prefill math is what's validated); vision
tower cosine 0.999998, end-to-end vision logits cosine 1.000000 with
identical argmax. Mutation sensitivity: NEURON_VISION_LEGACY_POS=1
collapses tower cosine to 0.75 and fails loudly.

One production fidelity fix the harness surfaced: the pos-embed
bilinear blend now accumulates in f32 and casts once at the end,
matching the reference (we previously rounded the weights to bf16
before blending).

Fixtures: 0.8B text + vision (f32), 27B text (bf16 — an f32 27B
forward needs ~108 GB; the automated comparison runs against the
0.8B, which executes the same arch modules). Regeneration documented
in tests/fixtures/numerical/README.md.

Closes #15

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 23:35:57 +03:00
49a8dbcd28 Merge pull request 'perf(neuron): parallel in-situ quantization + cold-load phase timing (#1)' (#40) from perf/1-parallel-isq into main
Some checks are pending
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Blocked by required conditions
build-prerelease / Resolve version stamps + change detection (push) Successful in 30s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m11s
build-prerelease / Build cortex binary (push) Successful in 2m21s
build-prerelease / Test (push) Successful in 3m56s
build-prerelease / Build neuron-blackwell (push) Successful in 9m38s
build-prerelease / Build neuron-ada (push) Successful in 14m21s
build-prerelease / Build neuron-ampere (push) Successful in 19m0s
build-prerelease / Package cortex RPM (push) Successful in 1m21s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 3m24s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 3m12s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 4m29s
2026-06-12 20:12:44 +00:00
90e971dcf5 perf(neuron): parallel in-situ quantization + cold-load phase timing (#1)
All checks were successful
CI / Format (push) Successful in 32s
CI / Format (pull_request) Successful in 35s
CI / CUDA type-check (push) Successful in 1m50s
CI / CUDA type-check (pull_request) Successful in 2m7s
CI / Clippy (pull_request) Successful in 2m18s
CI / Clippy (push) Successful in 2m46s
CI / Test (push) Successful in 5m33s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Test (pull_request) Successful in 5m33s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
QTensor::quantize runs its per-block math strictly sequentially on
one core (CUDA storage round-trips through the same CPU path), which
made Q6K ISQ the dominant phase of the 27B TP cold load. Blocks are
independent, so quantize_parallel re-implements the same encoding
through candle's public per-block API (k_quants::GgmlType::from_float)
with rayon fanning blocks across the CPU pool — byte-identical output,
pinned by parity tests against QTensor::quantize for Q6K/Q5K/Q4K/Q8_0.

Threading discipline holds: the device-to-host read and the
QStorage::from_data upload stay on the calling thread (device worker /
subprocess main); rayon workers touch host memory only.

Also adds the per-phase timing the issue asked for first: per-layer
debug + layer-loop total + lm_head info lines, so the next cold load
shows where the time actually goes.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 22:47:57 +03:00
92273eb936 chore(ci): retrigger build-prerelease — ampere/blackwell packaging skipped after transient build failure on 128b381
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 29s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m10s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Test (push) Successful in 3m54s
build-prerelease / Build neuron-blackwell (push) Successful in 9m36s
build-prerelease / Build neuron-ada (push) Successful in 14m6s
build-prerelease / Build neuron-ampere (push) Successful in 19m8s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 3m8s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 3m8s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 3m52s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 1m5s
2026-06-12 22:38:31 +03:00
128b3818cb Merge pull request 'perf(neuron): chunked delta-rule prefill for Gated DeltaNet (#23)' (#39) from perf/23-chunked-gdn-prefill into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 30s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m13s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Test (push) Successful in 3m57s
build-prerelease / Build neuron-blackwell (push) Successful in 10m3s
build-prerelease / Build neuron-ada (push) Successful in 14m11s
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been skipped
build-prerelease / Build neuron-ampere (push) Successful in 14m3s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 3m10s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 58s
2026-06-12 18:44:22 +00:00
812d191e50 fix(neuron): UT transform by forward substitution, not nilpotent squaring
All checks were successful
CI / Format (push) Successful in 32s
CI / Format (pull_request) Successful in 53s
CI / CUDA type-check (push) Successful in 1m52s
CI / CUDA type-check (pull_request) Successful in 2m12s
CI / Clippy (push) Successful in 2m18s
CI / Clippy (pull_request) Successful in 2m36s
CI / Test (push) Successful in 4m18s
CI / Test (pull_request) Successful in 4m22s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
Live A/B on beast produced NaN logits ("!!!" replies) on real prompts:
the nilpotent-squaring form of (I - T)^-1 computes raw powers of T,
whose entries grow combinatorially (path counts ~ C(62,31)) before
nilpotency collapses them — fine on uncorrelated test data, f32
precision death on real prompts whose repetitive text makes keys
highly correlated. The reference's forward-substitution loop never
forms raw powers; its intermediates are the convergent M entries.

Port the reference loop faithfully (rows accumulate into a fresh
tensor). New adversarial parity test with near-identical keys and
beta ~= 1 diverges to 8e30 under the squaring form and passes under
forward substitution.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 21:18:32 +03:00
2a9def6d2d perf(neuron): chunked delta-rule prefill for Gated DeltaNet (#23)
All checks were successful
CI / Format (push) Successful in 32s
CI / Format (pull_request) Successful in 24s
CI / CUDA type-check (push) Successful in 1m38s
CI / CUDA type-check (pull_request) Successful in 2m10s
CI / Clippy (push) Successful in 2m34s
CI / Test (push) Successful in 4m20s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Clippy (pull_request) Successful in 2m29s
CI / Test (pull_request) Successful in 4m21s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
Prefill (seq_len >= 64) now runs the chunk-parallel gated delta rule
ported from the HF reference torch_chunk_gated_delta_rule
(chunk_size=64): identical math reorganised into per-chunk batched
matmuls (cuBLAS/tensor cores on CUDA, gemm on CPU) instead of the
O(L)-sequential per-token recurrence. Decode steps and short prompts
keep the recurrent paths (CUDA kernel / Rust loop) unchanged.

One deliberate deviation from the reference: its in-place row-by-row
UT-transform computes (I - T)^-1 - I by forward substitution; T is
strictly lower triangular and therefore nilpotent at chunk size 64,
so the same inverse is the product of six squarings
prod_{j=0..5}(I + T^(2^j)) — batched matmuls instead of 63 sequential
row updates, which suits candle's immutable tensors. Chunk-local math
runs rank-3 over a flattened B*H*N batch dim (candle matmul supports
at most two batch dims).

Initial-state continuation is supported, so chunked prefill composes
with #11's restored prefix snapshots. Both single-GPU and TP paths
pick this up through the shared run_delta_rule dispatch.
NEURON_GDN_CHUNKED=0 forces the recurrent paths for A/B measurement.

Parity tests pin chunked against recurrent (2e-4 abs) across padding
(L=130), exact multiples with non-zero initial state (L=128 after a
50-token prefix), and a single exact chunk.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 20:51:51 +03:00
ddb331e1a3 Merge pull request 'docs(bench): record post-#11 fleet numbers' (#38) from docs/benchmarks-post-11 into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 30s
build-prerelease / Build neuron-blackwell (push) Has been skipped
build-prerelease / Build neuron-ampere (push) Has been skipped
build-prerelease / Build neuron-ada (push) Has been skipped
build-prerelease / Package helexa-neuron-ada RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been skipped
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m47s
build-prerelease / Test (push) Successful in 4m25s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been skipped
2026-06-12 17:14:00 +00:00
df0bf4c518 docs(bench): record post-#11 fleet numbers
All checks were successful
CI / Format (push) Successful in 37s
CI / Format (pull_request) Successful in 37s
CI / CUDA type-check (push) Successful in 1m31s
CI / CUDA type-check (pull_request) Successful in 2m7s
CI / Clippy (push) Successful in 2m24s
CI / Test (push) Successful in 4m17s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Clippy (pull_request) Successful in 2m24s
CI / Test (pull_request) Successful in 3m57s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
Appends the 2026-06-12 post-prefix-cache run: 27B @4k warm TTFT
7.07 s -> 1.43 s, no-cache control models unchanged, with a
methodology note that repeated-prompt cells now measure warm TTFT on
qwen3_5-arch models.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 20:06:53 +03:00
a1952a4522 Merge pull request 'fix(neuron): snapshot at the last special-token boundary (#11)' (#37) from fix/11-snapshot-cut-retokenization into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 30s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m37s
build-prerelease / Test (push) Successful in 4m21s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Build neuron-blackwell (push) Successful in 10m22s
build-prerelease / Build neuron-ampere (push) Successful in 13m8s
build-prerelease / Build neuron-ada (push) Successful in 21m31s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 2m57s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 3m0s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 3m46s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 1m2s
2026-06-12 16:24:15 +00:00
4f266dbd82 fix(neuron): snapshot at the last special-token boundary (#11)
All checks were successful
CI / Format (push) Successful in 42s
CI / Format (pull_request) Successful in 34s
CI / CUDA type-check (push) Successful in 1m31s
CI / Clippy (push) Successful in 2m19s
CI / CUDA type-check (pull_request) Successful in 2m10s
CI / Test (push) Successful in 4m13s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Clippy (pull_request) Successful in 2m9s
CI / Test (pull_request) Successful in 4m5s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
Second finding from live 27B validation: prompt-covering snapshots
still never matched. The rendered prompt ends with
`<|im_start|>assistant\n`, and when the next turn re-tokenizes that
text followed by the assistant's reply, BPE merges the trailing
newline with the reply's first characters — the final token(s) of the
cached sequence differ from the next prompt's, so the exact-prefix
match never fires. (A reply starting with an atomic special token
like <think> masks this, which is why the 0.8B check passed.)

Snapshot one past the last <|im_start|> instead: special tokens are
hard segmentation points, so ids up to and including it are provably
identical across renders. Prefill pauses at that boundary to capture
the snapshot, then finishes the ~2-token `assistant\n` tail. Applied
to all six request paths; unit tests for the cut helper.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 19:16:45 +03:00
43a6d96d5f Merge pull request 'fix(neuron): snapshot prefix cache at the prefill boundary (#11)' (#36) from fix/11-prefix-snapshot-at-prefill into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 36s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m16s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Test (push) Successful in 4m2s
build-prerelease / Build neuron-ampere (push) Successful in 13m22s
build-prerelease / Build neuron-blackwell (push) Successful in 13m31s
build-prerelease / Build neuron-ada (push) Successful in 14m25s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 3m6s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 3m12s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 3m50s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 1m6s
2026-06-12 15:34:59 +00:00
3fd1989b2b fix(neuron): snapshot prefix cache at the prefill boundary (#11)
All checks were successful
CI / Format (push) Successful in 41s
CI / Format (pull_request) Successful in 42s
CI / CUDA type-check (push) Successful in 1m39s
CI / CUDA type-check (pull_request) Successful in 2m6s
CI / Clippy (push) Successful in 3m10s
CI / Clippy (pull_request) Successful in 3m3s
CI / Test (pull_request) Successful in 4m2s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
CI / Test (push) Successful in 5m1s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
Live validation on beast's Qwen3.6-27B showed reused=0 on every turn:
the post-generation snapshot includes reasoning tokens (<think>...)
that get stripped when the client echoes the assistant message back,
so the cached sequence is never a token-prefix of the next prompt.
quadbrat's 0.8B only matched because its think block round-tripped as
literal text.

Snapshot after prefill instead (covering exactly the prompt tokens) —
that is the state the next turn provably extends under a stable chat
template, regardless of how reasoning or tool-call content is
transformed on echo. Taken after the first healthy sample so
NaN-poisoned prefills never cache their state; this also retires the
forwarded-token bookkeeping and the consumer-hangup store sites.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 18:29:00 +03:00
f7952547e7 Merge pull request 'feat(neuron): prefix KV caching for the TP path (#11)' (#35) from feat/11-prefix-kv-cache-tp into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 31s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m12s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Test (push) Successful in 3m58s
build-prerelease / Build neuron-blackwell (push) Successful in 9m5s
build-prerelease / Build neuron-ada (push) Successful in 14m22s
build-prerelease / Build neuron-ampere (push) Successful in 19m0s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 2m56s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 3m0s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 3m51s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 1m5s
2026-06-12 14:49:19 +00:00
7e66f77851 fix(neuron): CUDA type-check fixes for TP prefix cache
All checks were successful
CI / Format (push) Successful in 38s
CI / Format (pull_request) Successful in 39s
CI / CUDA type-check (pull_request) Successful in 1m26s
CI / CUDA type-check (push) Successful in 1m34s
CI / Clippy (push) Successful in 3m14s
CI / Clippy (pull_request) Successful in 3m18s
CI / Test (push) Successful in 5m15s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Test (pull_request) Successful in 3m56s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
Two errors only the cuda config surfaces: the TpSnapshotKv dispatch
arms mixed candle and anyhow error types, and restore_or_clear_tp held
the registry MutexGuard across the cleanup await inside a let-chain
(making the TP request futures non-Send). Bind the removed ref before
awaiting, same discipline as the other lock sites.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 17:39:32 +03:00
e629e1872c feat(neuron): prefix KV caching for the TP path (#11)
Some checks failed
CI / Format (push) Successful in 37s
CI / Format (pull_request) Successful in 31s
CI / CUDA type-check (push) Failing after 1m55s
CI / CUDA type-check (pull_request) Failing after 1m47s
CI / Clippy (push) Successful in 2m11s
CI / Test (push) Successful in 4m15s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Clippy (pull_request) Successful in 2m23s
CI / Test (pull_request) Successful in 4m0s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
Extends the prefix cache to tensor-parallel models — Qwen3.6-27B on
beast, where the TTFT win is largest. Closes #11.

Every rank holds its shard's snapshot under one pool-minted id: the
leader's lives in the device worker beside the TP slab
(Job::TpSnapshotKv / TpRestoreKv / TpDropKvSnapshot), each subprocess
rank stores its own in-process via new WorkerRequest variants
(SnapshotKvCache / RestoreKvCache / DropKvSnapshot). Shard state has
the same shape as single-GPU (attention ConcatKvCache + GDN
conv/recurrent state + rope_delta), so the snapshot types are reused;
all ranks sit at the same token boundary because step fan-out is
synchronous.

Consistency on partial failure: a failed restore falls back to
clear-all-ranks + full prefill (and drops the entry); a failed
snapshot drops the id on every rank so nothing half-stored leaks.
DropTp / UnloadModel invalidate a model's snapshots with it, covering
auto-recovery. Vision requests bypass as on single-GPU. Budget
accounting uses leader bytes x world_size (shards are symmetric).

Wired into both TP request paths (non-streaming inner + streaming
orchestration task); chunked_prefill_tp gains the restored-offset
start.

Closes #11

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 17:34:49 +03:00
bb558451db Merge pull request 'feat(neuron): prefix KV caching across requests — single-GPU + CPU paths (#11)' (#34) from feat/11-prefix-kv-cache into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 29s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m15s
build-prerelease / Test (push) Successful in 4m0s
build-prerelease / Build neuron-blackwell (push) Successful in 9m44s
build-prerelease / Build neuron-ampere (push) Successful in 12m47s
build-prerelease / Build neuron-ada (push) Successful in 19m6s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 4m2s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 3m10s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 4m8s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 1m2s
2026-06-12 14:20:24 +00:00
c5378d532d feat(neuron): prefix KV caching across requests — single-GPU + CPU paths (#11)
All checks were successful
CI / Format (push) Successful in 32s
CI / Format (pull_request) Successful in 34s
CI / Clippy (push) Successful in 2m29s
CI / CUDA type-check (pull_request) Successful in 1m31s
CI / CUDA type-check (push) Successful in 1m37s
CI / Clippy (pull_request) Successful in 2m32s
CI / Test (push) Successful in 4m24s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Test (pull_request) Successful in 4m23s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
Stop discarding cache state between requests. When an incoming
prompt's token sequence starts with the exact tokens of a stored
snapshot, restore it and prefill only the divergent suffix.

For the hybrid qwen3_5 arch a snapshot is attention ConcatKvCache k/v
+ GatedDeltaNet conv/recurrent state + the rope_delta counter, all at
one token boundary; the recurrent state cannot rewind, so matching is
exact-prefix only. GDN states are deep-copied both directions (the
CUDA delta-rule kernels mutate the state buffer in place); attention
k/v snapshots share storage safely (append-by-cat never mutates).

Snapshots live in the device worker's state next to the model slab
(Job::SnapshotKv / RestoreKv / DropKvSnapshot); the async side holds
only an opaque id + token sequence + byte size. DropArch drops a
model's snapshots with it, so unload and auto-recovery invalidate for
free. CPU loads hold snapshots inline on the legacy path.

Per-model LRU registry (harness/prefix_cache.rs) bounded by
[harness.candle.prefix_cache] budget_mb / max_entries, enabled by
default; inserting a snapshot drops entries it strictly extends.
Vision requests and candle-transformers archs bypass the cache
entirely (clear-every-request, unchanged).

Covers the single-GPU worker path (streaming + non-streaming) and the
CPU-local path. The TP path (Qwen3.6-27B on beast) is a follow-up PR
that closes #11 with before/after bench numbers.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 17:14:07 +03:00
9f383e7bc7 Merge pull request 'feat(gateway): Anthropic streaming SSE translation (#24)' (#33) from feat/gateway-24-anthropic-sse into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 33s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m13s
build-prerelease / Test (push) Successful in 3m59s
build-prerelease / Build cortex binary (push) Successful in 2m15s
build-prerelease / Build neuron-blackwell (push) Successful in 9m59s
build-prerelease / Build neuron-ada (push) Successful in 14m24s
build-prerelease / Build neuron-ampere (push) Successful in 19m3s
build-prerelease / Package cortex RPM (push) Successful in 1m24s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 3m5s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 3m6s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 3m59s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 1m8s
2026-06-12 12:57:09 +00:00
569c528c4b feat(gateway): Anthropic streaming SSE translation (#24)
All checks were successful
CI / Format (push) Successful in 36s
CI / CUDA type-check (push) Successful in 2m25s
CI / Clippy (push) Successful in 2m25s
CI / Format (pull_request) Successful in 41s
CI / CUDA type-check (pull_request) Successful in 2m9s
CI / Clippy (pull_request) Successful in 2m45s
CI / Test (push) Successful in 5m3s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Test (pull_request) Successful in 4m29s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
The /v1/messages handler translated request envelopes but proxied raw
OpenAI SSE frames back to streaming Anthropic clients — the gap
between the README's "point your tooling at it once" contract and
what Claude Code actually received.

cortex-core gains AnthropicStreamTranslator, a pure per-stream state
machine: OpenAI chunks in, ordered (event, payload) pairs out —
message_start → content_block_start/delta/stop (text and tool_use
blocks, indexed; tool_calls map to input_json_delta) → message_delta
(stop_reason mapped via the now-shared map_stop_reason, which also
teaches the non-streaming path tool_calls→tool_use) → message_stop.
Without an upstream usage frame the output count falls back to the
delta count (engine-exact for neuron's one-chunk-per-token streams,
#31); with one, input/output tokens ride message_delta.

cortex-gateway gains anthropic_sse: the wire pump that splits the
upstream byte stream into SSE events, parses data: payloads
(leniently — engines omit fields on special frames), feeds the
translator, and frames results as `event:`/`data:` pairs through a
bounded channel (slow client back-pressures the upstream read).
Upstream truncation without [DONE] still closes the Anthropic event
sequence. Nothing is buffered beyond the current event's bytes.

Tests: 5 state-machine unit tests (text flow, stop-reason mapping +
defaults, tool_use blocks, usage propagation, idempotent finish) and
2 gateway integration tests (full event sequence + text reassembly,
usage propagation into message_delta). Validated end-to-end by
running this branch's gateway against a production neuron and
streaming a live Anthropic request.

Closes #24

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 15:47:30 +03:00
06e4ffc25c Merge pull request 'feat(bench): reproducible benchmark harness + first fleet numbers (#22)' (#32) from feat/22-benchmark-harness into main
Some checks failed
build-prerelease / Build neuron-blackwell (push) Blocked by required conditions
build-prerelease / Build neuron-ampere (push) Blocked by required conditions
build-prerelease / Build neuron-ada (push) Blocked by required conditions
build-prerelease / Resolve version stamps + change detection (push) Successful in 32s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m27s
build-prerelease / Build cortex binary (push) Successful in 2m41s
build-prerelease / Package cortex RPM (push) Successful in 1m29s
build-prerelease / Test (push) Successful in 4m44s
build-prerelease / Package helexa-neuron-ada RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been cancelled
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been cancelled
2026-06-12 12:46:33 +00:00
a2e73a8907 feat(bench): reproducible batch-1 benchmark harness + first fleet numbers (#22)
All checks were successful
CI / Format (push) Successful in 40s
CI / Format (pull_request) Successful in 38s
CI / CUDA type-check (push) Successful in 2m8s
CI / CUDA type-check (pull_request) Successful in 2m8s
CI / Clippy (push) Successful in 2m23s
CI / Test (pull_request) Successful in 3m54s
CI / Test (push) Successful in 6m23s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Clippy (pull_request) Successful in 4m23s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
script/bench.py: stdlib-only, works against any OpenAI-compatible /v1
endpoint (helexa, llama.cpp, Ollama, vLLM) so cross-engine tables are
a concatenation via the --label column. Measures the operator-felt
trio per (model, prompt-size) cell: TTFT (first SSE content chunk),
decode tok/s (visible tokens over the first→last chunk window,
chunk-per-token engine invariant since streaming usage frames aren't
emitted yet — #31), total wall-clock. Medians over N runs after one
warmup; append-only JSONL for longitudinal tracking.

Measurement traps found against the live fleet and handled:
- thinking models burn the budget invisibly (reasoning deltas are
  off-wire by default) — the prompt appends Qwen's /no_think soft
  switch
- short coalesced replies collapse the decode window to one TCP read
  — rates require a ≥200 ms window and the prompt demands ~300 words

doc/benchmarks.md: method, fleet table, and the first published
numbers (2026-06-12, 8f6f1d3): 1.7B@3060 81 tok/s, 8B@4090 62 tok/s,
27B@2×5090 Q6K TP=2 35 tok/s with flat decode from 128→4k context —
and the 7.1 s 4k-prefill TTFT recorded as #23's before-number.

Refs #22 (competitor baselines still pending — the harness is ready
for them)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 15:39:13 +03:00
8f6f1d3205 feat(deploy): validate neuron capability after every deploy
Some checks failed
build-prerelease / Build neuron-ampere (push) Blocked by required conditions
build-prerelease / Build neuron-ada (push) Blocked by required conditions
build-prerelease / Package cortex RPM (push) Blocked by required conditions
build-prerelease / Resolve version stamps + change detection (push) Successful in 29s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m14s
build-prerelease / Build neuron-blackwell (push) Successful in 10m36s
build-prerelease / Build cortex binary (push) Successful in 2m35s
build-prerelease / Test (push) Successful in 6m35s
build-prerelease / Package helexa-neuron-ada RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been cancelled
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been cancelled
A deploy previously went green the moment systemd reported the
service started — a merge that broke model loading or inference
itself would deploy "successfully" and only surface when a human
noticed. Each neuron deploy now earns its green:

1. Wait for default models: poll /health until activation.state is
   ready, with per-host timeouts in the matrix (beast 900s for the
   27B Q6K TP=2 cold-load, benjy/quadbrat 300s). Any entry in
   activation.failed fails the deploy with the per-model error —
   the structured equivalent of watching the journal for
   "loaded default model", plus failure detail the journal line
   can't carry.
2. LLM smoke probe: ask the first loaded model to reply with one
   specific word (max_tokens 512 so thinking models have room,
   temperature 0) and grep the response for it. Not a quality bar —
   just proof the deploy didn't lobotomize inference.

Hosts whose package is already current still skip everything — the
validation cost is only paid when a restart actually happened. The
probe was dry-run against benjy's production neuron before landing.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 15:28:20 +03:00
b0d0b939af Merge pull request 'feat(gateway): per-request token metrics — TTFT and tok/s (#21)' (#30) from feat/gateway-21-token-metrics into main
Some checks failed
build-prerelease / Lint (fmt + clippy) (push) Blocked by required conditions
build-prerelease / Test (push) Blocked by required conditions
build-prerelease / Build cortex binary (push) Blocked by required conditions
build-prerelease / Build neuron-blackwell (push) Blocked by required conditions
build-prerelease / Build neuron-ampere (push) Blocked by required conditions
build-prerelease / Build neuron-ada (push) Blocked by required conditions
build-prerelease / Resolve version stamps + change detection (push) Successful in 33s
build-prerelease / Package cortex RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ada RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been cancelled
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been cancelled
2026-06-12 12:25:32 +00:00
6a36d15ef1 feat(gateway): per-request token metrics — TTFT and tok/s (#21)
All checks were successful
CI / Format (push) Successful in 45s
CI / Format (pull_request) Successful in 37s
CI / CUDA type-check (push) Successful in 2m25s
CI / Clippy (push) Successful in 2m37s
CI / Test (push) Successful in 4m22s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Clippy (pull_request) Successful in 2m23s
CI / Test (pull_request) Successful in 4m19s
CI / CUDA type-check (pull_request) Successful in 1m57s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
The deferred Phase 6b, and the unblock for the 7→8 milestone's
benchmark work (#22): until cortex measures itself per request,
nothing downstream can be benchmarked or graphed.

The proxy wraps the upstream byte stream in a pass-through inspector
(TokenMetricsStream): chunks are forwarded verbatim — never buffered
or re-serialised — while the inspector records arrival times and
keeps a bounded (64 KiB) tail of the body text. At stream end (or
client disconnect, via Drop) it extracts the final OpenAI usage
object — present on the last SSE chunk and non-streaming JSON bodies
alike — for engine-truth token counts.

Per request, labelled {model, node}:
- cortex_time_to_first_token_seconds (histogram) — first body chunk
- cortex_tokens_per_second (histogram) — completion tokens over the
  decode window (first→last chunk); falls back to total request
  duration for single-chunk non-streaming bodies
- cortex_prompt_tokens_total / cortex_completion_tokens_total
  (counters)

The extractor is pure and chunk-boundary-safe; quoted-needle matching
keeps completion_tokens_details from shadowing completion_tokens,
and the last usage object wins. Covers chat completions, completions,
the Responses API, and the Anthropic streaming path (which currently
proxies OpenAI SSE).

Tests: 4 extractor unit tests; integration test with a streaming
mock emitting a stream_options-style final usage chunk, asserting
both histograms and exact-or-greater counter values (the test
recorder is process-global and shared across the binary's tests).

Closes #21

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 15:11:52 +03:00
b463439416 Merge pull request 'feat(neuron): startup preflight for NVIDIA driver/library mismatch (#19)' (#29) from feat/neuron-19-driver-preflight into main
Some checks failed
build-prerelease / Build neuron-ampere (push) Blocked by required conditions
build-prerelease / Build neuron-ada (push) Blocked by required conditions
build-prerelease / Resolve version stamps + change detection (push) Successful in 29s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m11s
build-prerelease / Build cortex binary (push) Successful in 2m33s
build-prerelease / Test (push) Successful in 4m24s
build-prerelease / Package cortex RPM (push) Successful in 1m27s
build-prerelease / Build neuron-blackwell (push) Successful in 10m18s
build-prerelease / Package helexa-neuron-ada RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been cancelled
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been cancelled
2026-06-12 12:08:20 +00:00
716558c8ff feat(neuron): startup preflight for NVIDIA driver/library mismatch (#19)
All checks were successful
CI / Format (push) Successful in 38s
CI / Format (pull_request) Successful in 38s
CI / CUDA type-check (push) Successful in 2m11s
CI / Clippy (push) Successful in 2m13s
CI / Clippy (pull_request) Successful in 2m37s
CI / Test (push) Successful in 4m17s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Test (pull_request) Successful in 3m56s
CI / CUDA type-check (pull_request) Successful in 1m44s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
The un-rebooted driver update (userspace libs bumped, kernel module
still old) kills every CUDA call on the host including nvidia-smi,
and neuron surfaced it only as `Comm::from_rank ... NcclError` deep
inside the first model load — 30 minutes of forensics on beast
(2026-06-08) to diagnose. Make it instantly legible instead:

- discovery distinguishes nvidia-smi absent (CPU-only, fine) from
  present-but-failing, classifies the "Driver/library version
  mismatch" signature, and pairs the userspace NVML version with the
  loaded kernel-module version from /proc/driver/nvidia/version.
- DiscoveryResponse gains `cuda_unavailable_reason` (omitted when
  None — wire-compatible) so cortex can see why the node has no
  devices and route around it.
- startup logs one loud ERROR line with the actionable reason
  ("reboot the host to reload the kernel module") and skips default
  model loads entirely, marking each failed with that reason so
  /health activation shows the real cause.
- POST /models/load fast-rejects with 503 + code=cuda_unavailable on
  a mismatch host instead of dying minutes later in cuInit/NCCL.

No false positives: other nvidia-smi failures (no devices, perms)
keep their existing behaviour, CPU-only hosts stay silent.

Closes #19

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 15:00:00 +03:00
112e4e124a fix(ci): export RUSTC_WRAPPER in the build step itself — GITHUB_ENV doesn't propagate
Some checks failed
build-prerelease / Package helexa-neuron-ada RPM (push) Blocked by required conditions
build-prerelease / Package helexa-neuron-ampere RPM (push) Blocked by required conditions
build-prerelease / Package helexa-neuron-blackwell RPM (push) Blocked by required conditions
build-prerelease / Resolve version stamps + change detection (push) Successful in 32s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m22s
build-prerelease / Build cortex binary (push) Successful in 2m20s
build-prerelease / Test (push) Successful in 3m50s
build-prerelease / Build neuron-blackwell (push) Successful in 10m10s
build-prerelease / Package cortex RPM (push) Successful in 1m25s
build-prerelease / Build neuron-ada (push) Successful in 14m29s
build-prerelease / Build neuron-ampere (push) Successful in 14m31s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been cancelled
Run 375 proved the CUDA image ships sccache (probe step printed
"sccache enabled") but the wrapper never reached cargo: the runner
does not propagate GITHUB_ENV across steps, so the builds ran
unwrapped (server stats: 4 compile requests for a ~600-crate build,
durations unchanged). Probe and export inside the build step's own
shell instead, in both build-neuron and ci.yml's cuda-check.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 14:50:25 +03:00
dc6feec6dc fix(deploy): gate on the publish manifest, not unprivileged dnf check-update
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 31s
build-prerelease / Build cortex binary (push) Successful in 2m18s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m33s
build-prerelease / Test (push) Successful in 4m20s
build-prerelease / Package cortex RPM (push) Successful in 1m23s
build-prerelease / Build neuron-blackwell (push) Successful in 9m46s
build-prerelease / Build neuron-ampere (push) Successful in 13m57s
build-prerelease / Build neuron-ada (push) Successful in 15m29s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 3m5s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 3m5s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 3m49s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 1m8s
The f5fa840 deploy exposed both failure modes of gating with
`dnf check-update` as the gitea_ci user in one run: it hung
indefinitely on quadbrat (blocked process, 0 CPU, killed manually),
and on benjy/beast it silently reported "no updates" two minutes
after new RPMs were published — both hosts skipped a real (luckily
binary-identical) update.

Gate with data we own instead: fetch packages.json from
rpm.lair.cafe (plain curl, no privileges, no dnf locks), take the
newest release per package by buildTime, and skip the
stop/upgrade/start cycle only when it exactly equals
`rpm -q %{VERSION}-%{RELEASE}`. Unreachable or unparsable manifest
fails open to a full deploy. The dnf transaction itself still runs
under the scoped sudoers rules, unchanged.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 14:20:21 +03:00
02f20bc9e1 Merge pull request 'feat: keep auto-recovering models visible as recovering (#20)' (#28) from feat/neuron-20-recovering-status into main
Some checks failed
build-prerelease / Test (push) Blocked by required conditions
build-prerelease / Build neuron-blackwell (push) Blocked by required conditions
build-prerelease / Build neuron-ampere (push) Blocked by required conditions
build-prerelease / Build neuron-ada (push) Blocked by required conditions
build-prerelease / Resolve version stamps + change detection (push) Successful in 30s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m39s
build-prerelease / Build cortex binary (push) Successful in 3m46s
build-prerelease / Package cortex RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ada RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been cancelled
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been cancelled
2026-06-12 11:15:38 +00:00
2a231e49de merge main (sccache enablement supersedes branch cuda-check pin)
All checks were successful
CI / Format (push) Successful in 40s
CI / Format (pull_request) Successful in 37s
CI / Clippy (push) Successful in 2m17s
CI / CUDA type-check (push) Successful in 2m39s
CI / CUDA type-check (pull_request) Successful in 2m30s
CI / Test (push) Successful in 4m51s
CI / Clippy (pull_request) Successful in 2m12s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Test (pull_request) Successful in 4m49s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
# Conflicts:
#	.gitea/workflows/ci.yml
2026-06-12 14:05:55 +03:00
2dadea5d8d ci: enable sccache on the build jobs (conditional on the CUDA image)
Some checks failed
build-prerelease / Build neuron-blackwell (push) Blocked by required conditions
build-prerelease / Resolve version stamps + change detection (push) Successful in 34s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m57s
build-prerelease / Test (push) Has been cancelled
build-prerelease / Build cortex binary (push) Has been cancelled
build-prerelease / Build neuron-ampere (push) Has been cancelled
build-prerelease / Build neuron-ada (push) Has been cancelled
build-prerelease / Package cortex RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ada RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been cancelled
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been cancelled
The 3 CUDA flavour builds (10-14 min each, the critical path of every
full run) and build-cortex compiled entirely uncached. With the
gongfoo-side sccache hardening in place, wire them up:

- build-cortex: full sccache env (rust image ships it) + the standard
  escalation loop (retry -> server restart -> uncached final attempt).
- build-neuron: probe for sccache before enabling the wrapper — the
  CUDA image may not ship it, and a missing binary must degrade to an
  uncached build, not fail cargo at `sccache rustc -vV` (the original
  reason the wrapper was cleared here). rustc compilations are shared
  across all three flavours; candle-kernels' nvcc output stays
  uncached (build-script artifact).
- ci.yml cuda-check: same probe pattern replaces the blanket env
  clear; also pins CUDA_COMPUTE_CAP=86 since the image no longer
  ships nvidia-smi for candle-kernels' fallback detection (mirrors
  9bb9678 on the #20 branch).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 14:05:26 +03:00
9bb9678f93 fix(ci): pin CUDA_COMPUTE_CAP in cuda-check — builder image has no nvidia-smi
All checks were successful
CI / Format (push) Successful in 37s
CI / Format (pull_request) Successful in 38s
CI / CUDA type-check (push) Successful in 1m45s
CI / Clippy (push) Successful in 2m24s
CI / Clippy (pull_request) Successful in 2m19s
CI / Test (push) Successful in 4m40s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Test (pull_request) Successful in 4m35s
CI / CUDA type-check (pull_request) Successful in 1m50s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
candle-kernels' build script shells out to nvidia-smi for compute-cap
detection when CUDA_COMPUTE_CAP is unset; the current GPU-less builder
image doesn't ship it, so the type-check died in the build script
before borrow-checking anything. Pin an arbitrary valid cap — the
check is feature-gate compilation only; real caps live in
build-prerelease.yml's flavour matrix.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 13:55:23 +03:00
df9c490614 feat(neuron+gateway): keep auto-recovering models visible as recovering (#20)
Some checks failed
CI / Format (push) Successful in 37s
CI / CUDA type-check (pull_request) Failing after 28s
CI / Format (pull_request) Successful in 37s
CI / Clippy (push) Successful in 2m54s
CI / Clippy (pull_request) Successful in 3m36s
CI / Test (push) Successful in 4m37s
CI / Test (pull_request) Successful in 5m20s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
CI / CUDA type-check (push) Failing after 31s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
During the #17 auto-recovery window (unload → reload, minutes for a
large TP model) the model's registry slot is absent, so it vanished
from neuron's /models — and cortex, routing by /models presence,
answered "model not found on any node" while a direct request to
neuron would have correctly said "recovering, retry shortly".

neuron: the recovery set becomes a map carrying a devices/capabilities
snapshot taken at trigger time (while the registry slot still exists).
list_models reports `recovering` for models in the set — both while
the poisoned slot is still present and during the reload gap, where
the snapshot keeps the model listed.

gateway: ModelStatus grows a Recovering variant (parsed from the
wire); the router holds the route — new RouteError::ModelRecovering
mapped to 503 instead of 404 — and deliberately does not fall through
to the catalogue cold-load, which would race a second placement
against the in-flight recovery. The evictor already ignores
non-Loaded entries.

Tests: neuron unit test (recovering model stays listed with snapshot),
gateway integration tests (poller parses `recovering`; request gets
503 retry-shortly and the model stays on /v1/models).

Closes #20

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 13:42:03 +03:00
f5fa840dfb ci: escalate sccache retries — restart server, then fall back uncached
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 30s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m6s
build-prerelease / Test (push) Successful in 4m50s
build-prerelease / Build cortex binary (push) Successful in 3m45s
build-prerelease / Build neuron-blackwell (push) Successful in 9m59s
build-prerelease / Build neuron-ada (push) Successful in 14m11s
build-prerelease / Build neuron-ampere (push) Successful in 14m13s
build-prerelease / Package cortex RPM (push) Successful in 1m30s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 3m28s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 3m50s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 3m54s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 1m3s
Run 361's Test job failed all 3 attempts with the sccache
dead-server signature (sccache fatal error, ENOENT on its own tmp
files under target/debug/deps). Retrying the same invocation only
helps for transient races; against a wedged server every same-VM
retry fails identically — and under the new pipeline that blocks
publish and the deploy behind it.

Escalate instead: attempt 1 plain, attempt 2 after an sccache server
restart, attempt 3 with RUSTC_WRAPPER unset (uncached). A sick cache
now costs build minutes, never the deploy. Applied to the lint/test
jobs in build-prerelease.yml and ci.yml alike.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 13:24:02 +03:00
7557c5e877 ci: cut iteration latency — change-aware builds, gated deploys, dev fast path
Some checks failed
build-prerelease / Build neuron-blackwell (push) Blocked by required conditions
build-prerelease / Resolve version stamps + change detection (push) Successful in 28s
build-prerelease / Test (push) Failing after 1m16s
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m7s
build-prerelease / Build cortex binary (push) Successful in 3m57s
build-prerelease / Build neuron-ampere (push) Has been cancelled
build-prerelease / Build neuron-ada (push) Has been cancelled
build-prerelease / Package cortex RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ada RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been cancelled
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been cancelled
Push-to-testable was ~20.5 min for every commit (measured on the
2026-06-08 green chain) plus a ~5 min 27B cold-load, regardless of
what changed. Three structural fixes:

- build-prerelease: a change-detection step in `prepare` diffs HEAD
  against the git sha embedded in the last *published* unstable RPM
  (per package, from packages.json) and skips builds whose inputs
  didn't change. Docs-only commits build nothing; gateway-only
  commits skip the 3 CUDA flavour builds. Detection failures fall
  open to a full build.
- ci.yml no longer runs on pushes to main; fmt/clippy/test live in
  build-prerelease as parallel jobs gating publish. The two workflows
  previously queued against each other on the same runner labels,
  delaying the cortex build ~12 min. Branches, PRs, and tags keep the
  full ci.yml gate.
- deploy: each host self-gates with `dnf check-update` and leaves the
  service untouched when the installed package is already current —
  no more neuron restarts (and 27B cold-loads) for commits that
  didn't change neuron.
- deploy-dev (new): manual single-host fast path — build one CUDA
  flavour, scp the binary, restart the service. Skips packaging,
  signing, publish, and dnf entirely. Backed by a new exact-form
  sudoers rule in asset/sudoers.d/neuron-host.conf (already applied
  to all three hosts).

Expected loop times when runners behave: docs ≈ 1 min (nothing
deploys), gateway-only ≈ 6-8 min, single-neuron dev ≈ 8-10 min,
full fleet ≈ 13-15 min.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 13:17:22 +03:00
91e95ca979 docs: rewrite README around project positioning
Some checks failed
CI / CUDA type-check (push) Failing after 46s
CI / Format (push) Successful in 47s
CI / Clippy (push) Successful in 2m53s
CI / Test (push) Successful in 4m31s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Package helexa-neuron-ada RPM (push) Blocked by required conditions
build-prerelease / Package helexa-neuron-ampere RPM (push) Blocked by required conditions
build-prerelease / Package helexa-neuron-blackwell RPM (push) Blocked by required conditions
build-prerelease / Resolve version stamps (push) Successful in 39s
build-prerelease / Build cortex binary (push) Successful in 3m52s
build-prerelease / Package cortex RPM (push) Successful in 1m18s
build-prerelease / Build neuron-blackwell (push) Successful in 11m34s
build-prerelease / Build neuron-ampere (push) Successful in 15m31s
build-prerelease / Build neuron-ada (push) Successful in 15m37s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been cancelled
Lead with what helexa is for — near-frontier open-weight models on
consumer hardware you own — instead of a feature list. Adds the scope
section (intentional divergence from vLLM/SGLang; CUDA-only today as a
test-coverage constraint, not a principle), an engine section covering
the per-device worker threads and consumer-GPU tensor parallelism, the
previously-missing helexa-acp crate, and a status section pointing at
git.lair.cafe as the source of truth with GitHub as read-only mirror.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 11:37:00 +03:00
1a74cb0c56 chore: rename repo cortex -> helexa
Some checks failed
CI / CUDA type-check (push) Failing after 30s
build-prerelease / Resolve version stamps (push) Successful in 45s
CI / Format (push) Successful in 32s
build-prerelease / Build neuron-blackwell (push) Failing after 31s
build-prerelease / Build neuron-ada (push) Failing after 34s
build-prerelease / Build neuron-ampere (push) Failing after 38s
build-prerelease / Package helexa-neuron-ada RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been skipped
CI / Clippy (push) Failing after 1m11s
build-prerelease / Build cortex binary (push) Successful in 3m47s
CI / Test (push) Successful in 5m32s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
build-prerelease / Package cortex RPM (push) Successful in 1m22s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
helexa is the project; cortex (per-operator control plane / LLM proxy)
and neuron (per-host LLM harness) are its components. The Gitea repo
is now helexa/helexa. Update repository URLs in Cargo metadata, RPM
specs, and docs; make the CI changelog push URL rename-proof via the
github.repository context; reframe README.md and CLAUDE.md around the
project name. Binary, package, service, and config-path names are
unchanged.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 10:54:01 +03:00
60f5598542 build(neuron): bump cudarc fork to 63327a2 (idempotent abort + Comm Send+Sync)
Some checks failed
build-prerelease / Resolve version stamps (push) Successful in 29s
CI / CUDA type-check (push) Successful in 31s
CI / Format (push) Successful in 35s
CI / Test (push) Failing after 1m9s
CI / Clippy (push) Successful in 2m36s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Build neuron-blackwell (push) Successful in 6m10s
build-prerelease / Build neuron-ampere (push) Successful in 7m35s
build-prerelease / Build neuron-ada (push) Successful in 5m7s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 2m53s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 3m14s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 3m48s
build-prerelease / Build cortex binary (push) Successful in 4m33s
build-prerelease / Package cortex RPM (push) Successful in 1m21s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 1m1s
The fork's new commit makes `Comm: Send + Sync` (asserting NCCL's
thread-safety invariant upstream) and makes `Comm::abort` idempotent via
an `aborted` flag (so abort-then-Drop can't double-free) — strictly
better than the previous Drop-no-panic workaround, and the `abort()`
signature is unchanged so the watchdog call site is unaffected.

Because `Comm` is now `Send + Sync`, `Arc<Comm>` and the `SendComm` /
`NcclState` wrappers auto-derive `Send`/`Sync`, which conflicts (E0119)
with neuron's manual `unsafe impl`s. Remove the four now-redundant impls
— the safety assertion lives upstream in cudarc where it belongs. The
conflict is in cuda-gated code, so only the CUDA type-check catches it
(non-cuda build + clippy + tests stay green).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-08 16:33:14 +03:00
7945240646 chore: re-trigger deploy (#17 Stage 2, attempt 3)
All checks were successful
CI / CUDA type-check (push) Successful in 31s
build-prerelease / Resolve version stamps (push) Successful in 31s
CI / Format (push) Successful in 33s
CI / Clippy (push) Successful in 2m41s
build-prerelease / Build cortex binary (push) Successful in 4m45s
build-prerelease / Build neuron-blackwell (push) Successful in 5m50s
CI / Test (push) Successful in 6m44s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Package cortex RPM (push) Successful in 1m23s
build-prerelease / Build neuron-ampere (push) Successful in 8m38s
build-prerelease / Build neuron-ada (push) Successful in 5m36s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 2m55s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 2m59s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 3m43s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 59s
No code change. Each deploy run, the degraded CI runner kills a different
single arch build (blackwell, then ada) ~fast, and the all-arch-gated
packaging skips → no publish. Every arch HAS built green across runs
(blackwell  in 342, ampere , ada  in 339) and the gate + CUDA
type-check pass. Re-running to catch all three green in one run so the
Stage-2 RPMs publish. Runner FS/cache health is the real fix (separate
infra work).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-08 15:06:04 +03:00
0c74d89d15 chore: re-trigger deploy (#17 Stage 2)
Some checks failed
CI / CUDA type-check (push) Successful in 32s
build-prerelease / Resolve version stamps (push) Successful in 29s
CI / Format (push) Successful in 30s
build-prerelease / Build neuron-ada (push) Failing after 51s
CI / Clippy (push) Successful in 2m41s
build-prerelease / Build cortex binary (push) Successful in 4m28s
build-prerelease / Build neuron-blackwell (push) Successful in 6m32s
build-prerelease / Build neuron-ampere (push) Successful in 7m42s
build-prerelease / Package helexa-neuron-ada RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been skipped
CI / Test (push) Successful in 6m6s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Package cortex RPM (push) Successful in 1m21s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been skipped
No code change. The c94a2ae deploy's neuron-blackwell build died ~12min
into the Blackwell kernel compile on the degraded runner, while
neuron-ampere + neuron-ada built the identical Rust + patched cudarc
cleanly and the CUDA type-check passed. Transient infra; re-running to
get a healthy blackwell build so the RPMs publish and beast (Blackwell)
picks it up.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-08 14:45:16 +03:00
c94a2ae755 fix(neuron): correct nccl_state path on WorkerPool.leader_comm (#17 S2)
Some checks failed
CI / CUDA type-check (push) Successful in 32s
build-prerelease / Resolve version stamps (push) Successful in 35s
CI / Format (push) Successful in 44s
build-prerelease / Build cortex binary (push) Successful in 4m57s
build-prerelease / Package cortex RPM (push) Successful in 1m36s
CI / Test (push) Successful in 7m10s
CI / Clippy (push) Failing after 1m21s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Build neuron-ampere (push) Successful in 8m40s
build-prerelease / Build neuron-ada (push) Successful in 9m5s
build-prerelease / Build neuron-blackwell (push) Failing after 12m2s
build-prerelease / Package helexa-neuron-ada RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been cancelled
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been cancelled
`super::nccl_state` from tp/mod.rs resolves to `crate::harness::nccl_state`
(nonexistent); the module is the child `nccl_state` (cf. the existing
`nccl_state::generate_comm_id_hex` call). The field is cuda-gated so the
non-cuda build couldn't catch it; the branch CUDA type-check flaked on the
runner before compiling. Self-audited fix.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-08 14:21:43 +03:00
99920dd322 feat(neuron): TP step watchdog aborts wedged collectives (#17 Stage 2)
Some checks failed
CI / CUDA type-check (push) Failing after 47s
CI / Format (push) Successful in 31s
CI / Test (push) Failing after 1m3s
CI / Clippy (push) Successful in 2m44s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
Make a hung NCCL collective recoverable instead of a permanent brick.
Today a wedged collective hangs the in-process leader thread forever, and
even Stage 1's recovery can't help — its unload's DropTp queues behind the
stuck thread and hangs too.

- Cache the leader's NCCL Comm handle async-side at init (new cuda-gated
  Job::GetLeaderComm → DeviceWorkerHandle::get_leader_comm → stored on
  WorkerPool.leader_comm). Fetched while the thread is responsive — a
  wedged thread can't service the fetch, which is why it's cached up front.
- Wrap the leader forward in both generate_step and
  generate_step_with_images in tokio::time::timeout (default 120s,
  NEURON_TP_STEP_TIMEOUT_S). On expiry the watchdog calls
  Comm::abort() (ncclCommAbort) on the cached handle from the async
  thread — the one NCCL op sanctioned concurrently with an in-flight
  collective — which unblocks the leader thread, then fails the step
  WITHOUT draining (workers are wedged too; recovery's unload kills them).
  The error is a device fault → poison → Stage 1 auto-recovery, which now
  completes because the leader thread is responsive again.
- Bumps the cudarc patch to dbc425a (adds the Drop-must-not-panic fix so
  the post-abort comm teardown during recovery doesn't double-abort-panic).

Logs the whole sequence at ERROR with greppable `tp watchdog:` /
`ncclCommAbort` markers so a real-world hang leaves a forensic trail —
verification is by inspecting journals after real hangs, not a synthetic
harness. cuda-gated → validated by the blackwell build.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-08 14:15:29 +03:00
c4f239ceb9 build(neuron): patch cudarc to expose Comm::abort/get_async_error (#17 Stage 2)
All checks were successful
CI / CUDA type-check (push) Successful in 33s
CI / Format (push) Successful in 35s
CI / Clippy (push) Successful in 2m34s
CI / Test (push) Successful in 6m1s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
#17 Stage 2 (TP hang-recovery) needs to call ncclCommAbort on a LIVE
communicator from another thread — to unblock a collective wedged on a
dead/hung peer so the ranks can resync. No cudarc release (incl. main)
exposes this: the safe Comm only aborts in Drop, which can't fire while a
stuck thread holds an Arc<Comm> clone.

Pin neuron's cudarc 0.19.7 to a fork (grenade/cudarc @ nccl-comm-abort,
rev 4dff0be) adding three thin methods — Comm::abort, get_async_error,
and a raw comm() accessor — to be submitted upstream. The patch targets
0.19.x only; candle's transitive cudarc 0.17.8 stays on crates.io.

Foundation only; the watchdog + abort + comm-rebuild that consume these
land in follow-up commits (cuda-gated → validated by the blackwell build).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-08 13:49:59 +03:00
ac445c1569 chore: re-trigger deploy (#17 Stage 1)
Some checks failed
CI / CUDA type-check (push) Failing after 19s
CI / Format (push) Successful in 37s
build-prerelease / Resolve version stamps (push) Successful in 42s
CI / Clippy (push) Successful in 3m54s
build-prerelease / Build cortex binary (push) Successful in 4m43s
CI / Test (push) Successful in 6m35s
build-prerelease / Build neuron-blackwell (push) Successful in 5m58s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Package cortex RPM (push) Successful in 1m21s
build-prerelease / Build neuron-ampere (push) Successful in 8m10s
build-prerelease / Build neuron-ada (push) Successful in 5m21s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 2m56s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 3m1s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 3m46s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 1m4s
No code change. The abc6e60 deploy's neuron-ada build died on the
degraded CI runner (container dropped mid-checkout), skipping the
gated publish — even though neuron-blackwell + neuron-ampere compiled
the Stage-1 fault-recovery code cleanly. Re-running to get a healthy
ada build so the RPMs publish and beast picks up the build.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-08 09:34:20 +03:00
abc6e605b8 test(neuron): NEURON_DEBUG_POISON hook to verify auto-recovery (#17)
Some checks failed
CI / CUDA type-check (push) Failing after 19s
build-prerelease / Resolve version stamps (push) Successful in 43s
CI / Format (push) Successful in 50s
CI / Clippy (push) Failing after 57s
build-prerelease / Build neuron-ada (push) Failing after 48s
build-prerelease / Build cortex binary (push) Successful in 5m5s
build-prerelease / Build neuron-blackwell (push) Successful in 6m38s
build-prerelease / Package cortex RPM (push) Successful in 1m27s
build-prerelease / Build neuron-ampere (push) Successful in 7m27s
build-prerelease / Package helexa-neuron-ada RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been skipped
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been skipped
CI / Test (push) Successful in 10m27s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
One-shot, env-gated fault injector for beast verification: when
NEURON_DEBUG_POISON names a model, the first request for it triggers the
auto-recovery path as if a device fault had occurred — exercising
unload→reload→healthy without corrupting the GPU. Latched so it fires
exactly once (no recovery loop). No-op unless the env var is set; wired
into both the single-GPU and TP chat poison gates.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-08 09:08:40 +03:00
4f2957af9e feat(neuron): auto-recover poisoned models (#17 Stage 1c)
When an inference hit a device fault, the model was flagged poisoned and
every subsequent request rejected with "unload and reload the model to
recover" — until a *human* did exactly that. Now the harness rebuilds the
context automatically.

- Retain the loading `ModelSpec` on `LoadedModel`/`TpLoadedModel` (+
  `LoadedHandle::spec()`) so a poisoned model can be reloaded without an
  operator reconstructing the spec.
- A background recovery task (held via `Weak<CandleHarness>`, spawned in
  `new()` when a runtime is present) drains poisoned model ids and runs
  `unload_model` → `load_model(spec)`. Unload drops the model → cudarc
  `Comm::drop` aborts NCCL + releases the context; reload re-runs NCCL
  init + sanity inside the load path, so a successful reload yields a
  fresh, healthy model. A failed reload leaves it unloaded (next load
  retries) — never poisoned forever.
- The request-entry poison gates now `trigger_recovery` (single-flight
  per model via a `recovering` set) and return a transient "recovering,
  retry shortly" error instead of the manual-reload message. Requests
  that arrive during the brief reload gap (model absent from the registry)
  also get "recovering" rather than a misleading "not loaded".

`new()` now returns `Arc<Self>`. Recovery runs only on the background
task — never inline on the request path, which holds `inference_lock`
and would deadlock on the `models` write lock.

Stage 1c of the #17 plan (verified-healthy auto-recovery). Watchdog
(1b) + a fault-injection hook for beast verification follow. The
in-process rank-0 leader's own context fault still needs a reload that
can't rebind it (Stage 3); comm-desync + worker faults recover here.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-08 09:05:02 +03:00
75cd088b61 fix(neuron): cap vision max_pixels to the pos_embed patch budget (#14)
All checks were successful
CI / CUDA type-check (push) Successful in 31s
build-prerelease / Resolve version stamps (push) Successful in 29s
CI / Format (push) Successful in 30s
CI / Clippy (push) Successful in 2m32s
build-prerelease / Build neuron-blackwell (push) Successful in 6m5s
CI / Test (push) Successful in 5m49s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Build neuron-ampere (push) Successful in 8m11s
build-prerelease / Build neuron-ada (push) Successful in 5m40s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 3m4s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 3m2s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 3m57s
build-prerelease / Build cortex binary (push) Successful in 4m21s
build-prerelease / Package cortex RPM (push) Successful in 1m25s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 1m16s
Beast testing surfaced a real regression in the dynamic-resolution
default: a tall 808×1600 image resized (within the 1024² max_pixels) to a
90×44 patch grid = 3960 patches, exceeding the vision tower's hard
`num_position_embeddings = 2304` pos-embed budget. The per-rank
`patch count 3960 exceeds pos_embed budget 2304` error fired mid-TP-
forward and poisoned the device context, bricking the model until reload.

Hard-cap `max_pixels` to `2304 × 16² = 589_824` px (≤ 2304 patches →
≤ 576 LM tokens), clamping even the operator env override. `smart_resize`
floors the pixel count under the cap, so no resized image can ever exceed
the budget — the tower check never fires, no poison. The pos-embed grid
(48×48) is the resolution Qwen3.6 was trained at, so the cap is
principled, not just defensive. Still ~3× the old fixed 196 tokens, and
the book-cover OCR test (1176 patches) already reads full title+subtitle.

Test: a huge/tall/wide/extreme image battery stays within the 2304 patch
budget. (Per-rank-error poison robustness itself remains issue #17.)

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-04 23:30:47 +03:00
d311c8ca7a feat(neuron): operator pixel-budget env override + doc cleanup (#14 C5)
Some checks failed
CI / CUDA type-check (push) Successful in 32s
build-prerelease / Resolve version stamps (push) Successful in 38s
CI / Format (push) Successful in 45s
CI / Test (push) Failing after 58s
CI / Clippy (push) Successful in 2m41s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Build cortex binary (push) Successful in 4m14s
build-prerelease / Package cortex RPM (push) Successful in 1m23s
build-prerelease / Build neuron-blackwell (push) Successful in 6m20s
build-prerelease / Build neuron-ampere (push) Successful in 7m18s
build-prerelease / Build neuron-ada (push) Successful in 5m10s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 3m6s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 3m7s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 3m45s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 1m5s
- PreprocessProfile::qwen3_6() reads NEURON_VISION_MIN_PIXELS /
  NEURON_VISION_MAX_PIXELS (clamped to factor² ≤ min ≤ max), matching the
  NEURON_VISION_LEGACY_* / NEURON_MROPE knob convention. Defaults remain
  256²…1024² (64…1024 LM tokens/image).
- Test: a max-resolution source caps within the token budget (can't blow
  NEURON_MAX_PROMPT_TOKENS).
- Strip stale fixed-resolution / "MRoPE gap (#15)" / 14×14 language from
  the preprocess, mod, and rope doc-comments now that resolution is
  dynamic and M-RoPE is implemented.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-04 22:50:03 +03:00
c97a8654f5 feat(neuron): dynamic-resolution images via Qwen smart_resize (#14)
Some checks failed
CI / Clippy (push) Waiting to run
CI / Test (push) Waiting to run
CI / CUDA type-check (push) Successful in 32s
CI / Format (push) Successful in 34s
CI / Build cortex SRPM (push) Has been cancelled
CI / Build neuron SRPM (push) Has been cancelled
CI / Publish cortex to COPR (push) Has been cancelled
CI / Publish neuron to COPR (push) Has been cancelled
CI / Bump version in source (push) Has been cancelled
Replace the fixed 448×448-square preprocess with native-aspect
`smart_resize`, and thread the resulting per-image grid through the LM
so spatial structure survives non-square images (documents, screenshots,
charts, panoramas, OCR) instead of being squished into a square.

- preprocess.rs: port Qwen `smart_resize` (factor = patch×merge = 32;
  pixel budget [min,max], default 256²–1024² → 64–1024 LM tokens).
  `PreprocessProfile` drops the fixed target dims for `factor`/`min_pixels`/
  `max_pixels`; `preprocess`/`preprocess_data_uri` now return the resized
  `(h, w)`; add `resized_dims_for_uri` (decode + resize, no normalize) for
  the TP leader's token count.
- rope.rs: `compute_mrope_index`/`get_rope_index` take per-image
  `grids: &[(lm_gh, lm_gw)]` instead of assuming a square `isqrt(run)`.
  Walk image runs in order, validate `run == gh*gw`, emit row-major
  positions, resume the shared counter at `base + max(gh,gw)`. Correct
  for multiple images of differing grids interleaved with text.
- candle.rs: `VisionMeta`/`LoadedModel`/`TpLoadedModel` carry the
  `image_grid_factor` (patch×merge) instead of the constant 196; all four
  prompt-build sites compute per-image counts from each image's resized
  grid (single-GPU from the extracted `ImageInput.h/w`, TP from
  `resized_dims_for_uri`). `ModelArch` gains `vision_grid_factor`.
- single-GPU (`mod.rs`, `dispatch.rs`) and TP
  (`tp_qwen3_5.rs::prefill_with_images_chunked`, `dispatch.rs`,
  `tp/worker.rs`) thread the grids into `get_rope_index`. Each TP rank
  recomputes grids from its own deterministic preprocess — no rpc.rs
  change, single source of truth.

The vision tower itself was already grid-general (recent pos-embed
interpolation + 2D rotary fix). No patch-count cap: pos-embed is
interpolated to any grid; `max_pixels` bounds cost (O(patches²) ViT
attention + prefill) instead.

Tests: smart_resize (aspect/cap/floor/reject), `compute_mrope_index`
non-square + two-image + mismatch cases, square-grid regression guard.
Non-cuda build + clippy + full workspace tests green; TP load/dispatch
paths are cuda-gated → Gitea CUDA type-check. Operator pixel-budget
config + remaining doc cleanup follow in C5.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-04 22:47:27 +03:00
dc048ffcc9 fix(neuron): vision-tower 2D positions + M-RoPE default on
All checks were successful
CI / CUDA type-check (push) Successful in 32s
build-prerelease / Resolve version stamps (push) Successful in 32s
CI / Format (push) Successful in 33s
CI / Clippy (push) Successful in 2m36s
build-prerelease / Build cortex binary (push) Successful in 4m48s
build-prerelease / Build neuron-blackwell (push) Successful in 5m59s
CI / Test (push) Successful in 6m35s
build-prerelease / Build neuron-ampere (push) Successful in 7m51s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Package cortex RPM (push) Successful in 1m21s
build-prerelease / Build neuron-ada (push) Successful in 5m13s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 3m0s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 3m5s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 3m49s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 1m6s
Two fixes to the spatial handling of images, validated against the HF
transformers 4.57.1 qwen3_vl reference on beast.

**Vision tower (the real cause of poor spatial vision).** The Stage-A
tower encoded position two ways wrong, so the model saw image *content*
but not *layout* (a row of 5 people read as "a line of 23", sky
inverted), regardless of the LM-side rope:

- Learned pos-embed was a naive sequential lookup of the first
  `n_patches` rows of the 48×48 (`num_position_embeddings=2304`) grid —
  wrong stride for a 28×28 patch grid. Now bilinearly interpolates the
  grid to `gh×gw` (port of HF `fast_pos_embed_interpolate`), row-major.
- The 2D vision rotary was absent entirely. Added
  `VisionRotaryEmbedding` (θ=10000, dim=head_dim/2) applying per-patch
  `(row, col)` rotary to q/k in every ViT block via rope_slow, matching
  HF `apply_rotary_pos_emb_vision`.

Both default on; `NEURON_VISION_LEGACY_POS=1` / `NEURON_VISION_LEGACY_ROPE=1`
revert each for A/B (no rebuild). New unit tests: interpolation reduces
to the sequential lookup at the native grid; rotary row/col structure.

**M-RoPE default on.** The interleaved M-RoPE matches HF
apply_interleaved_mrope / get_rope_index exactly and A/B'd strictly ≥
plain. `NEURON_MROPE` is now a kill switch (`=0` for plain), not opt-in
— defaults should encode the model's trained behaviour, not freeze the
broken state.

Vision tower is plain candle (CPU-testable): built, clippy-clean, full
workspace tests green locally.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-04 20:53:07 +03:00
7ebcfba5ca fix(neuron): gate M-RoPE behind NEURON_MROPE (default off)
All checks were successful
CI / CUDA type-check (push) Successful in 33s
build-prerelease / Resolve version stamps (push) Successful in 32s
CI / Format (push) Successful in 33s
CI / Clippy (push) Successful in 2m34s
build-prerelease / Build cortex binary (push) Successful in 4m33s
build-prerelease / Build neuron-blackwell (push) Successful in 6m14s
CI / Test (push) Successful in 6m50s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Build neuron-ampere (push) Successful in 8m12s
build-prerelease / Package cortex RPM (push) Successful in 1m23s
build-prerelease / Build neuron-ada (push) Successful in 5m9s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 2m59s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 3m3s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 3m52s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 1m2s
On beast the interleaved M-RoPE degraded image understanding rather than
fixing it: the model misread spatial layout (a horizontal row of people
described as a "diagonal receding line"), got attributes wrong, and
rambled — a "how many people" follow-up generated 4459 tokens over 3.5
minutes, past agent-0's HTTP timeout (the "fails to respond without an
error"). The interleave is evidently not numerically correct, and it
can't be validated remotely without a transformers reference.

Gate it: `get_rope_index` now returns plain sequential identity
positions unless NEURON_MROPE is truthy, so mrope_cos_sin reduces to
plain RoPE and image tokens behave exactly as pre-M-RoPE (content
recognition works; spatial layout approximate; no rambling). The real
computation moves to `compute_mrope_index` (still unit-tested). Default
off restores the working vision and unblocks agent-0; the M-RoPE code
stays in place to debug + validate before flipping the default on.

Pure non-cuda change (rope.rs); both single-GPU and TP forwards call
the gated get_rope_index unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-04 19:32:59 +03:00
825bf4e905 feat(neuron): M-RoPE Stage 4 — wire interleaved M-RoPE into the TP path
All checks were successful
build-prerelease / Resolve version stamps (push) Successful in 30s
CI / CUDA type-check (push) Successful in 31s
CI / Format (push) Successful in 42s
build-prerelease / Build cortex binary (push) Successful in 5m9s
build-prerelease / Build neuron-blackwell (push) Successful in 6m4s
build-prerelease / Package cortex RPM (push) Successful in 1m32s
CI / Test (push) Successful in 7m19s
build-prerelease / Build neuron-ampere (push) Successful in 8m40s
build-prerelease / Build neuron-ada (push) Successful in 5m17s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 3m0s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 3m1s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 3m53s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 1m14s
CI / Clippy (push) Successful in 2m29s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
Mirror Stage 3 into the tensor-parallel Qwen3.6 model:

- TpQwen3_5Attention / DecoderLayer take (cos, sin) instead of a scalar
  offset and apply via apply_cos_sin.
- TpQwen3_5Model gains the replicated rotary + rope_delta (reset in
  clear_kv_cache, settable). forward_inner builds the cos/sin once —
  interleaved M-RoPE from explicit position_ids (vision) or plain at
  offset+rope_delta (text/decode). forward() and forward_with_positions()
  delegate; the old single-shot forward_with_vision is gone.
- prefill_with_images_chunked now computes get_rope_index over the whole
  prompt once, stores rope_delta on the base model, and slices the
  (3, prompt_len) position tensor per chunk — so every rank assigns image
  tokens their 14×14 grid coordinates and steps in lockstep (every chunk,
  text or image, carries the M-RoPE slice because the image shifts the
  surrounding text positions).

Also build the position-id tensor as f32 directly (positions are small
integers, exact in f32) to avoid an i64→f32 cast on the GPU.

The TP forward is cuda-gated — CI CUDA type-check is the compile gate.
Non-cuda build + clippy + full workspace tests green; rope math + the
plain-RoPE-reduction invariant covered by unit tests.

Completes the interleaved-M-RoPE work for the vision spatial misread.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-04 18:46:27 +03:00
4c12c7e2f0 feat(neuron): M-RoPE Stage 3 — wire interleaved M-RoPE into single-GPU
Qwen3_5Model now builds the rotary cos/sin once per forward and threads
(cos, sin) through the decoder → full-attention → rope, replacing the
scalar offset that reached RotaryEmbedding:

- vision forward computes get_rope_index over the (single-shot) prompt,
  sets rope_delta, and builds interleaved-M-RoPE cos/sin so image tokens
  carry their 14×14 grid (height/width) positions;
- text / decode take plain_cos_sin at offset + rope_delta — with
  rope_delta == 0 (no image) this is bit-for-bit the old plain RoPE, and
  the device→host id copy is skipped on the text decode hot path.

rope_delta is stored on the model and reset in clear_kv_cache, so decode
after a vision prefill resumes text positions from the image-compressed
counter. decoder.rs / full_attn.rs take (cos, sin) instead of offset;
linear-attention layers are unchanged (no RoPE). The TP path still uses
the retained apply(offset) — wired in Stage 4.

Full workspace tests green; the load-bearing invariant (M-RoPE == plain
for equal axes) keeps text unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-04 18:39:52 +03:00
ba1b5ba408 feat(neuron): M-RoPE Stage 2 — get_rope_index position-id helper
Pure function computing the interleaved-M-RoPE 3D position ids for a
prompt with image-placeholder runs, plus the decode rope_delta:
text tokens advance a single counter (all axes equal); each image run
gets [base+t, base+h, base+w] row-major over a square grid_t=1,
grid_h=grid_w=isqrt(run) (196 → 14×14); the counter resumes from
base + max(grid). rope_delta = final_counter - seq_len lets decode
resume text positions after the position-compressed image blocks.
Plus mrope_position_tensor to build the (3, seq) tensor.

Unit tests: text-only is sequential (delta 0); text+image+text matches
hand-computed grid ids + resume + delta; 196 → 14×14; non-square run
rejected; end-to-end through mrope_cos_sin tracks the height axis.

#[allow(dead_code)] until Stage 3/4 wire it into the forward.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-04 18:34:28 +03:00
5731f4c318 feat(neuron): M-RoPE Stage 1 — interleaved rope machinery + config
Parse + store mrope_section / mrope_interleaved in RopeParameters
(previously accepted-but-ignored). RotaryEmbedding gains:
- inv_freq + per-axis column masks (mask_t/h/w) built from mrope_section;
- plain_cos_sin(pos, seq_len): narrow the precomputed tables (text/decode);
- mrope_cos_sin(position_ids (3,seq)): per-axis freqs blended at the
  interleave columns (vision);
- apply_cos_sin(q,k,cos,sin): the rope_slow application, factored out.

The existing apply(q,k,offset) is retained (delegates to
plain_cos_sin + apply_cos_sin) so current callers are unchanged; Stages
3–4 move cos/sin construction into the model forward and thread the 3D
position ids for image tokens.

Tests: masks partition the half-dim; interleave drives the right axis
per column; and the load-bearing invariant — mrope_cos_sin reduces
bit-for-bit to plain_cos_sin when the three axes are equal (so text
inference is unchanged).

Refs the MRoPE-gap diagnosis (vision spatial misread). Pure non-cuda;
no behaviour change until wired.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-04 18:31:15 +03:00
fa013505d1 fix(neuron): chunked TP-vision prefill + pre-flight VRAM guard
All checks were successful
build-prerelease / Resolve version stamps (push) Successful in 29s
build-prerelease / Build cortex binary (push) Successful in 4m26s
build-prerelease / Package cortex RPM (push) Successful in 1m18s
build-prerelease / Build neuron-blackwell (push) Successful in 6m6s
build-prerelease / Build neuron-ampere (push) Successful in 8m30s
CI / Format (push) Successful in 38s
CI / CUDA type-check (push) Successful in 47s
CI / Clippy (push) Successful in 2m36s
build-prerelease / Build neuron-ada (push) Successful in 5m19s
CI / Test (push) Successful in 6m3s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 3m1s
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 3m32s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 3m47s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 59s
agent-0 sent a ~13k-token prompt + image; the TP vision prefill was
single-shot, so it tried to materialise activations for all 12,960
positions at once and OOM'd rank 1 mid-forward. Rank 1 died before
issuing its row-parallel AllReduce, stranding rank 0 on the collective
(it hung holding the pool lock). The text path survives the same size
because it chunks the prefill.

Chunk the vision prefill the same way:

- TpQwen3_5ForCausalLM::prefill_with_images_chunked encodes the image(s)
  once, then walks the pre-expanded prompt in prefill_chunk_tokens()
  windows, splicing the patch-embedding rows into whichever chunk(s)
  carry <|image_pad|> positions (pure-text chunks take the plain
  forward). Activation is bounded by the chunk, not the prompt.
- Every rank runs the identical chunk sequence (chunk_size threaded
  through GenerateStepWithImages / TpForwardLogitsWithImages /
  generate_step_with_images), so the per-chunk AllReduces stay paired
  across ranks with no extra sync — the KV cache accumulates via the
  growing offset, only the last chunk's logits are kept.

Pre-flight guard (validate_vision_prefill): even chunked, a long
prompt's KV cache can exhaust VRAM mid-forward, and on TP that hangs
the collective. Reject up front with a clean InsufficientVram when the
estimated footprint exceeds free VRAM, so a doomed request fails fast
instead of hanging the daemon. Heuristic + tunable
(NEURON_VISION_PREFILL_MB_PER_1K_TOKENS / _BASE_MB); default permissive
so the now-working 12,960-token case still passes. Applied to every
vision path (single-GPU + TP); single-GPU vision stays single-shot for
now, so the guard is its protection until it's chunked too.

Tests: pre-flight guard behaviour; RPC round-trip carries chunk_size.
The chunked forward is cuda-gated — CI CUDA type-check validates it.

Refs #16 / TP-vision. Operational note: a TP rank OOM still hangs the
daemon (needs restart); making a worker failure abort the leader's
collective is separate, broader TP hardening.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-04 17:21:36 +03:00
97 changed files with 12269 additions and 781 deletions

View File

@@ -1,11 +1,20 @@
name: build-prerelease
# Manually-dispatched workflow that builds CUDA-flavoured neuron binaries
# (and a single cortex binary), packages each as a Fedora RPM, signs
# them, and publishes to the `unstable` channel at rpm.lair.cafe.
# Builds CUDA-flavoured neuron binaries (and a single cortex binary),
# packages each as a Fedora RPM, signs them, and publishes to the
# `unstable` channel at rpm.lair.cafe.
#
# Trigger from the Gitea UI: Actions → build-prerelease → Run workflow.
# Optionally provide a `ref` to build from a non-default branch.
# Change-aware: the `prepare` job diffs HEAD against the git sha
# embedded in the most recently *published* unstable RPM (per package)
# and skips builds whose inputs didn't change. Docs-only commits build
# nothing; gateway-only commits skip the 3 CUDA builds (and, via
# deploy.yml's own check-update gate, the neuron restarts + model
# cold-loads). Diffing against the published sha — not the previous
# push — means a failed run can never cause a change to be missed.
#
# Lint (fmt+clippy) and test run here as parallel jobs and gate
# `publish`; ci.yml no longer runs on pushes to main (see its trigger
# comment), so the two workflows stop competing for the same runners.
#
# The published packages are versioned as e.g.
# helexa-neuron-blackwell-0.1.16-0.1.20260518T140530.gitabcdef0.fc43.x86_64
@@ -22,6 +31,7 @@ on:
push:
branches: [main]
# Manual dispatch still available to build from a non-main ref.
# Dispatched runs skip change detection and build everything.
workflow_dispatch:
inputs:
ref:
@@ -29,15 +39,15 @@ on:
required: false
default: ""
# Coalesce same-ref pushes: a newer push cancels the older in-flight
# run — the newest commit is the one we want on the fleet. The publish
# job keeps its own `rpm-publish` group (cancel=false) so an in-flight
# repo update is never interrupted. Runners are ephemeral (one VM per
# job) so concurrent runs no longer race on a shared workspace; the
# old shared `cortex-runner-pool` group with ci.yml is gone.
concurrency:
# Share the group with ci.yml so the two workflows can't run
# concurrently on the same `rust` runner (act reuses the workspace
# cache and races destroy each other's build files mid-compile).
# cancel-in-progress=false → workflows queue; if a newer push lands,
# the older run is still picked up by ci.yml's own ref-keyed
# concurrency (same group, queued).
group: cortex-runner-pool-${{ github.ref }}
cancel-in-progress: false
group: build-prerelease-${{ github.ref }}
cancel-in-progress: true
env:
CARGO_INCREMENTAL: "0"
@@ -45,13 +55,17 @@ env:
jobs:
prepare:
name: Resolve version stamps
name: Resolve version stamps + change detection
runs-on: rust
outputs:
version: ${{ steps.info.outputs.version }}
release: ${{ steps.info.outputs.release }}
short_sha: ${{ steps.info.outputs.short_sha }}
commit_timestamp: ${{ steps.info.outputs.commit_timestamp }}
build_cortex: ${{ steps.changes.outputs.build_cortex }}
build_neuron: ${{ steps.changes.outputs.build_neuron }}
build_bench: ${{ steps.changes.outputs.build_bench }}
check_rust: ${{ steps.changes.outputs.check_rust }}
steps:
- uses: actions/checkout@v4
with:
@@ -78,19 +92,228 @@ jobs:
echo "short_sha=${SHORT_SHA}" >> "$GITHUB_OUTPUT"
echo "commit_timestamp=${COMMIT_TIMESTAMP}" >> "$GITHUB_OUTPUT"
- id: changes
run: |
set -ux
# Default: build everything. Detection only ever narrows
# this, and any failure along the way (manifest unreachable,
# unparsable, sha not in history after a force-push) leaves
# the full build in place. Manual dispatches always build
# everything — predictable when building odd refs.
BUILD_CORTEX=true
BUILD_NEURON=true
BUILD_BENCH=true
CHECK_RUST=true
if [ "${GITHUB_EVENT_NAME}" = "push" ]; then
MANIFEST_URL="https://rpm.lair.cafe/fedora/43/x86_64/unstable/packages.json"
if curl -fsS --max-time 20 -o /tmp/packages.json "$MANIFEST_URL"; then
# Latest published sha per package, by buildTime.
base_for() {
python3 - "$1" <<'PY'
import json, re, sys
name = sys.argv[1]
try:
with open("/tmp/packages.json") as f:
pkgs = json.load(f)["packages"]
cands = [p for p in pkgs if p.get("name") == name]
if cands:
latest = max(cands, key=lambda p: p.get("buildTime", 0))
m = re.search(r"git\.?([0-9a-f]{7,40})", latest.get("release", ""))
if m:
print(m.group(1))
except Exception:
pass
PY
}
# true if no usable base, else true iff the diff since
# the published sha touches the given path pattern.
decide() {
local base="$1" pattern="$2"
if [ -z "$base" ] \
|| ! git cat-file -e "${base}^{commit}" 2>/dev/null \
|| ! git merge-base --is-ancestor "$base" HEAD 2>/dev/null; then
echo true; return
fi
if git diff --name-only "${base}..HEAD" | grep -qE "$pattern"; then
echo true
else
echo false
fi
}
# cortex-core is shared by both binaries; Cargo.{toml,lock}
# affect both; this workflow file affects both.
NEURON_RE='^crates/neuron/|^crates/cortex-core/|^Cargo\.toml$|^Cargo\.lock$|^rpm/helexa-neuron-prerelease\.spec$|^data/neuron|^neuron\.example\.toml$|^\.gitea/workflows/build-prerelease\.yml$'
CORTEX_RE='^crates/cortex-gateway/|^crates/cortex-cli/|^crates/cortex-core/|^Cargo\.toml$|^Cargo\.lock$|^rpm/cortex-prerelease\.spec$|^data/cortex|^cortex\.example\.toml$|^models\.example\.toml$|^\.gitea/workflows/build-prerelease\.yml$'
BENCH_RE='^crates/helexa-bench/|^crates/cortex-core/|^Cargo\.toml$|^Cargo\.lock$|^rpm/helexa-bench-prerelease\.spec$|^data/helexa-bench|^helexa-bench\.example\.toml$|^\.gitea/workflows/build-prerelease\.yml$'
# Any Rust change (incl. crates not packaged here, e.g.
# helexa-acp) still needs lint+test on main.
RUST_RE='\.rs$|^crates/|Cargo\.toml$|^Cargo\.lock$'
CORTEX_BASE=$(base_for cortex)
NEURON_BASE=$(base_for helexa-neuron-blackwell)
BENCH_BASE=$(base_for helexa-bench)
BUILD_CORTEX=$(decide "$CORTEX_BASE" "$CORTEX_RE")
BUILD_NEURON=$(decide "$NEURON_BASE" "$NEURON_RE")
BUILD_BENCH=$(decide "$BENCH_BASE" "$BENCH_RE")
if [ "$BUILD_CORTEX" = "true" ] || [ "$BUILD_NEURON" = "true" ] || [ "$BUILD_BENCH" = "true" ]; then
CHECK_RUST=true
else
CHECK_RUST=$(decide "$CORTEX_BASE" "$RUST_RE")
fi
fi
fi
echo "build_cortex=${BUILD_CORTEX}" >> "$GITHUB_OUTPUT"
echo "build_neuron=${BUILD_NEURON}" >> "$GITHUB_OUTPUT"
echo "build_bench=${BUILD_BENCH}" >> "$GITHUB_OUTPUT"
echo "check_rust=${CHECK_RUST}" >> "$GITHUB_OUTPUT"
echo "### change detection: build_cortex=${BUILD_CORTEX} build_neuron=${BUILD_NEURON} build_bench=${BUILD_BENCH} check_rust=${CHECK_RUST}"
# fmt + clippy + test moved here from ci.yml for main pushes so the
# two workflows stop queueing against each other (ci.yml's checks
# used to delay build-cortex by ~12 minutes on the shared runner
# pool). They run in parallel with the builds and gate `publish`,
# not the builds themselves — a clippy warning still can't reach the
# fleet, but it also doesn't serialize the pipeline.
lint:
name: Lint (fmt + clippy)
needs: prepare
if: needs.prepare.outputs.check_rust == 'true'
runs-on: rust
env:
RUSTC_WRAPPER: sccache
SCCACHE_BUCKET: sccache
SCCACHE_ENDPOINT: http://caveman.kosherinata.internal:9000
SCCACHE_REGION: auto
SCCACHE_S3_USE_SSL: "false"
AWS_ACCESS_KEY_ID: ${{ secrets.SCCACHE_S3_ACCESS_KEY }}
AWS_SECRET_ACCESS_KEY: ${{ secrets.SCCACHE_S3_SECRET_KEY }}
steps:
- uses: actions/checkout@v4
with:
ref: ${{ inputs.ref }}
- run: cargo fmt --check --all
# sccache failures come in two modes: transient races (a plain
# retry clears them) and a wedged/dead server, where every
# same-VM retry fails identically (sccache fatal error, ENOENT
# on its own tmp files). Escalate accordingly: retry → restart
# the server → final attempt uncached. A sick cache costs build
# time, never the run.
- name: Clippy (with sccache escalation)
run: |
for attempt in 1 2 3; do
echo "::group::clippy attempt ${attempt}"
if [ "${attempt}" -eq 3 ]; then
echo "final attempt: building without sccache"
export RUSTC_WRAPPER=""
fi
if cargo clippy --workspace -- -D warnings; then
echo "::endgroup::"
exit 0
fi
echo "::endgroup::"
echo "clippy failed on attempt ${attempt}"
if [ "${attempt}" -eq 1 ]; then
sccache --stop-server || true
sccache --start-server || true
fi
sleep 5
done
echo "clippy failed after 3 attempts"
exit 1
- run: sccache --show-stats || true
test:
name: Test
needs: prepare
if: needs.prepare.outputs.check_rust == 'true'
runs-on: rust
env:
RUSTC_WRAPPER: sccache
SCCACHE_BUCKET: sccache
SCCACHE_ENDPOINT: http://caveman.kosherinata.internal:9000
SCCACHE_REGION: auto
SCCACHE_S3_USE_SSL: "false"
AWS_ACCESS_KEY_ID: ${{ secrets.SCCACHE_S3_ACCESS_KEY }}
AWS_SECRET_ACCESS_KEY: ${{ secrets.SCCACHE_S3_SECRET_KEY }}
steps:
- uses: actions/checkout@v4
with:
ref: ${{ inputs.ref }}
# See the lint job for the escalation rationale.
- name: Test (with sccache escalation)
run: |
for attempt in 1 2 3; do
echo "::group::test attempt ${attempt}"
if [ "${attempt}" -eq 3 ]; then
echo "final attempt: building without sccache"
export RUSTC_WRAPPER=""
fi
if cargo test --workspace; then
echo "::endgroup::"
exit 0
fi
echo "::endgroup::"
echo "test failed on attempt ${attempt}"
if [ "${attempt}" -eq 1 ]; then
sccache --stop-server || true
sccache --start-server || true
fi
sleep 5
done
echo "test failed after 3 attempts"
exit 1
- run: sccache --show-stats || true
build-cortex:
name: Build cortex binary
needs: prepare
if: needs.prepare.outputs.build_cortex == 'true'
# runner-rust image already provides rust/cargo/clippy/rustfmt via
# dnf — no rustup install step needed.
runs-on: rust
env:
RUSTC_WRAPPER: sccache
SCCACHE_BUCKET: sccache
SCCACHE_ENDPOINT: http://caveman.kosherinata.internal:9000
SCCACHE_REGION: auto
SCCACHE_S3_USE_SSL: "false"
AWS_ACCESS_KEY_ID: ${{ secrets.SCCACHE_S3_ACCESS_KEY }}
AWS_SECRET_ACCESS_KEY: ${{ secrets.SCCACHE_S3_SECRET_KEY }}
steps:
- uses: actions/checkout@v4
with:
ref: ${{ inputs.ref }}
- name: Build cortex (release)
run: cargo build --release -p cortex-cli
# Escalation mirrors the lint/test jobs: retry → restart the
# sccache server → final attempt uncached. A sick cache costs
# build time, never the run.
- name: Build cortex (release, with sccache escalation)
run: |
for attempt in 1 2 3; do
echo "::group::build attempt ${attempt}"
if [ "${attempt}" -eq 3 ]; then
echo "final attempt: building without sccache"
export RUSTC_WRAPPER=""
fi
if cargo build --release -p cortex-cli; then
echo "::endgroup::"
sccache --show-stats || true
exit 0
fi
echo "::endgroup::"
echo "build failed on attempt ${attempt}"
if [ "${attempt}" -eq 1 ]; then
sccache --stop-server || true
sccache --start-server || true
fi
sleep 5
done
echo "build failed after 3 attempts"
exit 1
- name: Stage binary
run: |
@@ -104,9 +327,68 @@ jobs:
path: artifacts/cortex
retention-days: 1
build-bench:
name: Build helexa-bench binary
needs: prepare
if: needs.prepare.outputs.build_bench == 'true'
# Pure-Rust, non-CUDA binary — same runner as cortex.
runs-on: rust
env:
RUSTC_WRAPPER: sccache
SCCACHE_BUCKET: sccache
SCCACHE_ENDPOINT: http://caveman.kosherinata.internal:9000
SCCACHE_REGION: auto
SCCACHE_S3_USE_SSL: "false"
AWS_ACCESS_KEY_ID: ${{ secrets.SCCACHE_S3_ACCESS_KEY }}
AWS_SECRET_ACCESS_KEY: ${{ secrets.SCCACHE_S3_SECRET_KEY }}
steps:
- uses: actions/checkout@v4
with:
ref: ${{ inputs.ref }}
- name: Build helexa-bench (release, with sccache escalation)
run: |
# Stamp the SHA helexa-bench records as bench_sha against every
# run (option_env! in sweep.rs reads it at compile time).
export HELEXA_BUILD_SHA="$(git rev-parse HEAD)"
for attempt in 1 2 3; do
echo "::group::build attempt ${attempt}"
if [ "${attempt}" -eq 3 ]; then
echo "final attempt: building without sccache"
export RUSTC_WRAPPER=""
fi
if cargo build --release -p helexa-bench; then
echo "::endgroup::"
sccache --show-stats || true
exit 0
fi
echo "::endgroup::"
echo "build failed on attempt ${attempt}"
if [ "${attempt}" -eq 1 ]; then
sccache --stop-server || true
sccache --start-server || true
fi
sleep 5
done
echo "build failed after 3 attempts"
exit 1
- name: Stage binary
run: |
mkdir --parents artifacts
cp target/release/helexa-bench artifacts/helexa-bench
./artifacts/helexa-bench --version || true
- uses: actions/upload-artifact@v3
with:
name: bench-fc43
path: artifacts/helexa-bench
retention-days: 1
build-neuron:
name: Build neuron-${{ matrix.flavour }}
needs: prepare
if: needs.prepare.outputs.build_neuron == 'true'
strategy:
fail-fast: false
matrix:
@@ -117,34 +399,85 @@ jobs:
cuda_home: /usr/local/cuda-13.0
build_jobs: 8
nvcc_threads: 4
cargo_features: "cuda cudnn flash-attn"
cargo_features: "cuda cudnn"
- flavour: ada
compute_cap: "89"
runner: cuda-13.0
cuda_home: /usr/local/cuda-13.0
build_jobs: 8
nvcc_threads: 4
cargo_features: "cuda cudnn flash-attn"
cargo_features: "cuda cudnn"
- flavour: blackwell
compute_cap: "120"
runner: cuda-13.0
cuda_home: /usr/local/cuda-13.0
build_jobs: 8
nvcc_threads: 4
cargo_features: "cuda cudnn flash-attn"
cargo_features: "cuda cudnn"
runs-on: ${{ matrix.runner }}
env:
SCCACHE_BUCKET: sccache
SCCACHE_ENDPOINT: http://caveman.kosherinata.internal:9000
SCCACHE_REGION: auto
SCCACHE_S3_USE_SSL: "false"
AWS_ACCESS_KEY_ID: ${{ secrets.SCCACHE_S3_ACCESS_KEY }}
AWS_SECRET_ACCESS_KEY: ${{ secrets.SCCACHE_S3_SECRET_KEY }}
steps:
- uses: actions/checkout@v4
with:
ref: ${{ inputs.ref }}
# Escalation mirrors the lint/test jobs: retry → restart the
# sccache server → final attempt uncached.
#
# The CUDA image may or may not ship sccache — probe inside this
# step (NOT via GITHUB_ENV from a prior step, which this runner
# does not propagate; observed: probe step said "enabled", build
# ran unwrapped, server stats showed 4 compile requests). A
# missing binary degrades to an uncached build rather than
# failing cargo at `sccache rustc -vV`. The cache covers the
# ~600-crate host-side dep tree (the bulk of the 10-14 min
# build); rustc compilations are shared across all three
# flavours, so even one run seeds the next.
- name: Build neuron with CUDA (${{ matrix.flavour }})
run: |
set -eux
set -ux
if command -v sccache >/dev/null 2>&1; then
export RUSTC_WRAPPER=sccache
sccache --start-server 2>/dev/null || true
echo "sccache enabled"
else
echo "sccache not on PATH — building uncached"
fi
export PATH="${{ matrix.cuda_home }}/bin:${PATH}"
export LD_LIBRARY_PATH="${{ matrix.cuda_home }}/targets/x86_64-linux/lib:${{ matrix.cuda_home }}/lib64:${LD_LIBRARY_PATH:-}"
export LIBRARY_PATH="${{ matrix.cuda_home }}/targets/x86_64-linux/lib:${{ matrix.cuda_home }}/lib64:${LIBRARY_PATH:-}"
cargo build --release -p neuron --features "${{ matrix.cargo_features }}"
# Pin the build SHA neuron reports from GET /version. The git
# fallback in build.rs would also work on a full checkout, but
# injecting the exact checked-out commit is unambiguous under
# shallow/detached states and makes the artifact self-describing.
export HELEXA_BUILD_SHA="$(git rev-parse HEAD)"
for attempt in 1 2 3; do
echo "::group::build attempt ${attempt}"
if [ "${attempt}" -eq 3 ]; then
echo "final attempt: building without sccache"
export RUSTC_WRAPPER=""
fi
if cargo build --release -p neuron --features "${{ matrix.cargo_features }}"; then
echo "::endgroup::"
command -v sccache >/dev/null 2>&1 && sccache --show-stats || true
exit 0
fi
echo "::endgroup::"
echo "build failed on attempt ${attempt}"
if [ "${attempt}" -eq 1 ] && command -v sccache >/dev/null 2>&1; then
sccache --stop-server || true
sccache --start-server || true
fi
sleep 5
done
echo "build failed after 3 attempts"
exit 1
env:
CUDA_COMPUTE_CAP: ${{ matrix.compute_cap }}
CARGO_BUILD_JOBS: ${{ matrix.build_jobs }}
@@ -200,6 +533,42 @@ jobs:
path: ~/rpmbuild/RPMS/x86_64/*.rpm
retention-days: 7
package-bench:
name: Package helexa-bench RPM
needs: [prepare, build-bench]
runs-on: rpm
steps:
- uses: actions/checkout@v4
with:
ref: ${{ inputs.ref }}
- uses: actions/download-artifact@v3
with:
name: bench-fc43
path: artifacts/
- name: Build RPM
run: |
set -eux
rm -f ~/.rpmmacros
rpmdev-setuptree
cp artifacts/helexa-bench ~/rpmbuild/SOURCES/
cp data/helexa-bench.service ~/rpmbuild/SOURCES/
cp data/helexa-bench-sysusers.conf ~/rpmbuild/SOURCES/
cp helexa-bench.example.toml ~/rpmbuild/SOURCES/
cp LICENSE ~/rpmbuild/SOURCES/
rpmbuild -bb rpm/helexa-bench-prerelease.spec \
--define "bench_version ${{ needs.prepare.outputs.version }}" \
--define "bench_prerelease ${{ needs.prepare.outputs.release }}" \
--undefine dist \
--define "dist .fc43"
- uses: actions/upload-artifact@v3
with:
name: rpm-bench-fc43
path: ~/rpmbuild/RPMS/x86_64/*.rpm
retention-days: 7
package-neuron:
name: Package helexa-neuron-${{ matrix.flavour }} RPM
needs: [prepare, build-neuron]
@@ -247,7 +616,21 @@ jobs:
publish:
name: Publish to rpm.lair.cafe (unstable)
needs: [package-cortex, package-neuron]
needs: [lint, test, package-cortex, package-neuron, package-bench]
# Runs when at least one package was built and nothing failed.
# lint/test may be skipped (docs-only refs never get here because
# no packages build), but a real failure in any blocks the
# fleet from receiving the RPMs.
if: >-
${{
!cancelled()
&& (needs.lint.result == 'success' || needs.lint.result == 'skipped')
&& (needs.test.result == 'success' || needs.test.result == 'skipped')
&& (needs.package-cortex.result == 'success' || needs.package-neuron.result == 'success' || needs.package-bench.result == 'success')
&& needs.package-cortex.result != 'failure'
&& needs.package-neuron.result != 'failure'
&& needs.package-bench.result != 'failure'
}}
runs-on: rpm
concurrency:
group: rpm-publish

View File

@@ -1,21 +1,25 @@
name: CI
# Pushes to main are deliberately excluded: build-prerelease.yml runs
# its own lint/test jobs there (gating publish), and running both
# workflows on the same push made them queue against each other on the
# same runner labels — ~12 minutes of added latency per deploy. Feature
# branches, PRs to main, and release tags keep the full gate here.
on:
push:
branches: ["**"]
branches-ignore: [main]
tags: ["v*"]
pull_request:
branches: [main]
# Share a concurrency group with build-prerelease.yml so the two
# workflows don't race on the same `rust` runner workspace (act's
# /root/.cache/act/<hash>/hostexecutor/ is shared across concurrent
# jobs and one job's checkout step nukes another's in-flight build
# files). cancel-in-progress=false → they queue; same-ref pushes
# coalesce per workflow via cancel-in-progress on each.
# Coalesce same-ref pushes; a newer push supersedes the in-flight run.
# (The old shared `cortex-runner-pool` group with build-prerelease.yml
# is gone — the workflows no longer trigger on the same refs, and
# ephemeral one-VM-per-job runners removed the shared-workspace race
# that group existed to serialize.)
concurrency:
group: cortex-runner-pool-${{ github.ref }}
cancel-in-progress: false
group: ci-${{ github.ref }}
cancel-in-progress: true
env:
CARGO_INCREMENTAL: "0"
@@ -47,50 +51,64 @@ jobs:
runs-on: rust
steps:
- uses: actions/checkout@v4
# sccache occasionally fails with spurious race-condition errors;
# retrying the same invocation succeeds without code changes.
# Allow up to 3 attempts before declaring real failure.
- name: Clippy (with retry)
# sccache failures come in two modes: transient races (a plain
# retry clears them) and a wedged/dead server, where every
# same-VM retry fails identically. Escalate: retry → restart the
# server → final attempt uncached. A sick cache costs build
# time, never the run. Keep in sync with build-prerelease.yml.
- name: Clippy (with sccache escalation)
run: |
for attempt in 1 2 3; do
echo "::group::clippy attempt ${attempt}"
if [ "${attempt}" -eq 3 ]; then
echo "final attempt: building without sccache"
export RUSTC_WRAPPER=""
fi
if cargo clippy --workspace -- -D warnings; then
echo "::endgroup::"
exit 0
fi
echo "::endgroup::"
echo "clippy failed on attempt ${attempt}"
if [ "${attempt}" -lt 3 ]; then
sleep 5
if [ "${attempt}" -eq 1 ]; then
sccache --stop-server || true
sccache --start-server || true
fi
sleep 5
done
echo "clippy failed after 3 attempts"
exit 1
- run: sccache --show-stats
- run: sccache --show-stats || true
test:
name: Test
runs-on: rust
steps:
- uses: actions/checkout@v4
# See the clippy job for why this is retried.
- name: Test (with retry)
# See the clippy job for the escalation rationale.
- name: Test (with sccache escalation)
run: |
for attempt in 1 2 3; do
echo "::group::test attempt ${attempt}"
if [ "${attempt}" -eq 3 ]; then
echo "final attempt: building without sccache"
export RUSTC_WRAPPER=""
fi
if cargo test --workspace; then
echo "::endgroup::"
exit 0
fi
echo "::endgroup::"
echo "test failed on attempt ${attempt}"
if [ "${attempt}" -lt 3 ]; then
sleep 5
if [ "${attempt}" -eq 1 ]; then
sccache --stop-server || true
sccache --start-server || true
fi
sleep 5
done
echo "test failed after 3 attempts"
exit 1
- run: sccache --show-stats
- run: sccache --show-stats || true
# Type-check the CUDA-only code path. Borrow-check-only — we
# never run the tests here (the runner has no GPU). This catches
@@ -105,27 +123,34 @@ jobs:
cuda-check:
name: CUDA type-check
runs-on: cuda-13.0
# The workflow-level env sets `RUSTC_WRAPPER: sccache` for the
# `rust` runner (where fmt/clippy/test live and sccache is
# installed). The `cuda-13.0` runner doesn't have sccache on
# PATH, so inheriting the wrapper makes cargo bail with
# `could not execute process `sccache rustc -vV` (never executed)`
# before borrow-check even starts. Clear it locally. Also clear
# SCCACHE_* so cargo doesn't try to contact the cache (the
# remote auth headers come from secrets that aren't present on
# this runner either). Lose the cache, keep the gate.
# The workflow-level env sets `RUSTC_WRAPPER: sccache`
# unconditionally, which hard-fails cargo if the CUDA image
# doesn't ship sccache. Clear it at job level; the "Enable
# sccache when available" step opts back in only after probing
# for the binary. SCCACHE_*/AWS creds stay set — harmless when
# the wrapper is off, required when it's on.
env:
RUSTC_WRAPPER: ""
SCCACHE_BUCKET: ""
SCCACHE_ENDPOINT: ""
SCCACHE_REGION: ""
SCCACHE_S3_USE_SSL: ""
AWS_ACCESS_KEY_ID: ""
AWS_SECRET_ACCESS_KEY: ""
# candle-kernels' build script falls back to `nvidia-smi` for
# compute-cap detection when this is unset — and the GPU-less
# builder image doesn't ship nvidia-smi. Any valid cap works for
# a borrow-check; the real per-flavour caps live in
# build-prerelease.yml's matrix.
CUDA_COMPUTE_CAP: "86"
steps:
- uses: actions/checkout@v4
- name: cargo check --features cuda (with retry)
# sccache is probed inside this step (NOT via GITHUB_ENV from a
# prior step — this runner doesn't propagate it; see
# build-prerelease.yml for the observed failure).
- name: cargo check --features cuda (with sccache escalation)
run: |
if command -v sccache >/dev/null 2>&1; then
export RUSTC_WRAPPER=sccache
sccache --start-server 2>/dev/null || true
echo "sccache enabled"
else
echo "sccache not on PATH — building uncached"
fi
# act launches the step shell without /etc/profile, so the
# gitea_runner user's inherited PATH lacks /usr/local/cuda-13.0/bin.
# cudarc's build.rs:157 shells out to `nvcc --version` (because
@@ -135,17 +160,25 @@ jobs:
export PATH="/usr/local/cuda-13.0/bin:${PATH}"
export LD_LIBRARY_PATH="/usr/local/cuda-13.0/targets/x86_64-linux/lib:/usr/local/cuda-13.0/lib64:${LD_LIBRARY_PATH:-}"
export LIBRARY_PATH="/usr/local/cuda-13.0/targets/x86_64-linux/lib:/usr/local/cuda-13.0/lib64:${LIBRARY_PATH:-}"
# Escalation mirrors the lint/test jobs: plain retry →
# sccache server restart → final attempt uncached.
for attempt in 1 2 3; do
echo "::group::cuda-check attempt ${attempt}"
if [ "${attempt}" -eq 3 ]; then
echo "final attempt: building without sccache"
export RUSTC_WRAPPER=""
fi
if cargo check -p neuron --features cuda --all-targets; then
echo "::endgroup::"
exit 0
fi
echo "::endgroup::"
echo "cuda-check failed on attempt ${attempt}"
if [ "${attempt}" -lt 3 ]; then
sleep 5
if [ "${attempt}" -eq 1 ] && command -v sccache >/dev/null 2>&1; then
sccache --stop-server || true
sccache --start-server || true
fi
sleep 5
done
echo "cuda-check failed after 3 attempts"
exit 1
@@ -349,6 +382,6 @@ jobs:
echo "Nothing to commit for ${VERSION}"
else
git commit -m "chore: bump version to ${VERSION}"
git remote set-url origin "https://gitea-actions:${GITEA_TOKEN}@git.lair.cafe/helexa/cortex.git"
git remote set-url origin "https://gitea-actions:${GITEA_TOKEN}@git.lair.cafe/${{ github.repository }}.git"
git push origin HEAD:main
fi

View File

@@ -0,0 +1,134 @@
name: deploy-dev
# Fast-path iteration deploy for a SINGLE neuron host: build one CUDA
# flavour, copy the raw binary to the host, restart neuron.service.
# Skips the other two flavours, all RPM packaging, signing, repo
# publish, and dnf — push-to-testable drops from ~20 min to roughly
# one CUDA build plus a service restart.
#
# This is a DEV convenience, not a release path:
# - the binary lands at /usr/bin/neuron *outside* RPM ownership;
# the next regular deploy.yml run reconciles the host back to the
# packaged binary (dnf sees the newer RPM and reinstalls). `rpm -V
# helexa-neuron-<flavour>` flagging a modified /usr/bin/neuron in
# the interim is expected.
# - nothing is published; other hosts are untouched.
# - requires the `install` sudoers rule from
# asset/sudoers.d/neuron-host.conf (re-run script/infra-setup.sh
# after updating it).
#
# Trigger from the Gitea UI: Actions → deploy-dev → Run workflow,
# pick the target host. Defaults to the ref you dispatch from, so it
# works from feature branches without touching main.
on:
workflow_dispatch:
inputs:
target:
description: "neuron host to deploy to"
required: true
type: choice
options: [beast, benjy, quadbrat]
default: beast
# One dev deploy at a time; a newer dispatch for the same host wins.
concurrency:
group: deploy-dev-${{ inputs.target }}
cancel-in-progress: true
env:
CARGO_INCREMENTAL: "0"
CARGO_TERM_COLOR: "always"
jobs:
build:
name: Build neuron (${{ inputs.target }})
runs-on: cuda-13.0
outputs:
flavour: ${{ steps.map.outputs.flavour }}
steps:
- uses: actions/checkout@v4
# host → flavour → compute cap. Keep in sync with the
# build-neuron matrix in build-prerelease.yml and the
# deploy-neurons matrix in deploy.yml.
- id: map
run: |
case "${{ inputs.target }}" in
beast) flavour=blackwell cap=120 ;;
benjy) flavour=ada cap=89 ;;
quadbrat) flavour=ampere cap=86 ;;
*) echo "unknown target ${{ inputs.target }}"; exit 1 ;;
esac
echo "flavour=${flavour}" >> "$GITHUB_OUTPUT"
echo "cap=${cap}" >> "$GITHUB_OUTPUT"
- name: Build neuron with CUDA
run: |
set -eux
export PATH="/usr/local/cuda-13.0/bin:${PATH}"
export LD_LIBRARY_PATH="/usr/local/cuda-13.0/targets/x86_64-linux/lib:/usr/local/cuda-13.0/lib64:${LD_LIBRARY_PATH:-}"
export LIBRARY_PATH="/usr/local/cuda-13.0/targets/x86_64-linux/lib:/usr/local/cuda-13.0/lib64:${LIBRARY_PATH:-}"
cargo build --release -p neuron --features "cuda cudnn"
env:
CUDA_COMPUTE_CAP: ${{ steps.map.outputs.cap }}
CARGO_BUILD_JOBS: "8"
NVCC_THREADS: "4"
- name: Stage binary
run: |
mkdir --parents artifacts
cp target/release/neuron artifacts/neuron-dev
file artifacts/neuron-dev
- uses: actions/upload-artifact@v3
with:
name: neuron-dev-${{ inputs.target }}
path: artifacts/neuron-dev
retention-days: 1
deploy:
name: Deploy to ${{ inputs.target }}
needs: build
runs-on: fedora-43
env:
DEPLOY_KEY: |
${{ secrets.RSYNC_SSH_KEY }}
TARGET_HOST: ${{ inputs.target }}.hanzalova.internal
steps:
- name: SSH init
run: |
mkdir -p ~/.ssh
echo "${DEPLOY_KEY}" > ~/.ssh/id_ed25519
chmod 600 ~/.ssh/id_ed25519
ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=accept-new \
"gitea_ci@${TARGET_HOST}" 'hostname -f'
- uses: actions/download-artifact@v3
with:
name: neuron-dev-${{ inputs.target }}
path: artifacts/
- name: Copy binary to host
run: |
scp artifacts/neuron-dev "gitea_ci@${TARGET_HOST}:/var/lib/gitea_ci/neuron-dev"
- name: Install binary and restart neuron.service
run: |
ssh "gitea_ci@${TARGET_HOST}" '
set -eu
if systemctl is-active --quiet neuron.service; then
sudo /usr/bin/systemctl stop neuron.service
fi
# Exact command form required by the sudoers rule in
# asset/sudoers.d/neuron-host.conf — change both together.
sudo /usr/bin/install -o root -g root -m 0755 /var/lib/gitea_ci/neuron-dev /usr/bin/neuron
sudo /usr/bin/systemctl start neuron.service
rm -f /var/lib/gitea_ci/neuron-dev'
- name: Capture neuron.service startup journal
if: always()
run: |
sleep 10
ssh "gitea_ci@${TARGET_HOST}" \
'journalctl --unit neuron.service -I --no-pager'

View File

@@ -7,6 +7,12 @@ name: deploy
# point the new RPMs are live on rpm.lair.cafe/unstable), and also
# re-runnable manually from the Gitea UI.
#
# Each host self-gates: if dnf sees no newer package than what is
# installed, the service is left alone — no stop, no restart, no model
# cold-load. Combined with build-prerelease's change detection this
# means a docs- or gateway-only push never restarts the neurons (a
# neuron restart costs ~5 min of 27B cold-load, see issue #1).
#
# Per-host one-time setup (gitea_ci user, authorized_keys, scoped
# sudoers drop-in) lives in script/infra-setup.sh — run that once per
# host before this workflow can succeed.
@@ -48,27 +54,42 @@ jobs:
ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=accept-new \
gitea_ci@hanzalova.internal 'hostname -f'
- name: Stop cortex.service
# Gating compares `rpm -q` against the packages.json manifest the
# publish job maintains — NOT unprivileged `dnf check-update`,
# which proved unreliable as the gitea_ci user (hung on metadata
# locks on one host, silently reported "no updates" on others).
# An unreadable/unparsable manifest fails open: deploy proceeds.
- name: Deploy cortex (skips when already current)
run: |
ssh gitea_ci@hanzalova.internal '
ssh gitea_ci@hanzalova.internal 'bash -s' <<'DEPLOY'
set -eu
pkg=cortex
installed=$(rpm -q --qf '%{VERSION}-%{RELEASE}' "${pkg}" 2>/dev/null || echo "not-installed")
latest=$(curl -fsS --max-time 15 "https://rpm.lair.cafe/fedora/43/x86_64/unstable/packages.json" 2>/dev/null \
| python3 -c '
import json, sys
name = sys.argv[1]
cands = [p for p in json.load(sys.stdin)["packages"] if p.get("name") == name]
if cands:
p = max(cands, key=lambda p: p.get("buildTime", 0))
print(p["version"] + "-" + p["release"])
' "${pkg}" 2>/dev/null || true)
if [ -n "${latest}" ] && [ "${latest}" = "${installed}" ]; then
echo "${pkg}-${installed} already current — leaving service untouched"
exit 0
fi
echo "installed=${installed} published=${latest:-unknown} — deploying"
if systemctl is-active --quiet cortex.service; then
sudo /usr/bin/systemctl stop cortex.service
fi'
- name: Install / upgrade cortex from rpm.lair.cafe/unstable
run: |
ssh gitea_ci@hanzalova.internal '
if rpm -q cortex >/dev/null 2>&1; then
fi
if rpm -q "${pkg}" >/dev/null 2>&1; then
sudo /usr/bin/dnf upgrade --refresh --allowerasing -y cortex
else
sudo /usr/bin/dnf install --refresh --allowerasing -y cortex
fi'
- name: Start cortex.service
run: |
ssh gitea_ci@hanzalova.internal '
fi
sudo /usr/bin/systemctl daemon-reload
sudo /usr/bin/systemctl start cortex.service'
sudo /usr/bin/systemctl start cortex.service
DEPLOY
# Wait for the service to either come up or wedge, then capture
# the latest-invocation journal. Runs even on prior failure so a
@@ -90,12 +111,19 @@ jobs:
fail-fast: false
matrix:
include:
# load_timeout: how long to wait for default_models to finish
# loading after a restart. beast cold-loads Qwen3.6-27B Q6K
# TP=2 (~5-6 min typical, see #1); benjy/quadbrat load small
# single-GPU models in well under a minute.
- host: beast.hanzalova.internal
flavour: blackwell
load_timeout: 900
- host: benjy.hanzalova.internal
flavour: ada
load_timeout: 300
- host: quadbrat.hanzalova.internal
flavour: ampere
load_timeout: 300
steps:
- name: SSH init
run: |
@@ -105,21 +133,105 @@ jobs:
ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=accept-new \
gitea_ci@${{ matrix.host }} 'hostname -f'
- name: Stop neuron.service
# See deploy-cortex for why gating uses the publish manifest and
# not unprivileged `dnf check-update`.
- name: Deploy helexa-neuron-${{ matrix.flavour }} (skips when already current)
run: |
ssh gitea_ci@${{ matrix.host }} '
ssh gitea_ci@${{ matrix.host }} 'bash -s' <<'DEPLOY'
set -eu
pkg=helexa-neuron-${{ matrix.flavour }}
installed=$(rpm -q --qf '%{VERSION}-%{RELEASE}' "${pkg}" 2>/dev/null || echo "not-installed")
latest=$(curl -fsS --max-time 15 "https://rpm.lair.cafe/fedora/43/x86_64/unstable/packages.json" 2>/dev/null \
| python3 -c '
import json, sys
name = sys.argv[1]
cands = [p for p in json.load(sys.stdin)["packages"] if p.get("name") == name]
if cands:
p = max(cands, key=lambda p: p.get("buildTime", 0))
print(p["version"] + "-" + p["release"])
' "${pkg}" 2>/dev/null || true)
if [ -n "${latest}" ] && [ "${latest}" = "${installed}" ]; then
echo "${pkg}-${installed} already current — leaving service untouched"
exit 0
fi
echo "installed=${installed} published=${latest:-unknown} — deploying"
if systemctl is-active --quiet neuron.service; then
sudo /usr/bin/systemctl stop neuron.service
fi'
- name: Install / upgrade helexa-neuron-${{ matrix.flavour }}
run: |
ssh gitea_ci@${{ matrix.host }} "
if rpm -q helexa-neuron-${{ matrix.flavour }} >/dev/null 2>&1; then
sudo /usr/bin/dnf upgrade --refresh --allowerasing -y helexa-neuron-${{ matrix.flavour }}
fi
if rpm -q "${pkg}" >/dev/null 2>&1; then
sudo /usr/bin/dnf upgrade --refresh --allowerasing -y "${pkg}"
else
sudo /usr/bin/dnf install --refresh --allowerasing -y helexa-neuron-${{ matrix.flavour }}
fi"
sudo /usr/bin/dnf install --refresh --allowerasing -y "${pkg}"
fi
sudo /usr/bin/systemctl daemon-reload
sudo /usr/bin/systemctl start neuron.service
# ── Post-deploy validation ────────────────────────────────
# A deploy only goes green if the neuron (a) finishes loading
# its default models and (b) answers a trivial prompt like an
# LLM should. Catches the class of bug where the binary
# starts fine but model load or inference is broken — which
# previously surfaced only when a human noticed. The wait
# polls /health activation (the structured source of the
# "loaded default model" journal line, plus per-model failure
# detail); the journal-capture step below still runs for
# forensics either way.
load_timeout=${{ matrix.load_timeout }}
echo "waiting for default models (timeout ${load_timeout}s)"
deadline=$(( $(date +%s) + load_timeout ))
health=""
while :; do
health=$(curl -fsS --max-time 5 http://localhost:13131/health 2>/dev/null || true)
state=$(printf %s "${health}" | python3 -c '
import json, sys
try:
print(json.load(sys.stdin).get("activation", {}).get("state", ""))
except Exception:
print("")
')
if [ "${state}" = "ready" ]; then
break
fi
if [ "$(date +%s)" -ge "${deadline}" ]; then
echo "FAIL: activation not ready within ${load_timeout}s (last state: ${state:-unreachable})"
exit 1
fi
sleep 10
done
model=$(printf %s "${health}" | python3 -c '
import json, sys
a = json.load(sys.stdin).get("activation", {})
failed = a.get("failed", [])
if failed:
for f in failed:
msg = "FAILED " + str(f.get("model_id")) + ": " + str(f.get("error", ""))[:400]
sys.stderr.write(msg + chr(10))
sys.exit(1)
completed = a.get("completed", [])
print(completed[0] if completed else "")
')
if [ -z "${model}" ]; then
echo "no default models configured — skipping LLM probe"
exit 0
fi
echo "LLM probe against ${model}"
probe_body=$(printf '{"model":"%s","messages":[{"role":"user","content":"Reply with exactly one word: pineapple"}],"max_tokens":512,"temperature":0}' "${model}")
resp=$(curl -fsS --max-time 180 -H "content-type: application/json" \
-d "${probe_body}" http://localhost:13131/v1/chat/completions) || {
echo "FAIL: probe request errored"
exit 1
}
if printf %s "${resp}" | grep -qi pineapple; then
echo "LLM probe passed"
else
echo "FAIL: probe response missing expected token"
printf %s "${resp}" | head -c 2000
echo
exit 1
fi
DEPLOY
- name: Ensure firewalld allows helexa-neuron
run: |
@@ -129,12 +241,6 @@ jobs:
sudo /usr/bin/firewall-cmd --reload
fi'
- name: Start neuron.service
run: |
ssh gitea_ci@${{ matrix.host }} '
sudo /usr/bin/systemctl daemon-reload
sudo /usr/bin/systemctl start neuron.service'
# Wait for the service to either come up or wedge, then capture
# the latest-invocation journal. Runs even on prior failure so a
# failed start step still leaves a usable record in the deploy log.

View File

@@ -1,16 +1,26 @@
# CLAUDE.md — cortex
# CLAUDE.md — helexa
## Project overview
cortex is a Rust reverse-proxy that sits in front of multiple
mistral.rs inference nodes and presents a unified OpenAI + Anthropic
compatible API surface. It handles model routing, lifecycle management
(load/unload/evict), request translation, and metrics collection.
helexa is a self-hosted LLM serving stack for multi-node GPU inference
clusters. It has two components:
- **cortex** — the per-operator control plane and LLM proxy. A Rust
reverse-proxy that sits in front of the fleet and presents a unified
OpenAI + Anthropic compatible API surface. It handles model routing,
lifecycle management (load/unload/evict), request translation, and
metrics collection.
- **neuron** — the per-host LLM harness. One instance runs on every GPU
host, serving candle-based in-process inference and managing local
hardware discovery and model lifecycle.
(Historical note: cortex originally proxied to mistral.rs nodes; neuron
replaced that — see the 2026-05-18 candle-native addendum below.)
## Repository layout
```
cortex/
helexa/
├── Cargo.toml # workspace root
├── cortex.toml # example gateway config
├── README.md
@@ -548,7 +558,7 @@ and the hardcoded `vram_mb` per node.
## Revised repository layout
```
cortex/
helexa/
├── Cargo.toml
├── cortex.toml # gateway config (neurons only)
├── models.toml # model catalogue
@@ -754,3 +764,39 @@ Landed in four PRs:
from Phases 2/3 deleted; `SendComm` newtype no longer needed in the
load path. `grep -rn spawn_blocking crates/neuron/src/harness/`
returns only deliberate CPU-fallback hits after this PR.
## 2026-06-13 addendum: build metadata + helexa-bench
Two coupled additions so fleet performance can be tracked automatically
across neuron updates instead of by hand-running `script/bench.py` and
editing `doc/benchmarks.md`.
**neuron build metadata + `GET /version`.** neuron's `build.rs` now also
captures build identity (`HELEXA_GIT_SHA` — preferring a CI/RPM-injected
`HELEXA_BUILD_SHA`, falling back to git, else `unknown` — plus dirty
flag, build timestamp, rustc version, profile, enabled cargo features,
and a best-effort `candle-core` version from `Cargo.lock`). These are
exposed as `cortex_core::build_info::BuildInfo` (new module) from a new
`GET /version` endpoint (`neuron/src/version.rs`, wired in `api.rs`) and
in clap's `--version` long form. The SHA is injected in CI
(`build-prerelease.yml` build-neuron step: `export HELEXA_BUILD_SHA=$(git
rev-parse HEAD)`) and via `--define helexa_commit` in the source-build
spec, so tarball-built RPMs report the real SHA. `/version` is now the
canonical "which build is live" probe (supersedes the per-host RPM-sha
check in the fleet-validation flow).
**`crates/helexa-bench`** — a new binary: a continuous, version-aware
benchmark harness (one systemd unit, typically on the metrics host). It
hits each neuron **directly** on `:13131`, exercises each **warm**
(`status == "loaded"`) model with an extensible `Scenario` suite (phase
1: the chat-latency family ported verbatim from `bench.py` — synthetic
128/4096-tok prompts, `/no_think`, streamed TTFT + decode-window
tok/s), and records each run into a SQLite system-of-record stamped with
the neuron's full `BuildInfo`. The loop is **version-aware**: it skips
any (target, build SHA, model, scenario) cell already at
`samples_per_version`, so a steady fleet costs only cheap `/version` +
`/models` polls until a new SHA ships. `helexa-bench report` regenerates
the `benchmarks.md`-style table from the DB. `kind = "openai"` targets
(mistral.rs/llama.cpp comparison) are scaffolded but not yet wired.
Packaged as the `helexa-bench` RPM (prebuilt-binary spec, outbound-only
so no firewalld service) via the same `build-prerelease.yml` pipeline.

84
Cargo.lock generated
View File

@@ -905,8 +905,7 @@ dependencies = [
[[package]]
name = "cudarc"
version = "0.19.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1cea5f10a99e025c1b44ae2354c2d8326b25ddbd0baf76bde8e55cfd4018a2cc"
source = "git+https://github.com/grenade/cudarc?rev=63327a256059f8252641ae46c6bb9eefe707f382#63327a256059f8252641ae46c6bb9eefe707f382"
dependencies = [
"float8",
"half",
@@ -1218,6 +1217,18 @@ dependencies = [
"pin-project-lite",
]
[[package]]
name = "fallible-iterator"
version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649"
[[package]]
name = "fallible-streaming-iterator"
version = "0.1.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a"
[[package]]
name = "fancy-regex"
version = "0.17.0"
@@ -1251,8 +1262,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8cb01cd46b0cf372153850f4c6c272d9cbea2da513e07538405148f95bd789f3"
dependencies = [
"atomic",
"parking_lot",
"pear",
"serde",
"tempfile",
"toml",
"uncased",
"version_check",
@@ -1808,6 +1821,15 @@ version = "0.12.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888"
[[package]]
name = "hashbrown"
version = "0.14.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1"
dependencies = [
"ahash",
]
[[package]]
name = "hashbrown"
version = "0.15.5"
@@ -1836,6 +1858,15 @@ version = "0.17.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4f467dd6dccf739c208452f8014c75c18bb8301b050ad1cfb27153803edb0f51"
[[package]]
name = "hashlink"
version = "0.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6ba4ff7128dee98c7dc9794b6a411377e1404dba1c97deb8d1a55297bd25d8af"
dependencies = [
"hashbrown 0.14.5",
]
[[package]]
name = "heck"
version = "0.5.0"
@@ -1866,6 +1897,29 @@ dependencies = [
"url",
]
[[package]]
name = "helexa-bench"
version = "0.1.16"
dependencies = [
"anyhow",
"async-trait",
"axum",
"chrono",
"clap",
"cortex-core",
"eventsource-stream",
"figment",
"futures",
"reqwest",
"rusqlite",
"serde",
"serde_json",
"tokio",
"tokio-stream",
"tracing",
"tracing-subscriber",
]
[[package]]
name = "hermit-abi"
version = "0.5.2"
@@ -2358,6 +2412,17 @@ dependencies = [
"libc",
]
[[package]]
name = "libsqlite3-sys"
version = "0.30.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2e99fb7a497b1e3339bc746195567ed8d3e24945ecd636e3619d20b9de9e9149"
dependencies = [
"cc",
"pkg-config",
"vcpkg",
]
[[package]]
name = "linux-raw-sys"
version = "0.12.1"
@@ -2616,6 +2681,7 @@ dependencies = [
"image",
"minijinja",
"minijinja-contrib",
"rayon",
"reqwest",
"safetensors 0.7.0",
"serde",
@@ -3431,6 +3497,20 @@ dependencies = [
"syn",
]
[[package]]
name = "rusqlite"
version = "0.32.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7753b721174eb8ff87a9a0e799e2d7bc3749323e773db92e0984debb00019d6e"
dependencies = [
"bitflags",
"fallible-iterator",
"fallible-streaming-iterator",
"hashlink",
"libsqlite3-sys",
"smallvec",
]
[[package]]
name = "rustc-hash"
version = "2.1.2"

View File

@@ -6,13 +6,14 @@ members = [
"crates/cortex-cli",
"crates/neuron",
"crates/helexa-acp",
"crates/helexa-bench",
]
[workspace.package]
version = "0.1.16"
edition = "2024"
license = "GPL-3.0-or-later"
repository = "https://git.lair.cafe/helexa/cortex"
repository = "https://git.lair.cafe/helexa/helexa"
[workspace.dependencies]
# async runtime
@@ -61,3 +62,12 @@ eventsource-stream = "0.2"
# workspace crates
cortex-core = { path = "crates/cortex-core" }
cortex-gateway = { path = "crates/cortex-gateway" }
# Patched cudarc (affects neuron's 0.19.x only; candle's 0.17.x is
# untouched since the fork is 0.19.7 and doesn't satisfy a 0.17 req). Adds
# Comm::abort / get_async_error / raw comm() — needed for #17 Stage 2 TP
# hang-recovery (abort a wedged collective from another thread, then
# rebuild the comm). Pinned to a fork revision pending upstream review
# (grenade/cudarc @ nccl-comm-abort).
[patch.crates-io]
cudarc = { git = "https://github.com/grenade/cudarc", rev = "63327a256059f8252641ae46c6bb9eefe707f382" }

190
README.md
View File

@@ -1,25 +1,68 @@
# cortex
# helexa
A Rust reverse-proxy and fleet management layer for multi-node GPU inference
clusters. Cortex sits in front of one or more `neuron` daemons (each running
candle-based inference on a local GPU host) and presents a unified OpenAI +
Anthropic compatible API surface.
**Near-frontier AI for mortals.**
## Problem
helexa is a self-hosted LLM serving stack, written in Rust, for people
who run open-weight models on their own consumer GPUs. It has two
components:
Running local LLMs across multiple GPU nodes (different VRAM tiers, different
model affinities) requires a unified API surface that:
- **cortex** — the per-operator control plane and LLM proxy. It sits in
front of your GPU fleet and presents a unified OpenAI + Anthropic
compatible API surface, handling model routing, lifecycle management
(load / unload / evict), request translation, and metrics.
- **neuron** — the per-host LLM harness. One instance runs on every GPU
host, serving candle-based in-process inference and managing local
hardware discovery and model lifecycle.
- Presents a **single `/v1/models` catalogue** merging every model that can be
served by any neuron in the fleet.
- **Routes requests** to the correct node based on where a model is loaded
(or can be loaded), handling cold-load and eviction transparently.
- Manages **model lifecycle** — load on demand, unload cold models, pin
critical ones — by calling each neuron's `/models/{load,unload}` API.
- Translates between **OpenAI and Anthropic** request/response envelopes so
every client speaks whichever dialect it prefers.
- Captures **per-request metrics** (tokens, tok/s, TTFT, latency) and exposes
them as Prometheus counters/histograms.
## Why
Two principles constrain everything in this repository:
1. **Frontier or close to it.** helexa serves the open-weight models
that get nearest to frontier capability — not every architecture
ever published.
2. **Consumer hardware.** Everything must run on the cards mortals can
actually buy: a 3060 here, a 4090 there, a 5090 if you got lucky.
Mixed VRAM tiers across mismatched boxes are the expected topology,
not a degraded case.
GPU acquisition is harder than it was a year ago, and the gap between
what cloud providers charge and what your own silicon costs keeps
widening. The intersection of those two principles — near-frontier
models, squeezed onto hardware you own — is helexa's entire niche.
The secondary objective is **predictable consumption**. If you own the
hardware, your tooling shouldn't break because a cloud provider changed
billing, deprecated a model, or reshaped an API. cortex's OpenAI and
Anthropic surfaces are a stability contract: point your editor, agent,
or CLI at it once, and it keeps working.
## What helexa is not
This is an intentionally different path from vLLM, SGLang, and peers —
not a smaller version of them. Out of scope, permanently:
- Any-model breadth. Architectures are ported because they're at or
near the frontier, not to complete a compatibility matrix.
- Datacenter-class scheduling. No sophisticated continuous-batching /
paged-attention machinery — the workload is a handful of operators
and their agents, not 200 QPS.
- Wrapping external inference engines. neuron builds directly on
[candle](https://github.com/huggingface/candle); every model
architecture it serves is implemented in this repository, ported
against the HuggingFace reference.
One thing that is *not* a principle: CUDA exclusivity. All high-end
consumer hardware is in scope. helexa is CUDA-only today because
that's the hardware on the bench — nothing ships untested — and ROCm
or other consumer accelerators join as soon as there's real hardware
to build against.
In scope, and where the engineering effort goes: aggressive
quantization (GGUF Q4_K_M / Q6_K / Q8_0), NCCL tensor parallelism
across heterogeneous consumer GPUs, careful CUDA failure handling, and
single-request latency — the performance that one operator at a
keyboard actually feels.
## Architecture
@@ -29,7 +72,7 @@ model affinities) requires a unified API surface that:
└──────┬───────┘ └─────┬────┘ └──────┬─────┘ └──────┬─────┘
│ │ │ │
└────────────────┴──────┬───────┴───────────────┘
OpenAI + Anthropic APIs
┌──────────▼──────────┐
│ cortex │
│ (cortex-gateway) │
@@ -46,40 +89,59 @@ model affinities) requires a unified API surface that:
private network (.internal)
```
cortex discovers each neuron's hardware (devices, VRAM, compute
capability) at runtime and matches it against a model catalogue
(`models.toml`) to decide placement: which models fit where, what to
evict when VRAM is tight, where to route a request right now. Adding a
GPU host to the fleet is one `[[neurons]]` entry — no device specs in
config.
### Crates
| Crate | Purpose |
|---|---|
| `cortex-core` | Shared types: config, node/model state, metrics, OpenAI/Anthropic envelopes, harness trait, discovery types |
| `cortex-gateway` | Axum HTTP server: proxy, router, evictor, poller, metrics exporter |
| `neuron` | Per-node daemon: GPU discovery, in-process candle inference, model lifecycle API |
| `neuron` | Per-host daemon: GPU discovery, in-process candle inference, NCCL tensor parallelism, model lifecycle API |
| `cortex-cli` | CLI entrypoint (`cortex serve`, `cortex status`, etc.) |
| `helexa-acp` | Agent Client Protocol bridge — connects ACP editors (Zed, etc.) to any OpenAI-compatible endpoint, cortex by default |
## Node setup
## The engine
Each GPU node runs `neuron` (listening on `:13131`). Neuron uses
huggingface/candle for in-process inference — there is no external
inference subprocess to manage.
neuron runs inference in-process on candle — there is no external
inference server to babysit. The parts that earn their keep:
Inside the daemon, every CUDA device gets one dedicated OS thread
(named `cuda-dev-N`) that owns the device's CUDA context for the
daemon's lifetime. Model loads, forward passes, KV-cache resets,
NCCL collectives, VRAM queries, and unloads all route through that
thread via a job channel; tensors never escape it alive. This pins
context binding to a known thread, makes the CUDA Drop contract
structurally safe, and isolates driver-error poisoning to one worker
rather than the whole process. See `CLAUDE.md` for the design
rationale and `crates/neuron/src/harness/device_worker/` for the code.
- **Per-device worker threads.** Every CUDA device gets one dedicated
OS thread that owns its CUDA context for the daemon's lifetime. All
loads, forward passes, KV-cache resets, NCCL collectives, VRAM
queries, and unloads route through it; tensors never escape it
alive. Context binding is pinned to a known thread, the CUDA `Drop`
contract is structurally safe, and a driver error poisons one worker
— visibly — instead of hanging the whole process.
- **Tensor parallelism on consumer cards.** Megatron-style row/column
parallel layers with NCCL all-reduce, spanning the mismatched GPUs
you actually have. A step watchdog aborts wedged collectives instead
of letting a request hang forever.
- **Current model focus: the Qwen3 family** — dense and GGUF-quantized,
including the hybrid linear-attention (Gated DeltaNet) generation.
Vision support is in progress. Each architecture is ported against
its HuggingFace reference implementation.
The neuron RPM (`helexa-neuron`) ships a systemd unit:
See `CLAUDE.md` for design rationale and
`crates/neuron/src/harness/device_worker/` for the worker narrative.
## Install
Pre-built RPMs for Fedora:
```sh
dnf copr enable helexa/helexa
dnf install helexa-neuron
systemctl enable --now neuron
dnf install cortex # on the gateway host
dnf install helexa-neuron # on each GPU host
systemctl enable --now cortex # or neuron, respectively
```
## Gateway config
## Configure
```toml
# /etc/cortex/cortex.toml
@@ -100,29 +162,10 @@ name = "benjy"
endpoint = "http://benjy.internal:13131"
```
Model placement profiles live in `models.toml` — see `models.example.toml`.
Model placement profiles (VRAM requirements, quant, device minimums,
pinning) live in `models.toml` — see `models.example.toml`.
## Building
```sh
cargo build --release
```
## CI
Every push triggers format, lint, and test checks. Ensure these pass
locally before pushing:
```sh
cargo fmt --check --all # must be clean
cargo clippy --workspace -- -D warnings # warnings are errors
cargo test --workspace # all tests must pass
```
Tagged releases (`v*`) additionally build SRPMs for both `cortex` and
`helexa-neuron` and publish to COPR.
## Running
## Run
```sh
# start the gateway
@@ -131,10 +174,37 @@ cortex serve --config /etc/cortex/cortex.toml
# check fleet status
cortex status
# list all models across nodes
# one catalogue across every node
curl http://localhost:31313/v1/models
```
## Build from source
```sh
cargo build --release
```
CI runs on every push; keep it green locally:
```sh
cargo fmt --check --all # must be clean
cargo clippy --workspace -- -D warnings # warnings are errors
cargo test --workspace # all tests must pass
```
Tagged releases (`v*`) build SRPMs for `cortex` and `helexa-neuron`
and publish to COPR.
## Status
Pre-1.0 and moving fast. The gateway path (routing, eviction,
translation, metrics) is stable and tested; the candle-native engine
is under active development — expect the supported-model list to track
the open-weight frontier, deliberately narrowly.
Development happens at <https://git.lair.cafe/helexa/helexa>;
<https://github.com/helexa-ai/helexa> is a read-only mirror.
## License
GPL-3.0

View File

@@ -31,3 +31,8 @@ gitea_ci ALL=(root) NOPASSWD: /usr/bin/dnf config-manager addrepo --from-repofil
gitea_ci ALL=(root) NOPASSWD: /usr/bin/dnf install -y libcudnn9-cuda-13
gitea_ci ALL=(root) NOPASSWD: /usr/bin/firewall-cmd --add-service=helexa-neuron --permanent
gitea_ci ALL=(root) NOPASSWD: /usr/bin/firewall-cmd --reload
# deploy-dev.yml fast path: install a freshly-built dev binary over the
# packaged one. Exact source path + args; the workflow must use this
# command form verbatim. The next deploy.yml run reconciles the host
# back to the RPM-owned binary.
gitea_ci ALL=(root) NOPASSWD: /usr/bin/install -o root -g root -m 0755 /var/lib/gitea_ci/neuron-dev /usr/bin/neuron

View File

@@ -4,7 +4,7 @@ Release: 1%{?dist}
Summary: Inference gateway for multi-node GPU clusters
License: GPL-3.0-or-later
URL: https://git.lair.cafe/helexa/cortex
URL: https://git.lair.cafe/helexa/helexa
Source0: %{name}-%{version}.tar.gz
Source1: %{name}-%{version}-vendor.tar.gz

View File

@@ -0,0 +1,119 @@
//! Build/version metadata shared between cortex and neuron.
//!
//! neuron captures these facts at compile time in its `build.rs`
//! (git SHA, enabled cargo features, rustc/candle versions, …) and
//! serves them from `GET /version`. cortex and `helexa-bench`
//! deserialize the same struct so a benchmark run can be attributed to
//! the exact daemon build that produced it — not just the host's CUDA
//! and driver versions that `/discovery` already reports.
//!
//! Every field beyond the always-present package version is
//! `#[serde(default)]` so a newer reader stays compatible with an
//! older neuron that omits a field (and vice versa) — the same
//! forward/backward-compat discipline as
//! [`crate::discovery::ActivationStatus`].
use serde::{Deserialize, Serialize};
/// Build-time identity of a neuron daemon.
///
/// Returned by `GET /version`. The `git_sha` is the canonical "which
/// build is live" key — benchmark records are bucketed by it, so a
/// regression can be pinned to a daemon change rather than a host
/// change. When neuron is built from a source tarball with no git
/// metadata available (and no `HELEXA_BUILD_SHA` injected by CI/RPM),
/// `git_sha` is the string `"unknown"`.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct BuildInfo {
/// Crate version from `CARGO_PKG_VERSION` (e.g. `"0.1.16"`).
pub package_version: String,
/// Short git SHA, or `"unknown"` when unavailable at build time.
#[serde(default = "unknown")]
pub git_sha: String,
/// Full 40-char git SHA when available.
#[serde(default)]
pub git_sha_long: Option<String>,
/// Whether the working tree had uncommitted changes at build time.
/// `false` when the SHA is unknown (tarball build).
#[serde(default)]
pub git_dirty: bool,
/// RFC3339 build timestamp.
#[serde(default)]
pub build_timestamp: Option<String>,
/// `rustc --version` output of the compiler used.
#[serde(default)]
pub rustc_version: Option<String>,
/// Cargo build profile: `"release"` or `"debug"`.
#[serde(default)]
pub profile: Option<String>,
/// Target triple the binary was compiled for.
#[serde(default)]
pub target: Option<String>,
/// Enabled cargo features (e.g. `["cuda", "cudnn"]`). These define
/// the performance envelope, so they are recorded against every
/// benchmark run.
#[serde(default)]
pub features: Vec<String>,
/// Locked `candle-core` version, best-effort from `Cargo.lock`.
#[serde(default)]
pub candle_version: Option<String>,
}
fn unknown() -> String {
"unknown".to_string()
}
impl BuildInfo {
/// A placeholder used by non-neuron benchmark targets (and tests)
/// that have no build metadata to report.
pub fn unknown() -> Self {
BuildInfo {
package_version: env!("CARGO_PKG_VERSION").to_string(),
git_sha: unknown(),
git_sha_long: None,
git_dirty: false,
build_timestamp: None,
rustc_version: None,
profile: None,
target: None,
features: Vec::new(),
candle_version: None,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn round_trips_full() {
let info = BuildInfo {
package_version: "0.1.16".into(),
git_sha: "30d50d6".into(),
git_sha_long: Some("30d50d6abc123".into()),
git_dirty: true,
build_timestamp: Some("2026-06-13T10:00:00+00:00".into()),
rustc_version: Some("rustc 1.85.0".into()),
profile: Some("release".into()),
target: Some("x86_64-unknown-linux-gnu".into()),
features: vec!["cuda".into(), "cudnn".into()],
candle_version: Some("0.10.2".into()),
};
let json = serde_json::to_string(&info).unwrap();
let back: BuildInfo = serde_json::from_str(&json).unwrap();
assert_eq!(info, back);
}
#[test]
fn deserializes_minimal_payload() {
// An older neuron might send only the package version; every
// other field must default rather than fail.
let back: BuildInfo = serde_json::from_str(r#"{"package_version":"0.1.0"}"#).unwrap();
assert_eq!(back.package_version, "0.1.0");
assert_eq!(back.git_sha, "unknown");
assert!(!back.git_dirty);
assert!(back.features.is_empty());
assert!(back.candle_version.is_none());
}
}

View File

@@ -22,6 +22,17 @@ pub struct DiscoveryResponse {
pub driver_version: Option<String>,
pub devices: Vec<DeviceInfo>,
pub harnesses: Vec<String>,
/// Set when the host has an NVIDIA stack that is currently
/// unusable — specifically the userspace↔kernel-module version
/// skew after an un-rebooted driver update ("Driver/library
/// version mismatch"), where every CUDA call including nvidia-smi
/// fails (#19). `None` on healthy hosts AND on hosts with no
/// NVIDIA stack at all (CPU-only is not an error). Carries an
/// operator-actionable description; cortex can read it to route
/// around the node instead of cold-loading into a guaranteed
/// failure.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cuda_unavailable_reason: Option<String>,
}
/// Runtime health metrics for a single GPU device.

View File

@@ -1,4 +1,5 @@
pub mod anthropic;
pub mod build_info;
pub mod catalogue;
pub mod config;
pub mod discovery;

View File

@@ -61,6 +61,12 @@ pub enum ModelStatus {
Unloaded,
Reloading,
Loading,
/// Reported by neuron while a poisoned model auto-recovers via
/// unload→reload (#17/#20). Temporarily unservable but NOT
/// evicted: the gateway holds the route, answers with a transient
/// retry error instead of 404, and must not race a second
/// placement elsewhere.
Recovering,
}
/// Unified model entry as exposed by the gateway's `/v1/models` endpoint.

View File

@@ -71,10 +71,18 @@ pub struct ChatCompletionChoice {
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletionChunk {
#[serde(default)]
pub id: String,
#[serde(default)]
pub object: String,
#[serde(default)]
pub created: u64,
// Lenient deserialization throughout: the gateway parses chunks
// from arbitrary OpenAI-compatible upstreams, and some engines
// omit fields on special frames (e.g. usage-only final chunks).
#[serde(default)]
pub model: String,
#[serde(default)]
pub choices: Vec<ChunkChoice>,
#[serde(skip_serializing_if = "Option::is_none")]
pub usage: Option<Usage>,

View File

@@ -75,11 +75,7 @@ pub fn openai_to_anthropic(resp: ChatCompletionResponse) -> MessagesResponse {
MessageContent::Text(t) => t,
MessageContent::Parts(parts) => serde_json::to_string(&parts).unwrap_or_default(),
};
let stop = c.finish_reason.map(|r| match r.as_str() {
"stop" => "end_turn".to_string(),
"length" => "max_tokens".to_string(),
other => other.to_string(),
});
let stop = c.finish_reason.map(|r| map_stop_reason(&r));
(text, stop)
}
None => (String::new(), None),
@@ -108,3 +104,374 @@ pub fn openai_to_anthropic(resp: ChatCompletionResponse) -> MessagesResponse {
extra: Value::Null,
}
}
// ── Streaming SSE translation (#24) ──────────────────────────────────
/// Map an OpenAI `finish_reason` to an Anthropic `stop_reason`.
pub fn map_stop_reason(openai: &str) -> String {
match openai {
"stop" => "end_turn".to_string(),
"length" => "max_tokens".to_string(),
"tool_calls" => "tool_use".to_string(),
other => other.to_string(),
}
}
/// Stateful OpenAI-SSE → Anthropic-SSE event translator.
///
/// Feed each parsed OpenAI [`crate::openai::ChatCompletionChunk`] to
/// [`on_chunk`](Self::on_chunk) and call [`finish`](Self::finish) on
/// `[DONE]` (or upstream EOF); both return ordered
/// `(event_name, payload)` pairs ready to be framed as
/// `event: <name>\ndata: <payload>\n\n`. The translation is stateless
/// across requests — one instance per stream — and never buffers
/// content: every text delta maps to a `content_block_delta`
/// immediately.
///
/// Event sequence produced (per Anthropic's streaming spec):
/// `message_start` → `content_block_start` / `content_block_delta`* /
/// `content_block_stop` (text and `tool_use` blocks, indexed) →
/// `message_delta` (stop_reason + output usage) → `message_stop`.
#[derive(Debug, Default)]
pub struct AnthropicStreamTranslator {
started: bool,
finished: bool,
/// Index of the currently-open content block, with its kind.
open_block: Option<(u32, OpenBlock)>,
next_index: u32,
stop_reason: Option<String>,
usage: Option<Usage>,
/// Visible text deltas counted as an output-token estimate for
/// streams whose upstream never sends a usage frame (neuron emits
/// one chunk per token, so this is exact there).
text_deltas: u64,
}
#[derive(Debug, PartialEq, Eq)]
enum OpenBlock {
Text,
ToolUse,
}
impl AnthropicStreamTranslator {
pub fn new() -> Self {
Self::default()
}
pub fn on_chunk(&mut self, chunk: &crate::openai::ChatCompletionChunk) -> Vec<(String, Value)> {
let mut out = Vec::new();
if !self.started {
self.started = true;
out.push((
"message_start".to_string(),
json!({
"type": "message_start",
"message": {
// Upstream ids are opaque to Anthropic clients;
// prefix for shape-compatibility with msg_* ids.
"id": format!("msg_{}", chunk.id),
"type": "message",
"role": "assistant",
"content": [],
"model": chunk.model,
"stop_reason": null,
"stop_sequence": null,
// Input tokens are unknown until (if ever) a
// usage frame arrives; corrected in
// message_delta. Anthropic clients sum deltas.
"usage": { "input_tokens": 0, "output_tokens": 0 }
}
}),
));
}
if let Some(usage) = &chunk.usage {
self.usage = Some(usage.clone());
}
for choice in &chunk.choices {
if let Some(text) = choice.delta.get("content").and_then(Value::as_str)
&& !text.is_empty()
{
self.ensure_text_block(&mut out);
self.text_deltas += 1;
let index = self.open_block.as_ref().map(|(i, _)| *i).unwrap_or(0);
out.push((
"content_block_delta".to_string(),
json!({
"type": "content_block_delta",
"index": index,
"delta": { "type": "text_delta", "text": text }
}),
));
}
if let Some(calls) = choice.delta.get("tool_calls").and_then(Value::as_array) {
for call in calls {
let name = call
.get("function")
.and_then(|f| f.get("name"))
.and_then(Value::as_str);
let arguments = call
.get("function")
.and_then(|f| f.get("arguments"))
.and_then(Value::as_str)
.unwrap_or_default();
if let Some(name) = name {
// A named entry begins a new tool_use block.
self.close_open_block(&mut out);
let id = call
.get("id")
.and_then(Value::as_str)
.unwrap_or("toolu_unknown");
let index = self.next_index;
self.next_index += 1;
self.open_block = Some((index, OpenBlock::ToolUse));
out.push((
"content_block_start".to_string(),
json!({
"type": "content_block_start",
"index": index,
"content_block": {
"type": "tool_use",
"id": id,
"name": name,
"input": {}
}
}),
));
}
if !arguments.is_empty()
&& let Some((index, OpenBlock::ToolUse)) = &self.open_block
{
out.push((
"content_block_delta".to_string(),
json!({
"type": "content_block_delta",
"index": index,
"delta": {
"type": "input_json_delta",
"partial_json": arguments
}
}),
));
}
}
}
if let Some(reason) = &choice.finish_reason {
self.stop_reason = Some(map_stop_reason(reason));
}
}
out
}
/// Close the stream: emits the trailing block-stop, message_delta
/// (stop_reason + output usage) and message_stop. Idempotent.
pub fn finish(&mut self) -> Vec<(String, Value)> {
let mut out = Vec::new();
if self.finished || !self.started {
self.finished = true;
return out;
}
self.finished = true;
self.close_open_block(&mut out);
let output_tokens = self
.usage
.as_ref()
.map(|u| u.completion_tokens)
.unwrap_or(self.text_deltas);
let mut usage = json!({ "output_tokens": output_tokens });
if let Some(u) = &self.usage {
usage["input_tokens"] = json!(u.prompt_tokens);
}
out.push((
"message_delta".to_string(),
json!({
"type": "message_delta",
"delta": {
"stop_reason": self.stop_reason.as_deref().unwrap_or("end_turn"),
"stop_sequence": null
},
"usage": usage
}),
));
out.push((
"message_stop".to_string(),
json!({ "type": "message_stop" }),
));
out
}
fn ensure_text_block(&mut self, out: &mut Vec<(String, Value)>) {
match &self.open_block {
Some((_, OpenBlock::Text)) => {}
_ => {
self.close_open_block(out);
let index = self.next_index;
self.next_index += 1;
self.open_block = Some((index, OpenBlock::Text));
out.push((
"content_block_start".to_string(),
json!({
"type": "content_block_start",
"index": index,
"content_block": { "type": "text", "text": "" }
}),
));
}
}
}
fn close_open_block(&mut self, out: &mut Vec<(String, Value)>) {
if let Some((index, _)) = self.open_block.take() {
out.push((
"content_block_stop".to_string(),
json!({ "type": "content_block_stop", "index": index }),
));
}
}
}
#[cfg(test)]
mod stream_tests {
use super::*;
use crate::openai::{ChatCompletionChunk, ChunkChoice};
fn chunk(delta: Value, finish: Option<&str>) -> ChatCompletionChunk {
ChatCompletionChunk {
id: "abc123".into(),
object: "chat.completion.chunk".into(),
created: 1,
model: "Qwen/Qwen3-8B".into(),
choices: vec![ChunkChoice {
index: 0,
delta,
finish_reason: finish.map(String::from),
extra: Value::Null,
}],
usage: None,
extra: Value::Null,
}
}
fn names(events: &[(String, Value)]) -> Vec<&str> {
events.iter().map(|(n, _)| n.as_str()).collect()
}
#[test]
fn text_stream_produces_full_anthropic_sequence() {
let mut t = AnthropicStreamTranslator::new();
let mut all = Vec::new();
all.extend(t.on_chunk(&chunk(json!({"role": "assistant"}), None)));
all.extend(t.on_chunk(&chunk(json!({"content": "Hel"}), None)));
all.extend(t.on_chunk(&chunk(json!({"content": "lo"}), None)));
all.extend(t.on_chunk(&chunk(json!({}), Some("stop"))));
all.extend(t.finish());
assert_eq!(
names(&all),
vec![
"message_start",
"content_block_start",
"content_block_delta",
"content_block_delta",
"content_block_stop",
"message_delta",
"message_stop",
]
);
// message_start carries role/model; deltas carry the text.
assert_eq!(all[0].1["message"]["model"], "Qwen/Qwen3-8B");
assert_eq!(all[2].1["delta"]["text"], "Hel");
assert_eq!(all[3].1["delta"]["text"], "lo");
// stop → end_turn; without a usage frame the output count
// falls back to the delta count (engine-exact for neuron's
// one-chunk-per-token streams).
let md = &all[5].1;
assert_eq!(md["delta"]["stop_reason"], "end_turn");
assert_eq!(md["usage"]["output_tokens"], 2);
}
#[test]
fn length_maps_to_max_tokens_and_missing_finish_defaults_to_end_turn() {
let mut t = AnthropicStreamTranslator::new();
t.on_chunk(&chunk(json!({"content": "x"}), Some("length")));
let fin = t.finish();
assert_eq!(fin[1].1["delta"]["stop_reason"], "max_tokens");
let mut t2 = AnthropicStreamTranslator::new();
t2.on_chunk(&chunk(json!({"content": "x"}), None));
let fin2 = t2.finish();
assert_eq!(fin2[1].1["delta"]["stop_reason"], "end_turn");
}
#[test]
fn tool_call_becomes_tool_use_block() {
let mut t = AnthropicStreamTranslator::new();
let mut all = Vec::new();
all.extend(t.on_chunk(&chunk(json!({"content": "Let me check."}), None)));
all.extend(t.on_chunk(&chunk(
json!({"tool_calls": [{
"index": 0,
"id": "call_7",
"function": {"name": "get_weather", "arguments": "{\"city\":\"Brno\"}"}
}]}),
None,
)));
all.extend(t.on_chunk(&chunk(json!({}), Some("tool_calls"))));
all.extend(t.finish());
assert_eq!(
names(&all),
vec![
"message_start",
"content_block_start", // text
"content_block_delta", // text delta
"content_block_stop", // text closed by tool block
"content_block_start", // tool_use
"content_block_delta", // input_json_delta
"content_block_stop",
"message_delta",
"message_stop",
]
);
let tool_start = &all[4].1;
assert_eq!(tool_start["content_block"]["type"], "tool_use");
assert_eq!(tool_start["content_block"]["id"], "call_7");
assert_eq!(tool_start["content_block"]["name"], "get_weather");
assert_eq!(tool_start["index"], 1);
assert_eq!(all[5].1["delta"]["partial_json"], "{\"city\":\"Brno\"}");
assert_eq!(all[7].1["delta"]["stop_reason"], "tool_use");
}
#[test]
fn usage_frame_feeds_message_delta() {
let mut t = AnthropicStreamTranslator::new();
t.on_chunk(&chunk(json!({"content": "hi"}), Some("stop")));
let mut usage_chunk = chunk(json!({}), None);
usage_chunk.choices.clear();
usage_chunk.usage = Some(crate::openai::Usage {
prompt_tokens: 225,
completion_tokens: 42,
total_tokens: 267,
});
t.on_chunk(&usage_chunk);
let fin = t.finish();
let md = &fin[1].1;
assert_eq!(md["usage"]["output_tokens"], 42);
assert_eq!(md["usage"]["input_tokens"], 225);
}
#[test]
fn finish_is_idempotent_and_silent_without_start() {
let mut t = AnthropicStreamTranslator::new();
assert!(t.finish().is_empty(), "no events for an empty stream");
assert!(t.finish().is_empty());
let mut t2 = AnthropicStreamTranslator::new();
t2.on_chunk(&chunk(json!({"content": "x"}), None));
assert!(!t2.finish().is_empty());
assert!(t2.finish().is_empty(), "second finish must emit nothing");
}
}

View File

@@ -0,0 +1,178 @@
//! Streaming Anthropic SSE translation (#24).
//!
//! The `/v1/messages` handler translates the request envelope to
//! OpenAI before proxying (see `cortex_core::translate`); this module
//! completes the round trip for `stream: true` — the upstream OpenAI
//! SSE stream is re-framed, event by event, into Anthropic's
//! `message_start` / `content_block_*` / `message_delta` /
//! `message_stop` sequence as it arrives. True streaming: each
//! upstream chunk is translated and forwarded immediately; nothing is
//! buffered beyond the current SSE event's bytes.
//!
//! The translation state machine itself is pure and lives in
//! [`cortex_core::translate::AnthropicStreamTranslator`]; this module
//! owns the wire concerns — splitting the upstream byte stream into
//! SSE events, parsing `data:` payloads, and framing the translated
//! events as `event: <name>\ndata: <json>\n\n`.
use axum::body::Body;
use axum::http::StatusCode;
use axum::response::Response;
use bytes::Bytes;
use cortex_core::openai::ChatCompletionChunk;
use cortex_core::translate::AnthropicStreamTranslator;
use futures::StreamExt;
use tokio_stream::wrappers::ReceiverStream;
/// Forward the translated OpenAI request to the upstream node and
/// return the response translated to Anthropic SSE framing.
pub async fn stream_translated(
client: &reqwest::Client,
endpoint: &str,
openai_body: axum::body::Bytes,
model_id: &str,
node_name: &str,
) -> Response {
let url = format!("{endpoint}/v1/chat/completions");
tracing::info!(
handler = "anthropic_messages",
model = %model_id,
node = %node_name,
url = %url,
"proxying streaming request (anthropic SSE translation)"
);
let upstream = match client
.post(&url)
.header("content-type", "application/json")
.body(openai_body)
.send()
.await
{
Ok(r) => r,
Err(e) => {
tracing::warn!(
handler = "anthropic_messages",
node = %node_name,
url = %url,
error = %e,
"anthropic stream: upstream request failed"
);
return anthropic_error(StatusCode::BAD_GATEWAY, "upstream request failed");
}
};
let status = upstream.status();
if !status.is_success() {
tracing::warn!(
handler = "anthropic_messages",
node = %node_name,
url = %url,
status = status.as_u16(),
"anthropic stream: upstream returned non-2xx"
);
return anthropic_error(
StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY),
"upstream returned an error",
);
}
// Bounded channel: a slow client back-pressures the pump task,
// which back-pressures the upstream read — same propagation
// discipline as neuron's own projectors.
let (tx, rx) = tokio::sync::mpsc::channel::<Result<Bytes, std::convert::Infallible>>(32);
let node = node_name.to_string();
tokio::spawn(async move {
let mut upstream = upstream.bytes_stream();
let mut translator = AnthropicStreamTranslator::new();
let mut buf: Vec<u8> = Vec::new();
let mut done = false;
'outer: while let Some(block) = upstream.next().await {
let block = match block {
Ok(b) => b,
Err(e) => {
tracing::warn!(node = %node, error = %e, "anthropic stream: upstream read failed mid-stream");
break;
}
};
buf.extend_from_slice(&block);
// SSE events are separated by a blank line.
while let Some(pos) = find_event_boundary(&buf) {
let event: Vec<u8> = buf.drain(..pos + 2).collect();
let text = String::from_utf8_lossy(&event);
for line in text.lines() {
let Some(data) = line.strip_prefix("data:") else {
continue;
};
let data = data.trim();
if data == "[DONE]" {
done = true;
if !send_frames(&tx, translator.finish()).await {
break 'outer;
}
continue;
}
let Ok(chunk) = serde_json::from_str::<ChatCompletionChunk>(data) else {
tracing::debug!(node = %node, "anthropic stream: unparsable upstream frame skipped");
continue;
};
if !send_frames(&tx, translator.on_chunk(&chunk)).await {
break 'outer;
}
}
}
}
// Upstream ended without [DONE] (error or truncation): still
// close the Anthropic event sequence so clients aren't left
// with an unterminated message.
if !done {
let _ = send_frames(&tx, translator.finish()).await;
}
});
Response::builder()
.status(StatusCode::OK)
.header("content-type", "text/event-stream")
.header("cache-control", "no-cache")
.body(Body::from_stream(ReceiverStream::new(rx)))
.unwrap_or_else(|_| {
anthropic_error(
StatusCode::INTERNAL_SERVER_ERROR,
"failed to build response",
)
})
}
/// `\n\n` boundary of the first complete SSE event in `buf`, if any.
fn find_event_boundary(buf: &[u8]) -> Option<usize> {
buf.windows(2).position(|w| w == b"\n\n")
}
/// Render translated events as SSE frames and send them. Returns
/// `false` when the client has gone away (receiver dropped).
async fn send_frames(
tx: &tokio::sync::mpsc::Sender<Result<Bytes, std::convert::Infallible>>,
events: Vec<(String, serde_json::Value)>,
) -> bool {
for (name, payload) in events {
let frame = format!("event: {name}\ndata: {payload}\n\n");
if tx.send(Ok(Bytes::from(frame))).await.is_err() {
return false;
}
}
true
}
/// Anthropic-shaped error body (`{"type":"error","error":{...}}`).
fn anthropic_error(status: StatusCode, message: &str) -> Response {
let body = serde_json::json!({
"type": "error",
"error": { "type": "api_error", "message": message }
});
Response::builder()
.status(status)
.header("content-type", "application/json")
.body(Body::from(body.to_string()))
.expect("static error response must build")
}

View File

@@ -57,7 +57,7 @@ async fn chat_completions(
// ("model 'X' not found...", "no healthy nodes available")
// — fine to surface to the caller. The warn above carries
// any extra context for operators.
return error_response(404, &e.to_string());
return error_response(e.http_status(), &e.to_string());
}
};
@@ -109,7 +109,7 @@ async fn responses(
error = %e,
"route resolve failed"
);
return error_response(404, &e.to_string());
return error_response(e.http_status(), &e.to_string());
}
};
@@ -157,7 +157,7 @@ async fn completions(
// ("model 'X' not found...", "no healthy nodes available")
// — fine to surface to the caller. The warn above carries
// any extra context for operators.
return error_response(404, &e.to_string());
return error_response(e.http_status(), &e.to_string());
}
};
@@ -178,7 +178,7 @@ async fn completions(
/// `POST /v1/messages` — accept Anthropic format, translate, proxy, translate back.
async fn anthropic_messages(
State(fleet): State<Arc<CortexState>>,
headers: HeaderMap,
_headers: HeaderMap,
body: Bytes,
) -> Response {
// Parse as Anthropic request.
@@ -225,7 +225,7 @@ async fn anthropic_messages(
// ("model 'X' not found...", "no healthy nodes available")
// — fine to surface to the caller. The warn above carries
// any extra context for operators.
return error_response(404, &e.to_string());
return error_response(e.http_status(), &e.to_string());
}
};
@@ -247,28 +247,23 @@ async fn anthropic_messages(
let start = Instant::now();
if is_streaming {
// TODO: streaming Anthropic translation requires converting SSE format.
// For now, proxy the OpenAI SSE stream directly (clients that can handle
// OpenAI SSE will work; full Anthropic SSE translation is a follow-up).
let result = proxy::forward_request(
// Anthropic SSE translation (#24): upstream speaks OpenAI SSE;
// re-frame it event-by-event into Anthropic's message_start /
// content_block_* / message_delta / message_stop sequence.
let resp = crate::anthropic_sse::stream_translated(
&fleet.http_client,
&route,
"/v1/chat/completions",
headers,
&route.endpoint,
openai_body,
&model_id,
&route.node_name,
)
.await;
metrics::histogram!("cortex_request_duration_seconds", &labels)
.record(start.elapsed().as_secs_f64());
match result {
Ok(resp) => resp,
Err(e) => {
if !resp.status().is_success() {
metrics::counter!("cortex_request_errors_total", &labels).increment(1);
// forward_request already warn'd with the wire-level
// detail; no need to log again here.
e.into_response()
}
}
resp
} else {
// Non-streaming: proxy, buffer full response, translate back to Anthropic.
let target_url = format!("{}/v1/chat/completions", route.endpoint);
@@ -591,7 +586,8 @@ async fn proxy_with_metrics(
}
let start = Instant::now();
let result = proxy::forward_request(&fleet.http_client, route, path, headers, body).await;
let result =
proxy::forward_request(&fleet.http_client, route, path, headers, body, model_id).await;
let duration = start.elapsed();
match result {

View File

@@ -1,3 +1,4 @@
pub mod anthropic_sse;
pub mod evictor;
pub mod handlers;
pub mod metrics;

View File

@@ -46,6 +46,14 @@ fn describe_metrics() {
"Generation throughput in tokens per second"
);
metrics::describe_counter!("cortex_requests_total", "Total number of proxied requests");
metrics::describe_counter!(
"cortex_prompt_tokens_total",
"Total prompt tokens reported by upstream usage objects"
);
metrics::describe_counter!(
"cortex_completion_tokens_total",
"Total completion tokens reported by upstream usage objects"
);
metrics::describe_counter!(
"cortex_request_errors_total",
"Total number of failed proxy requests"

View File

@@ -197,6 +197,7 @@ fn parse_status(s: &str) -> ModelStatus {
"unloaded" => ModelStatus::Unloaded,
"reloading" => ModelStatus::Reloading,
"loading" => ModelStatus::Loading,
"recovering" => ModelStatus::Recovering,
_ => ModelStatus::Loaded,
}
}

View File

@@ -9,7 +9,12 @@ use anyhow::Result;
use axum::body::Body;
use axum::http::{HeaderMap, StatusCode};
use axum::response::{IntoResponse, Response};
use futures::Stream;
use futures::stream::BoxStream;
use reqwest::Client;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Instant;
/// Proxy a request body to the resolved backend node and stream the response.
///
@@ -25,7 +30,9 @@ pub async fn forward_request(
path: &str,
headers: HeaderMap,
body: bytes::Bytes,
model_id: &str,
) -> Result<Response, ProxyError> {
let request_start = Instant::now();
let url = format!("{}{}", route.endpoint, path);
tracing::info!(
node = %route.node_name,
@@ -73,7 +80,10 @@ pub async fn forward_request(
let status = StatusCode::from_u16(upstream_status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
let resp_headers = upstream_resp.headers().clone();
let stream = upstream_resp.bytes_stream();
let stream = TokenMetricsStream::new(
Box::pin(upstream_resp.bytes_stream()),
TokenMetrics::new(model_id, &route.node_name, request_start),
);
let body = Body::from_stream(stream);
@@ -119,3 +129,224 @@ impl IntoResponse for ProxyError {
(status, axum::Json(body)).into_response()
}
}
// ── Per-request token metrics (#21) ─────────────────────────────────
//
// The proxy never buffers or re-serialises the upstream body — chunks
// are forwarded verbatim. For metrics it observes each chunk's arrival
// time and keeps a bounded tail of the body text, from which the final
// OpenAI `usage` object (present on the last SSE chunk and on
// non-streaming JSON bodies alike) yields engine-truth token counts.
//
// Emitted per request, labelled {model, node}:
// cortex_time_to_first_token_seconds (histogram) — first body chunk
// cortex_tokens_per_second (histogram) — completion tokens
// over the decode window (first→last chunk); falls back to the
// full request duration for single-chunk (non-streaming) bodies
// cortex_prompt_tokens_total / cortex_completion_tokens_total (counters)
/// Cap on the retained body tail. The usage object rides on the final
/// chunk, so a generous tail is plenty; the cap bounds memory on huge
/// non-streaming bodies.
const TAIL_CAP_BYTES: usize = 64 * 1024;
/// Find the value of the LAST `"key": <integer>` occurrence in `tail`.
/// Pure and chunk-boundary-safe (the tail is contiguous appended text).
/// The quoted-needle form means `completion_tokens` never matches
/// `completion_tokens_details`.
pub(crate) fn last_count_for(tail: &str, key: &str) -> Option<u64> {
let needle = format!("\"{key}\"");
let mut result = None;
for (idx, _) in tail.match_indices(&needle) {
let rest = tail[idx + needle.len()..].trim_start();
let Some(rest) = rest.strip_prefix(':') else {
continue;
};
let rest = rest.trim_start();
let digits: &str = &rest[..rest
.char_indices()
.find(|(_, c)| !c.is_ascii_digit())
.map(|(i, _)| i)
.unwrap_or(rest.len())];
if let Ok(v) = digits.parse::<u64>() {
result = Some(v);
}
}
result
}
struct TokenMetrics {
labels: [(&'static str, String); 2],
request_start: Instant,
first_chunk: Option<Instant>,
last_chunk: Option<Instant>,
tail: String,
finished: bool,
}
impl TokenMetrics {
fn new(model_id: &str, node_name: &str, request_start: Instant) -> Self {
Self {
labels: [
("model", model_id.to_string()),
("node", node_name.to_string()),
],
request_start,
first_chunk: None,
last_chunk: None,
tail: String::new(),
finished: false,
}
}
fn observe(&mut self, chunk: &[u8]) {
let now = Instant::now();
self.first_chunk.get_or_insert(now);
self.last_chunk = Some(now);
self.tail.push_str(&String::from_utf8_lossy(chunk));
if self.tail.len() > TAIL_CAP_BYTES {
// Keep the newest half; the usage object is always at the
// very end of the body. Split at a char boundary.
let mut cut = self.tail.len() - TAIL_CAP_BYTES / 2;
while !self.tail.is_char_boundary(cut) {
cut += 1;
}
self.tail.drain(..cut);
}
}
/// Emit the metrics exactly once — called on clean stream end and
/// from Drop (client disconnect mid-stream still records what we
/// saw).
fn finish(&mut self) {
if self.finished {
return;
}
self.finished = true;
let Some(first) = self.first_chunk else {
return; // no body ever arrived — nothing to record
};
let ttft = first.duration_since(self.request_start).as_secs_f64();
metrics::histogram!("cortex_time_to_first_token_seconds", &self.labels).record(ttft);
if let Some(prompt) = last_count_for(&self.tail, "prompt_tokens") {
metrics::counter!("cortex_prompt_tokens_total", &self.labels).increment(prompt);
}
let Some(completion) = last_count_for(&self.tail, "completion_tokens") else {
return;
};
if completion == 0 {
return;
}
metrics::counter!("cortex_completion_tokens_total", &self.labels).increment(completion);
let last = self.last_chunk.unwrap_or(first);
let decode_window = last.duration_since(first).as_secs_f64();
// Streaming: rate over the decode window (first→last chunk).
// Non-streaming bodies arrive as ~one chunk (window ≈ 0), where
// the only honest denominator is the full request duration.
let secs = if decode_window >= 0.1 {
decode_window
} else {
last.duration_since(self.request_start).as_secs_f64()
};
if secs > 0.0 {
metrics::histogram!("cortex_tokens_per_second", &self.labels)
.record(completion as f64 / secs);
}
}
}
/// Pass-through stream wrapper that feeds [`TokenMetrics`]. Emits on
/// clean end-of-stream; the Drop impl covers client disconnects.
struct TokenMetricsStream {
inner: BoxStream<'static, Result<bytes::Bytes, reqwest::Error>>,
metrics: TokenMetrics,
}
impl TokenMetricsStream {
fn new(
inner: BoxStream<'static, Result<bytes::Bytes, reqwest::Error>>,
metrics: TokenMetrics,
) -> Self {
Self { inner, metrics }
}
}
impl Stream for TokenMetricsStream {
type Item = Result<bytes::Bytes, reqwest::Error>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
match this.inner.as_mut().poll_next(cx) {
Poll::Ready(Some(Ok(chunk))) => {
this.metrics.observe(&chunk);
Poll::Ready(Some(Ok(chunk)))
}
Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
Poll::Ready(None) => {
this.metrics.finish();
Poll::Ready(None)
}
Poll::Pending => Poll::Pending,
}
}
}
impl Drop for TokenMetricsStream {
fn drop(&mut self) {
self.metrics.finish();
}
}
#[cfg(test)]
mod tests {
use super::last_count_for;
#[test]
fn extracts_counts_from_final_sse_usage_chunk() {
let tail = concat!(
"data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n",
"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":225,",
"\"completion_tokens\":42,\"total_tokens\":267}}\n\n",
"data: [DONE]\n\n"
);
assert_eq!(last_count_for(tail, "prompt_tokens"), Some(225));
assert_eq!(last_count_for(tail, "completion_tokens"), Some(42));
}
#[test]
fn extracts_counts_from_non_streaming_body() {
let tail = "{\"choices\":[{\"message\":{\"content\":\"hi\"}}],\
\"usage\":{\"prompt_tokens\": 12, \"completion_tokens\": 7}}";
assert_eq!(last_count_for(tail, "prompt_tokens"), Some(12));
assert_eq!(last_count_for(tail, "completion_tokens"), Some(7));
}
#[test]
fn ignores_details_variants_and_takes_last_occurrence() {
// completion_tokens_details must not shadow completion_tokens,
// and the LAST usage object wins (matters when content echoes
// a usage-shaped string earlier in the stream).
let tail = concat!(
"data: {\"usage\":{\"completion_tokens\":1}}\n\n",
"data: {\"usage\":{\"completion_tokens\":99,",
"\"completion_tokens_details\":{\"reasoning_tokens\":3}}}\n\n"
);
assert_eq!(last_count_for(tail, "completion_tokens"), Some(99));
}
#[test]
fn absent_keys_yield_none() {
assert_eq!(
last_count_for("data: [DONE]\n\n", "completion_tokens"),
None
);
assert_eq!(last_count_for("", "prompt_tokens"), None);
// key present but non-numeric value
assert_eq!(
last_count_for("\"completion_tokens\": null", "completion_tokens"),
None
);
}
}

View File

@@ -56,6 +56,22 @@ pub enum RouteError {
node: String,
message: String,
},
#[error(
"model '{model_id}' is recovering on node '{node}' (device context rebuild in progress) — retry shortly"
)]
ModelRecovering { model_id: String, node: String },
}
impl RouteError {
/// HTTP status the gateway should answer with. `ModelRecovering`
/// is the one transient case (503, retry the same request);
/// everything else keeps the long-standing 404 behaviour.
pub fn http_status(&self) -> u16 {
match self {
RouteError::ModelRecovering { .. } => 503,
_ => 404,
}
}
}
/// Resolve which node should serve a request for the given model.
@@ -76,11 +92,12 @@ pub async fn resolve(
"alias resolved"
);
}
// Snapshot loaded / unloaded state from the poller cache.
let (loaded_route, unloaded_route, any_healthy) = {
// Snapshot loaded / unloaded / recovering state from the poller cache.
let (loaded_route, unloaded_route, recovering_node, any_healthy) = {
let nodes = fleet.nodes.read().await;
let mut loaded_route = None;
let mut unloaded_route = None;
let mut recovering_node = None;
let mut any_healthy = false;
for node in nodes.values() {
if !node.healthy {
@@ -98,6 +115,17 @@ pub async fn resolve(
unloaded_route = Some((node.name.clone(), node.endpoint.clone(), true));
}
}
// Auto-recovering (#17/#20): the model is rebuilding
// its device context on this node. Hold the route —
// answer "retry shortly" rather than 404, and do NOT
// fall through to the catalogue cold-load, which
// would race a second placement (and a second copy's
// worth of VRAM) against the in-flight recovery.
ModelStatus::Recovering => {
if recovering_node.is_none() {
recovering_node = Some(node.name.clone());
}
}
// Loading is gateway-synthesised from neuron's
// activation snapshot; it never appears on the
// wire from neuron's `/models`. Skip — the model
@@ -110,7 +138,7 @@ pub async fn resolve(
}
}
}
(loaded_route, unloaded_route, any_healthy)
(loaded_route, unloaded_route, recovering_node, any_healthy)
};
if !any_healthy {
@@ -122,12 +150,20 @@ pub async fn resolve(
return finish(fleet, &node_name, &neuron_endpoint, model_id, cold_start).await;
}
// Priority 2: known to neuron but unloaded (neuron's lazy load).
// Priority 2: recovering somewhere — transient hold, not a reroute.
if let Some(node) = recovering_node {
return Err(RouteError::ModelRecovering {
model_id: model_id.to_string(),
node,
});
}
// Priority 3: known to neuron but unloaded (neuron's lazy load).
if let Some((node_name, neuron_endpoint, cold_start)) = unloaded_route {
return finish(fleet, &node_name, &neuron_endpoint, model_id, cold_start).await;
}
// Priority 3: catalogue × topology cold-load.
// Priority 4: catalogue × topology cold-load.
if let Some(profile) = fleet.catalogue.get(model_id) {
let (node_name, neuron_endpoint) = pick_feasible_neuron(fleet, profile).await?;
cold_load(fleet, &node_name, &neuron_endpoint, profile).await?;

View File

@@ -123,3 +123,124 @@ async fn test_anthropic_invalid_request() {
assert_eq!(resp.status(), 400);
}
/// #24: a streaming Anthropic request gets a translated Anthropic SSE
/// stream — not raw OpenAI frames. Verifies the full event sequence,
/// text reassembly, and the content type.
#[tokio::test]
async fn test_anthropic_streaming_sse_translation() {
let mock_url =
common::spawn_streaming_mock_neuron(4, std::time::Duration::from_millis(20)).await;
let gw_url = common::spawn_gateway(&mock_url).await;
let client = reqwest::Client::new();
let resp = client
.post(format!("{gw_url}/v1/messages"))
.header("content-type", "application/json")
.json(&json!({
"model": "test-model",
"max_tokens": 64,
"stream": true,
"messages": [{"role": "user", "content": "Hi"}]
}))
.send()
.await
.expect("request should succeed");
assert_eq!(resp.status(), 200);
assert!(
resp.headers()
.get("content-type")
.and_then(|v| v.to_str().ok())
.unwrap_or("")
.starts_with("text/event-stream"),
"anthropic stream must be SSE"
);
let body = resp.text().await.expect("stream should complete");
assert!(
!body.contains("chat.completion.chunk"),
"raw OpenAI frames must not leak through:\n{body}"
);
let event_names: Vec<&str> = body
.lines()
.filter_map(|l| l.strip_prefix("event: "))
.collect();
assert_eq!(
event_names,
vec![
"message_start",
"content_block_start",
"content_block_delta",
"content_block_delta",
"content_block_delta",
"content_block_delta",
"content_block_stop",
"message_delta",
"message_stop",
],
"unexpected event sequence:\n{body}"
);
// Reassemble the text deltas: the mock emits token0..token3.
let text: String = body
.lines()
.filter_map(|l| l.strip_prefix("data: "))
.filter_map(|d| serde_json::from_str::<serde_json::Value>(d).ok())
.filter(|v| v["type"] == "content_block_delta")
.filter_map(|v| v["delta"]["text"].as_str().map(String::from))
.collect();
assert_eq!(text, "token0token1token2token3");
// The mock sends no finish_reason — stop_reason defaults to
// end_turn, and output_tokens falls back to the delta count.
let message_delta = body
.lines()
.filter_map(|l| l.strip_prefix("data: "))
.filter_map(|d| serde_json::from_str::<serde_json::Value>(d).ok())
.find(|v| v["type"] == "message_delta")
.expect("message_delta event present");
assert_eq!(message_delta["delta"]["stop_reason"], "end_turn");
assert_eq!(message_delta["usage"]["output_tokens"], 4);
}
/// #24: an upstream usage frame (stream_options include_usage shape)
/// rides into message_delta as input/output token counts.
#[tokio::test]
async fn test_anthropic_streaming_usage_propagation() {
let mock_url = common::spawn_streaming_mock_neuron_with_usage(
3,
std::time::Duration::from_millis(10),
225,
42,
)
.await;
let gw_url = common::spawn_gateway(&mock_url).await;
let client = reqwest::Client::new();
let body = client
.post(format!("{gw_url}/v1/messages"))
.header("content-type", "application/json")
.json(&json!({
"model": "test-model",
"max_tokens": 64,
"stream": true,
"messages": [{"role": "user", "content": "Hi"}]
}))
.send()
.await
.expect("request should succeed")
.text()
.await
.expect("stream should complete");
let message_delta = body
.lines()
.filter_map(|l| l.strip_prefix("data: "))
.filter_map(|d| serde_json::from_str::<serde_json::Value>(d).ok())
.find(|v| v["type"] == "message_delta")
.expect("message_delta event present");
assert_eq!(message_delta["usage"]["output_tokens"], 42);
assert_eq!(message_delta["usage"]["input_tokens"], 225);
}

View File

@@ -196,6 +196,91 @@ pub async fn spawn_streaming_mock_neuron(chunk_count: usize, chunk_delay: Durati
base_url
}
/// Like `spawn_streaming_mock_neuron`, but the stream ends with an
/// OpenAI `stream_options.include_usage`-style final chunk (empty
/// choices + usage object) before `[DONE]` — the shape the gateway's
/// token metrics (#21) extract counts from.
pub async fn spawn_streaming_mock_neuron_with_usage(
chunk_count: usize,
chunk_delay: Duration,
prompt_tokens: u64,
completion_tokens: u64,
) -> String {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let base_url = format!("http://{addr}");
let inference_url = base_url.clone();
let app = Router::new()
.route("/models", get(mock_neuron_list_models))
.route(
"/models/{model_id}/endpoint",
get(move |Path(_model_id): Path<String>| {
let url = inference_url.clone();
async move { Json(json!({"url": url})) }
}),
)
.route(
"/v1/chat/completions",
post(move |Json(body): Json<Value>| async move {
let model = body
.get("model")
.and_then(|v| v.as_str())
.unwrap_or("unknown")
.to_string();
let mut chunks: Vec<String> = (0..chunk_count)
.map(|i| {
let chunk = json!({
"id": "chatcmpl-stream-002",
"object": "chat.completion.chunk",
"created": 1700000000_u64,
"model": model,
"choices": [{
"index": 0,
"delta": { "content": format!("token{i}") },
"finish_reason": null
}]
});
format!("data: {chunk}\n\n")
})
.collect();
let usage_chunk = json!({
"id": "chatcmpl-stream-002",
"object": "chat.completion.chunk",
"created": 1700000000_u64,
"model": model,
"choices": [],
"usage": {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens
}
});
chunks.push(format!("data: {usage_chunk}\n\n"));
chunks.push("data: [DONE]\n\n".to_string());
let delay = chunk_delay;
let stream = stream::iter(chunks).then(move |chunk| async move {
tokio::time::sleep(delay).await;
Ok::<_, std::convert::Infallible>(chunk)
});
Response::builder()
.header(header::CONTENT_TYPE, "text/event-stream")
.header(header::CACHE_CONTROL, "no-cache")
.body(Body::from_stream(stream))
.unwrap()
}),
);
tokio::spawn(async move {
axum::serve(listener, app).await.unwrap();
});
base_url
}
/// Spawns a mock neuron with a custom models list.
pub async fn spawn_mock_neuron_with_models(models_response: Value) -> String {
spawn_mock_neuron_with_models_and_health(models_response, default_health_response()).await

View File

@@ -1,20 +1,26 @@
mod common;
use serde_json::json;
use std::sync::OnceLock;
/// The metrics recorder is a process-wide global; both tests in this
/// binary run against one shared install. Assertions must therefore be
/// order-independent (presence of names / monotonic counters, not
/// "empty before").
fn recorder() -> &'static metrics_exporter_prometheus::PrometheusHandle {
static HANDLE: OnceLock<metrics_exporter_prometheus::PrometheusHandle> = OnceLock::new();
HANDLE.get_or_init(|| {
cortex_gateway::metrics::install_test_recorder().expect("recorder should install")
})
}
#[tokio::test]
async fn test_metrics_emitted_after_proxy() {
let handle = cortex_gateway::metrics::install_test_recorder().expect("recorder should install");
let handle = recorder();
let mock_url = common::spawn_mock_neuron().await;
let gw_url = common::spawn_gateway(&mock_url).await;
let before = handle.render();
assert!(
!before.contains("cortex_requests_total"),
"no request metrics before any requests"
);
let client = reqwest::Client::new();
let resp = client
.post(format!("{gw_url}/v1/chat/completions"))
@@ -44,3 +50,72 @@ async fn test_metrics_emitted_after_proxy() {
"no errors expected for a successful request"
);
}
#[tokio::test]
async fn test_token_metrics_emitted_for_streamed_request() {
// #21: a streamed chat completion with a final usage chunk must
// produce TTFT + tok/s histograms and prompt/completion token
// counters, labelled with model and node. The recorder is global
// per-process, so this test runs in its own binary invocation —
// cargo's per-file integration binaries give us that as long as
// only one test in this file installs the recorder... it isn't:
// test_metrics_emitted_after_proxy also installs. Whichever wins
// the race, both render from the same recorder, so assert on
// delta-able names rather than exact totals.
let handle = recorder();
let mock_url = common::spawn_streaming_mock_neuron_with_usage(
5,
std::time::Duration::from_millis(40),
225,
42,
)
.await;
let gw_url = common::spawn_gateway(&mock_url).await;
let client = reqwest::Client::new();
let resp = client
.post(format!("{gw_url}/v1/chat/completions"))
.header("content-type", "application/json")
.json(&json!({
"model": "test-model",
"messages": [{"role": "user", "content": "Hi"}],
"stream": true
}))
.send()
.await
.expect("request should succeed");
assert_eq!(resp.status(), 200);
let body = resp.text().await.expect("stream should complete");
assert!(body.contains("[DONE]"));
let rendered = handle.render();
for needle in [
"cortex_time_to_first_token_seconds",
"cortex_tokens_per_second",
] {
assert!(
rendered.contains(needle),
"{needle} should be present.\nMetrics:\n{rendered}"
);
}
// The recorder is shared with the sibling test (same model/node
// labels), so counters are lower bounds, not exact values: this
// request contributed prompt=225 / completion=42.
let counter_value = |name: &str| -> u64 {
rendered
.lines()
.find(|l| l.starts_with(name) && l.contains(r#"model="test-model""#))
.and_then(|l| l.rsplit(' ').next())
.and_then(|v| v.parse().ok())
.unwrap_or_else(|| panic!("{name} should be present.\nMetrics:\n{rendered}"))
};
assert!(
counter_value("cortex_prompt_tokens_total") >= 225,
"prompt token counter should include this request's 225.\nMetrics:\n{rendered}"
);
assert!(
counter_value("cortex_completion_tokens_total") >= 42,
"completion token counter should include this request's 42.\nMetrics:\n{rendered}"
);
}

View File

@@ -375,3 +375,39 @@ async fn test_poller_captures_activation_from_health() {
assert_eq!(activation.in_progress.as_deref(), Some("Qwen/model-x"));
assert_eq!(activation.pending, vec!["Qwen/model-y".to_string()]);
}
#[tokio::test]
async fn test_poller_parses_recovering_status() {
// #20: a model auto-recovering on a neuron (poisoned → unload →
// reload, #17) is reported with status "recovering" and must land
// in gateway state as the dedicated Recovering status — not fall
// through the parser's catch-all to Loaded.
let mock_url = common::spawn_mock_neuron_with_models(json!([
{"id": "model-r", "harness": "candle", "status": "recovering", "devices": [0, 1], "vram_used_mb": null}
]))
.await;
let config = GatewayConfig {
gateway: GatewaySettings {
listen: "127.0.0.1:0".into(),
metrics_listen: "127.0.0.1:0".into(),
},
eviction: EvictionSettings {
strategy: EvictionStrategy::Lru,
defrag_after_cycles: 0,
},
neurons: vec![NeuronEndpoint {
name: "test-node".into(),
endpoint: mock_url,
}],
models_config: "/dev/null".into(),
};
let fleet = Arc::new(CortexState::from_config(&config));
cortex_gateway::poller::poll_once(&fleet).await;
let nodes = fleet.nodes.read().await;
let node = nodes.get("test-node").unwrap();
let model_r = node.models.get("model-r").expect("model-r should exist");
assert_eq!(model_r.status, ModelStatus::Recovering);
}

View File

@@ -171,3 +171,64 @@ async fn test_missing_model_field() {
let body: serde_json::Value = resp.json().await.unwrap();
assert!(body["error"]["message"].as_str().unwrap().contains("model"));
}
#[tokio::test]
async fn test_recovering_model_returns_503_and_stays_listed() {
// #20: while a model auto-recovers on a neuron, the gateway must
// hold the route — transient 503 ("retry shortly"), not the 404
// "not found on any node" that makes a recovering model look
// evicted — and keep listing it on /v1/models.
let mock_url = common::spawn_mock_neuron().await;
let (fleet, gw_url) = common::spawn_gateway_with_state(&mock_url).await;
{
let mut nodes = fleet.nodes.write().await;
let node = nodes.get_mut("mock-node").expect("node must exist");
node.models.insert(
"recovering-model".into(),
cortex_core::node::ModelEntry {
id: "recovering-model".into(),
status: cortex_core::node::ModelStatus::Recovering,
last_accessed: None,
vram_estimate_mb: Some(8000),
capabilities: Vec::new(),
},
);
}
let client = reqwest::Client::new();
let resp = client
.post(format!("{gw_url}/v1/chat/completions"))
.header("content-type", "application/json")
.json(&json!({
"model": "recovering-model",
"messages": [{"role": "user", "content": "Hi"}]
}))
.send()
.await
.expect("request should succeed");
assert_eq!(resp.status(), 503);
let body: serde_json::Value = resp.json().await.unwrap();
let message = body["error"]["message"].as_str().unwrap();
assert!(
message.contains("recovering") && message.contains("retry"),
"503 body must say recovering/retry, got: {message}"
);
// The model must still be visible on the unified models endpoint.
let models: serde_json::Value = client
.get(format!("{gw_url}/v1/models"))
.send()
.await
.expect("models request should succeed")
.json()
.await
.unwrap();
let listed = models["data"]
.as_array()
.unwrap()
.iter()
.any(|m| m["id"] == "recovering-model");
assert!(listed, "recovering model must stay listed on /v1/models");
}

View File

@@ -3,7 +3,7 @@ name = "helexa-acp"
version = "0.1.16"
edition = "2024"
license = "Apache-2.0"
repository = "https://git.lair.cafe/helexa/cortex"
repository = "https://git.lair.cafe/helexa/helexa"
description = """
Agent Client Protocol bridge for the helexa self-hosted LLM stack.
Speaks ACP to ACP-compatible editor clients (Zed, etc.) and forwards

View File

@@ -58,8 +58,8 @@ one vendor's agent client.
### From source
```sh
git clone https://git.lair.cafe/helexa/cortex.git
cd cortex
git clone https://git.lair.cafe/helexa/helexa.git
cd helexa
cargo install --path crates/helexa-acp
# Binary lands at ~/.cargo/bin/helexa-acp
```
@@ -536,7 +536,7 @@ Cargo.toml-only.
## Contributing
Repository: https://git.lair.cafe/helexa/cortex (`crates/helexa-acp/`).
Repository: https://git.lair.cafe/helexa/helexa (`crates/helexa-acp/`).
Issues / PRs welcome. The canonical staged plan is in
`~/.claude/plans/plan-the-per-device-worker-abstract-micali.md` on
the maintainer's machine; the substages 3a3e and 6a/6b that the

View File

@@ -0,0 +1,38 @@
[package]
name = "helexa-bench"
version.workspace = true
edition.workspace = true
license.workspace = true
repository.workspace = true
[[bin]]
name = "helexa-bench"
path = "src/main.rs"
[dependencies]
cortex-core = { workspace = true }
tokio = { workspace = true }
reqwest = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
figment = { workspace = true }
anyhow = { workspace = true }
async-trait = { workspace = true }
clap = { workspace = true }
tracing = { workspace = true }
tracing-subscriber = { workspace = true }
chrono = { workspace = true }
futures = { workspace = true }
tokio-stream = { workspace = true }
eventsource-stream = { workspace = true }
# SQLite system-of-record. `bundled` compiles SQLite from source so the
# binary has no libsqlite3 runtime dependency — matches the project's
# single-static-binary packaging.
rusqlite = { version = "0.32", features = ["bundled"] }
[dev-dependencies]
axum = { workspace = true }
# Jail (isolated cwd + env) for config tests.
figment = { workspace = true, features = ["test"] }

View File

@@ -0,0 +1,159 @@
//! Outbound calls to a benchmark target: build identity, host discovery,
//! and warm-model enumeration. Neuron targets use the native neuron API;
//! `openai` targets use the OpenAI-compatible surface (preliminary).
use crate::config::{TargetConfig, TargetKind};
use anyhow::{Context, Result};
use cortex_core::build_info::BuildInfo;
use cortex_core::discovery::DiscoveryResponse;
use cortex_core::harness::ModelInfo;
use cortex_core::openai::ModelsResponse;
use std::time::Duration;
/// How long to wait on the cheap metadata polls (version/discovery/models).
const META_TIMEOUT: Duration = Duration::from_secs(10);
pub struct TargetClient {
http: reqwest::Client,
}
impl TargetClient {
pub fn new(request_timeout: Duration) -> Result<Self> {
let http = reqwest::Client::builder()
.timeout(request_timeout)
.build()
.context("building HTTP client")?;
Ok(TargetClient { http })
}
pub fn http(&self) -> &reqwest::Client {
&self.http
}
/// Chat-completions URL for the target.
pub fn chat_url(&self, target: &TargetConfig) -> String {
let base = target.endpoint.trim_end_matches('/');
match target.kind {
// neuron exposes OpenAI routes under /v1.
TargetKind::Neuron => format!("{base}/v1/chat/completions"),
// openai endpoint is the /v1 base already (bench.py convention).
TargetKind::Openai => format!("{base}/chat/completions"),
}
}
/// Build identity. Neuron: `GET /version`. Openai: a synthetic
/// placeholder keyed by `"external"` so the version-aware skip logic
/// treats it as one stable build (comparison runs are manual anyway).
pub async fn fetch_version(&self, target: &TargetConfig) -> Result<BuildInfo> {
match target.kind {
TargetKind::Neuron => {
let base = target.endpoint.trim_end_matches('/');
let info = self
.http
.get(format!("{base}/version"))
.timeout(META_TIMEOUT)
.send()
.await
.context("GET /version")?
.error_for_status()
.context("GET /version status")?
.json::<BuildInfo>()
.await
.context("decoding /version")?;
Ok(info)
}
TargetKind::Openai => {
let mut info = BuildInfo::unknown();
info.git_sha = "external".to_string();
Ok(info)
}
}
}
/// Host discovery (neuron only).
pub async fn fetch_discovery(
&self,
target: &TargetConfig,
) -> Result<Option<DiscoveryResponse>> {
if target.kind != TargetKind::Neuron {
return Ok(None);
}
let base = target.endpoint.trim_end_matches('/');
let disco = self
.http
.get(format!("{base}/discovery"))
.timeout(META_TIMEOUT)
.send()
.await
.context("GET /discovery")?
.error_for_status()
.context("GET /discovery status")?
.json::<DiscoveryResponse>()
.await
.context("decoding /discovery")?;
Ok(Some(disco))
}
/// Warm models — those ready to serve without a cold load.
///
/// Neuron: `GET /models` filtered to `status == "loaded"` (skips
/// `recovering`/`poisoned`). Openai: `GET /models`, honouring the
/// helexa `loaded` extension when present, else treating all listed
/// models as warm.
pub async fn warm_models(&self, target: &TargetConfig) -> Result<Vec<ModelInfo>> {
let base = target.endpoint.trim_end_matches('/');
match target.kind {
TargetKind::Neuron => {
let models = self
.http
.get(format!("{base}/models"))
.timeout(META_TIMEOUT)
.send()
.await
.context("GET /models")?
.error_for_status()
.context("GET /models status")?
.json::<Vec<ModelInfo>>()
.await
.context("decoding /models")?;
Ok(models
.into_iter()
.filter(|m| m.status == "loaded")
.collect())
}
TargetKind::Openai => {
let resp = self
.http
.get(format!("{base}/models"))
.timeout(META_TIMEOUT)
.send()
.await
.context("GET /models")?
.error_for_status()
.context("GET /models status")?
.json::<ModelsResponse>()
.await
.context("decoding /models")?;
Ok(resp
.data
.into_iter()
.filter(|m| {
// honour the helexa `loaded` extension if present
m.extra
.get("loaded")
.and_then(|v| v.as_bool())
.unwrap_or(true)
})
.map(|m| ModelInfo {
id: m.id,
harness: "openai".to_string(),
status: "loaded".to_string(),
devices: Vec::new(),
vram_used_mb: None,
capabilities: Vec::new(),
})
.collect())
}
}
}
}

View File

@@ -0,0 +1,210 @@
//! Bench configuration: loaded from `helexa-bench.toml` with figment,
//! `BENCH_`-prefixed env overrides (mirrors `NeuronConfig::load`).
use figment::{
Figment,
providers::{Env, Format, Toml},
};
use serde::{Deserialize, Serialize};
use std::path::Path;
use std::time::Duration;
/// Top-level bench config.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenchConfig {
#[serde(default)]
pub bench: BenchSettings,
#[serde(default)]
pub scenarios: ScenarioConfig,
/// Endpoints to benchmark. At least one is required for `run`/`once`.
#[serde(default)]
pub targets: Vec<TargetConfig>,
}
/// Loop/timing knobs.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenchSettings {
/// Pause between full sweeps of all targets.
#[serde(default = "default_sweep_interval")]
pub sweep_interval_secs: u64,
/// Target number of measured samples to record for a given
/// (target, build SHA, model, scenario). Once met, later sweeps skip
/// that cell — so a fully-sampled build costs only cheap version
/// polls until a new SHA ships.
#[serde(default = "default_samples")]
pub samples_per_version: u32,
/// Pause between successive measured iterations against one model.
#[serde(default = "default_iter_pause")]
pub iteration_pause_secs: u64,
/// Per-request timeout (cold lazy-loads can be slow; generous like
/// bench.py's 600s default).
#[serde(default = "default_timeout")]
pub request_timeout_secs: u64,
/// SQLite system-of-record path.
#[serde(default = "default_db_path")]
pub db_path: String,
}
impl Default for BenchSettings {
fn default() -> Self {
BenchSettings {
sweep_interval_secs: default_sweep_interval(),
samples_per_version: default_samples(),
iteration_pause_secs: default_iter_pause(),
request_timeout_secs: default_timeout(),
db_path: default_db_path(),
}
}
}
impl BenchSettings {
pub fn iteration_pause(&self) -> Duration {
Duration::from_secs(self.iteration_pause_secs)
}
pub fn request_timeout(&self) -> Duration {
Duration::from_secs(self.request_timeout_secs)
}
pub fn sweep_interval(&self) -> Duration {
Duration::from_secs(self.sweep_interval_secs)
}
}
/// Which scenarios to run and their shared parameters.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScenarioConfig {
/// Approximate prompt sizes (in tokens) — one chat-latency scenario
/// is generated per size, e.g. `chat:128`, `chat:4096`. This is the
/// per-cell dimension that the version-aware skip logic keys on.
#[serde(default = "default_prompt_sizes")]
pub prompt_sizes: Vec<u32>,
/// Max generated tokens per request.
#[serde(default = "default_max_tokens")]
pub max_tokens: u64,
}
impl Default for ScenarioConfig {
fn default() -> Self {
ScenarioConfig {
prompt_sizes: default_prompt_sizes(),
max_tokens: default_max_tokens(),
}
}
}
/// One endpoint to benchmark.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TargetConfig {
/// Stable label used as the engine column and in the DB.
pub name: String,
/// Which protocol/metadata surface the target exposes.
#[serde(default)]
pub kind: TargetKind,
/// Base URL. For `neuron`: the daemon root (e.g.
/// `http://beast.internal:13131`). For `openai`: the OpenAI `/v1`
/// base (e.g. `http://host:8080/v1`).
pub endpoint: String,
/// Optional display label override for reports (defaults to `name`).
#[serde(default)]
pub label: Option<String>,
}
impl TargetConfig {
pub fn display_label(&self) -> &str {
self.label.as_deref().unwrap_or(&self.name)
}
}
/// The two target surfaces. `neuron` gets rich build metadata and warm
/// model discovery via the native neuron API; `openai` is the seam for
/// later comparison against mistral.rs / llama.cpp / vLLM (phase 1
/// implements `neuron` fully; `openai` is preliminary plumbing).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum TargetKind {
#[default]
Neuron,
Openai,
}
impl BenchConfig {
pub fn load(path: impl AsRef<Path>) -> Result<Self, Box<figment::Error>> {
Figment::new()
.merge(Toml::file(path))
.merge(Env::prefixed("BENCH_").split("__"))
.extract()
.map_err(Box::new)
}
}
fn default_sweep_interval() -> u64 {
1800
}
fn default_samples() -> u32 {
5
}
fn default_iter_pause() -> u64 {
2
}
fn default_timeout() -> u64 {
600
}
fn default_db_path() -> String {
"/var/lib/helexa-bench/bench.sqlite".to_string()
}
fn default_prompt_sizes() -> Vec<u32> {
vec![128, 4096]
}
fn default_max_tokens() -> u64 {
256
}
#[cfg(test)]
// Jail's closure must return figment::Result; the large-Err type is
// figment's, not ours, so suppress the lint here.
#[allow(clippy::result_large_err)]
mod tests {
use super::*;
use figment::Jail;
#[test]
fn loads_minimal_with_defaults() {
Jail::expect_with(|jail| {
jail.create_file(
"helexa-bench.toml",
r#"
[[targets]]
name = "beast"
endpoint = "http://beast.internal:13131"
"#,
)?;
let cfg = BenchConfig::load("helexa-bench.toml").unwrap();
assert_eq!(cfg.targets.len(), 1);
assert_eq!(cfg.targets[0].kind, TargetKind::Neuron);
assert_eq!(cfg.bench.samples_per_version, 5);
assert_eq!(cfg.scenarios.prompt_sizes, vec![128, 4096]);
Ok(())
});
}
#[test]
fn env_overrides_apply() {
Jail::expect_with(|jail| {
jail.create_file(
"helexa-bench.toml",
r#"
[bench]
samples_per_version = 3
[[targets]]
name = "benjy"
kind = "openai"
endpoint = "http://benjy:8080/v1"
"#,
)?;
jail.set_env("BENCH_BENCH__SAMPLES_PER_VERSION", "9");
let cfg = BenchConfig::load("helexa-bench.toml").unwrap();
assert_eq!(cfg.bench.samples_per_version, 9);
assert_eq!(cfg.targets[0].kind, TargetKind::Openai);
Ok(())
});
}
}

View File

@@ -0,0 +1,12 @@
//! helexa-bench — a continuous, version-aware benchmark harness for the
//! neuron fleet. It hits each neuron directly, exercises an extensible
//! scenario suite against every warm model, and records each run with
//! full build/version provenance into SQLite so improvements can be
//! tracked automatically across neuron implementation updates.
pub mod client;
pub mod config;
pub mod report;
pub mod scenario;
pub mod store;
pub mod sweep;

View File

@@ -0,0 +1,126 @@
//! helexa-bench CLI.
//!
//! - `run` — continuous daemon (systemd default): sweep, sleep, repeat.
//! - `once` — a single sweep, then exit (manual / CI).
//! - `report` — render the SQLite store as a results table.
//!
//! Runs on a single-threaded runtime: the workload is batch-1 sequential
//! (one request at a time, the regime we measure), and it lets the
//! SQLite connection live across awaits without `Sync` gymnastics.
use anyhow::{Context, Result};
use clap::{Parser, Subcommand};
use helexa_bench::config::BenchConfig;
use helexa_bench::report;
use helexa_bench::store::Store;
use helexa_bench::sweep::Sweeper;
use tracing_subscriber::EnvFilter;
#[derive(Parser)]
#[command(name = "helexa-bench")]
#[command(about = "Continuous version-aware benchmark harness for the neuron fleet")]
#[command(version)]
struct Cli {
#[command(subcommand)]
command: Command,
}
#[derive(Subcommand)]
enum Command {
/// Run sweeps continuously, pausing `sweep_interval_secs` between them.
Run {
#[arg(short, long, default_value = "helexa-bench.toml")]
config: String,
},
/// Run a single sweep over all targets, then exit.
Once {
#[arg(short, long, default_value = "helexa-bench.toml")]
config: String,
},
/// Render recorded results. Uses `--db` if given, else the db_path
/// from `--config`.
Report {
#[arg(short, long, default_value = "helexa-bench.toml")]
config: String,
/// Override the SQLite path (skips reading the config file).
#[arg(long)]
db: Option<String>,
/// Output format.
#[arg(long, default_value = "md")]
format: Format,
},
}
#[derive(Clone, Copy, clap::ValueEnum)]
enum Format {
Md,
Json,
}
fn main() -> Result<()> {
tracing_subscriber::fmt()
.with_env_filter(
EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")),
)
.init();
let cli = Cli::parse();
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.context("building tokio runtime")?;
rt.block_on(run(cli))
}
async fn run(cli: Cli) -> Result<()> {
match cli.command {
Command::Run { config } => {
let cfg = load_config(&config)?;
require_targets(&cfg)?;
let sweeper = Sweeper::new(cfg)?;
tracing::info!("helexa-bench started; entering continuous sweep loop");
sweeper.run_forever().await
}
Command::Once { config } => {
let cfg = load_config(&config)?;
require_targets(&cfg)?;
let sweeper = Sweeper::new(cfg)?;
let summary = sweeper.run_once().await?;
tracing::info!(
measured = summary.measured,
skipped = summary.skipped,
failed = summary.failed,
unreachable = summary.targets_unreachable,
"single sweep complete"
);
Ok(())
}
Command::Report { config, db, format } => {
let db_path = match db {
Some(p) => p,
None => load_config(&config)?.bench.db_path,
};
let store = Store::open(&db_path)?;
let rows = store.report_rows()?;
let rendered = match format {
Format::Md => report::render_markdown(&rows),
Format::Json => report::render_json(&rows)?,
};
println!("{rendered}");
Ok(())
}
}
}
fn load_config(path: &str) -> Result<BenchConfig> {
BenchConfig::load(path)
.map_err(|e| anyhow::anyhow!("{e}"))
.with_context(|| format!("loading config {path}"))
}
fn require_targets(cfg: &BenchConfig) -> Result<()> {
if cfg.targets.is_empty() {
anyhow::bail!("no targets configured — add at least one [[targets]] entry");
}
Ok(())
}

View File

@@ -0,0 +1,106 @@
//! Render the SQLite store as a results table — the automated
//! replacement for hand-editing `doc/benchmarks.md`. Columns match that
//! doc: engine, model, prompt tok, TTFT (s), decode tok/s, total (s),
//! plus the build SHA each cell was measured against.
use crate::store::ReportRow;
use anyhow::Result;
pub fn render_markdown(rows: &[ReportRow]) -> String {
let mut out = String::new();
out.push_str(
"| engine | model | prompt tok | TTFT (s) | decode tok/s | total (s) | build | n |\n",
);
out.push_str("|---|---|---:|---:|---:|---:|---|---:|\n");
for r in rows {
let ptok = r
.prompt_tokens
.map(|t| t.to_string())
.unwrap_or_else(|| format!("~{}", r.prompt_size_approx));
out.push_str(&format!(
"| {} | {} | {} | {} | {} | {} | `{}` | {} |\n",
r.target_name,
r.model_id,
ptok,
fmt_opt(r.ttft_s_median, 3),
fmt_opt(r.decode_tps_median, 1),
fmt_opt(r.total_s_median, 3),
r.git_sha,
r.samples,
));
}
out
}
pub fn render_json(rows: &[ReportRow]) -> Result<String> {
let arr: Vec<serde_json::Value> = rows
.iter()
.map(|r| {
serde_json::json!({
"engine": r.target_name,
"model": r.model_id,
"scenario": r.scenario_id,
"prompt_size_approx": r.prompt_size_approx,
"prompt_tokens": r.prompt_tokens,
"ttft_s_median": r.ttft_s_median,
"decode_tps_median": r.decode_tps_median,
"total_s_median": r.total_s_median,
"git_sha": r.git_sha,
"samples": r.samples,
})
})
.collect();
Ok(serde_json::to_string_pretty(&arr)?)
}
fn fmt_opt(v: Option<f64>, places: usize) -> String {
match v {
Some(x) => format!("{x:.places$}"),
None => "".to_string(),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn markdown_has_header_and_row() {
let rows = vec![ReportRow {
target_name: "beast".into(),
model_id: "Qwen/Qwen3.6-27B".into(),
scenario_id: "chat:128".into(),
prompt_size_approx: 128,
git_sha: "30d50d6".into(),
prompt_tokens: Some(130),
ttft_s_median: Some(0.123),
decode_tps_median: Some(45.6),
total_s_median: Some(1.234),
samples: 5,
}];
let md = render_markdown(&rows);
assert!(md.contains("| engine |"));
assert!(md.contains("beast"));
assert!(md.contains("`30d50d6`"));
assert!(md.contains("0.123"));
}
#[test]
fn missing_decode_renders_dash() {
let rows = vec![ReportRow {
target_name: "benjy".into(),
model_id: "m".into(),
scenario_id: "chat:128".into(),
prompt_size_approx: 128,
git_sha: "abc".into(),
prompt_tokens: None,
ttft_s_median: Some(0.1),
decode_tps_median: None,
total_s_median: Some(0.5),
samples: 1,
}];
let md = render_markdown(&rows);
assert!(md.contains("~128"));
assert!(md.contains(""));
}
}

View File

@@ -0,0 +1,238 @@
//! The extensible test suite.
//!
//! A [`Scenario`] puts one warm model through one shaped request and
//! reports operator-felt metrics (TTFT, decode tok/s, total). Phase 1
//! ships the chat-latency family ported faithfully from `script/bench.py`;
//! the trait is the seam for future families (vision, concurrency,
//! long-generation, cold-start) selected per model via [`Scenario::applies_to`].
use crate::config::ScenarioConfig;
use anyhow::{Context, Result, anyhow};
use async_trait::async_trait;
use cortex_core::harness::ModelInfo;
use cortex_core::openai::ChatCompletionChunk;
use eventsource_stream::Eventsource;
use futures::StreamExt;
use serde_json::json;
use std::time::{Duration, Instant};
/// A paragraph of filler re-used to synthesise prompts of a target
/// approximate token count (~4 chars/token heuristic — close enough for
/// bucketing; real token counts are read back from the usage object).
/// Mirrors `script/bench.py::FILLER`.
const FILLER: &str = "The quick brown fox jumps over the lazy dog while the band plays \
a slow waltz in the background and somebody counts the beats. ";
/// `/no_think`: Qwen3-family soft switch keeping thinking models from
/// burning the token budget invisibly. Harmless for non-thinking models.
const QUESTION: &str = "\n\nRetell the scene above as a vivid story of about 300 words. /no_think";
/// Build a synthetic prompt of approximately `approx_tokens` tokens.
/// Ported from `bench.py::build_prompt`.
pub fn build_prompt(approx_tokens: u32) -> String {
let target_chars = (approx_tokens.max(16) as usize) * 4;
let reps = target_chars / FILLER.len() + 1;
let mut body = FILLER.repeat(reps);
body.truncate(target_chars);
body.push_str(QUESTION);
body
}
/// Per-request inputs shared by every scenario.
pub struct RunCtx<'a> {
pub client: &'a reqwest::Client,
/// Fully-qualified chat-completions URL for the target.
pub chat_url: String,
pub model_id: String,
pub max_tokens: u64,
pub timeout: Duration,
}
/// Operator-felt metrics for a single measured request.
#[derive(Debug, Clone)]
pub struct ScenarioMetrics {
/// Time to first content chunk (seconds).
pub ttft_s: f64,
/// Completion tokens / decode window. `None` when the window is too
/// short to be honest (≤ 200 ms), matching bench.py.
pub decode_tps: Option<f64>,
/// Wall-clock for the whole request (seconds).
pub total_s: f64,
/// Prompt tokens from the final `usage` object, if the server sent one.
pub prompt_tokens: Option<u64>,
/// Completion tokens: from `usage` when present, else content-chunk count.
pub completion_tokens: u64,
}
#[async_trait]
pub trait Scenario: Send + Sync {
/// Stable id, e.g. `chat:128`. Used as the version-aware skip key
/// dimension and recorded against every run.
fn id(&self) -> &str;
/// Approximate prompt size in tokens (the cell dimension), recorded
/// for reporting.
fn prompt_size(&self) -> u32;
/// Whether this scenario should run against the given model. Default
/// runs against everything; vision/audio scenarios will gate on
/// [`ModelInfo::capabilities`].
fn applies_to(&self, _model: &ModelInfo) -> bool {
true
}
/// Issue one shaped request and measure it.
async fn run(&self, ctx: &RunCtx) -> Result<ScenarioMetrics>;
}
/// Build the active scenario set from config. One chat-latency scenario
/// per configured prompt size.
pub fn build_scenarios(cfg: &ScenarioConfig) -> Vec<Box<dyn Scenario>> {
cfg.prompt_sizes
.iter()
.map(|&size| {
Box::new(ChatLatencyScenario {
id: format!("chat:{size}"),
approx_prompt_tokens: size,
}) as Box<dyn Scenario>
})
.collect()
}
/// Streamed single-request chat-completions latency probe — the batch-1
/// regime bench.py measures.
pub struct ChatLatencyScenario {
id: String,
approx_prompt_tokens: u32,
}
#[async_trait]
impl Scenario for ChatLatencyScenario {
fn id(&self) -> &str {
&self.id
}
fn prompt_size(&self) -> u32 {
self.approx_prompt_tokens
}
async fn run(&self, ctx: &RunCtx) -> Result<ScenarioMetrics> {
let prompt = build_prompt(self.approx_prompt_tokens);
let payload = json!({
"model": ctx.model_id,
"messages": [{"role": "user", "content": prompt}],
"max_tokens": ctx.max_tokens,
"temperature": 0,
"stream": true,
"stream_options": {"include_usage": true},
});
let fut = stream_and_measure(ctx, &payload);
tokio::time::timeout(ctx.timeout, fut)
.await
.map_err(|_| anyhow!("request timed out after {:?}", ctx.timeout))?
}
}
/// The SSE-timing core, ported from `bench.py::one_run`. Kept free of the
/// `Scenario` trait so it's unit-testable against a mock byte stream.
async fn stream_and_measure(
ctx: &RunCtx<'_>,
payload: &serde_json::Value,
) -> Result<ScenarioMetrics> {
let start = Instant::now();
let resp = ctx
.client
.post(&ctx.chat_url)
.json(payload)
.send()
.await
.context("sending chat request")?;
if !resp.status().is_success() {
let status = resp.status();
let body = resp.text().await.unwrap_or_default();
return Err(anyhow!("upstream returned {status}: {}", body.trim()));
}
let mut stream = resp.bytes_stream().eventsource();
let mut first: Option<Instant> = None;
let mut last: Option<Instant> = None;
let mut chunk_count: u64 = 0;
let mut prompt_tokens: Option<u64> = None;
let mut completion_tokens: Option<u64> = None;
while let Some(event) = stream.next().await {
let event = event.context("reading SSE stream")?;
let now = Instant::now();
let data = event.data.trim();
if data.is_empty() || data == "[DONE]" {
continue;
}
let chunk: ChatCompletionChunk = match serde_json::from_str(data) {
Ok(c) => c,
Err(_) => continue, // tolerate non-JSON keepalive frames
};
if let Some(choice) = chunk.choices.first()
&& choice
.delta
.get("content")
.and_then(|c| c.as_str())
.is_some_and(|s| !s.is_empty())
{
if first.is_none() {
first = Some(now);
}
last = Some(now);
chunk_count += 1;
}
if let Some(usage) = chunk.usage {
prompt_tokens = Some(usage.prompt_tokens);
completion_tokens = Some(usage.completion_tokens);
}
}
let end = Instant::now();
let first = first.ok_or_else(|| anyhow!("no content chunks received"))?;
// neuron emits one SSE chunk per visible token, so chunk_count is an
// engine-truth count when no usage frame is sent.
let tokens = completion_tokens.filter(|&t| t > 0).unwrap_or(chunk_count);
// decode rate is only meaningful over a real inter-chunk window.
let window = last
.filter(|&l| l > first)
.map(|l| (l - first).as_secs_f64())
.unwrap_or(0.0);
Ok(ScenarioMetrics {
ttft_s: (first - start).as_secs_f64(),
decode_tps: if window > 0.2 {
Some(tokens as f64 / window)
} else {
None
},
total_s: (end - start).as_secs_f64(),
prompt_tokens,
completion_tokens: tokens,
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn prompt_grows_with_token_target() {
let small = build_prompt(128);
let big = build_prompt(4096);
assert!(big.len() > small.len());
// ~4 chars/token + the trailing question.
assert!(small.len() >= 128 * 4);
assert!(small.ends_with("/no_think"));
}
#[test]
fn prompt_floor_for_tiny_targets() {
// max(approx,16) floor means even 0 yields a non-trivial prompt.
let p = build_prompt(0);
assert!(p.len() >= 16 * 4);
}
}

View File

@@ -0,0 +1,400 @@
//! SQLite system-of-record. One row per measured iteration, keyed so a
//! benchmark can be attributed to the exact neuron build that produced
//! it. Replaces hand edits to `doc/benchmarks.md`.
//!
//! Calls are synchronous (SQLite is local and the sweep is batch-1
//! sequential), so the connection is used inline between `await` points,
//! never held across one.
use anyhow::{Context, Result};
use rusqlite::{Connection, params};
use std::path::Path;
/// A single measured (or failed) iteration, with full provenance.
#[derive(Debug, Clone)]
pub struct RunRecord {
pub ts: String, // RFC3339
// target
pub target_name: String,
pub target_kind: String,
pub endpoint: String,
// host (from /discovery)
pub hostname: Option<String>,
pub driver_version: Option<String>,
pub cuda_version: Option<String>,
pub gpus_json: Option<String>,
// neuron build (from /version)
pub git_sha: String,
pub git_sha_long: Option<String>,
pub package_version: String,
pub git_dirty: bool,
pub build_timestamp: Option<String>,
pub rustc_version: Option<String>,
pub profile: Option<String>,
pub features_json: String,
pub candle_version: Option<String>,
// bench's own build
pub bench_version: String,
pub bench_sha: String,
// model
pub model_id: String,
pub harness: String,
pub capabilities_json: String,
pub devices_json: String,
// scenario
pub scenario_id: String,
pub prompt_size_approx: u32,
pub prompt_tokens_actual: Option<u64>,
pub max_tokens: u64,
// metrics
pub ttft_s: Option<f64>,
pub decode_tps: Option<f64>,
pub total_s: Option<f64>,
pub completion_tokens: Option<u64>,
// outcome
pub ok: bool,
pub error: Option<String>,
}
pub struct Store {
conn: Connection,
}
impl Store {
/// Open (creating parent dirs + schema as needed).
pub fn open(path: impl AsRef<Path>) -> Result<Self> {
let path = path.as_ref();
if let Some(parent) = path.parent()
&& !parent.as_os_str().is_empty()
{
std::fs::create_dir_all(parent)
.with_context(|| format!("creating db dir {}", parent.display()))?;
}
let conn = Connection::open(path)
.with_context(|| format!("opening sqlite db {}", path.display()))?;
Self::init(&conn)?;
Ok(Store { conn })
}
/// In-memory store for tests.
#[cfg(test)]
pub fn open_in_memory() -> Result<Self> {
let conn = Connection::open_in_memory()?;
Self::init(&conn)?;
Ok(Store { conn })
}
fn init(conn: &Connection) -> Result<()> {
conn.execute_batch(
r#"
CREATE TABLE IF NOT EXISTS runs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
ts TEXT NOT NULL,
target_name TEXT NOT NULL,
target_kind TEXT NOT NULL,
endpoint TEXT NOT NULL,
hostname TEXT,
driver_version TEXT,
cuda_version TEXT,
gpus_json TEXT,
git_sha TEXT NOT NULL,
git_sha_long TEXT,
package_version TEXT NOT NULL,
git_dirty INTEGER NOT NULL,
build_timestamp TEXT,
rustc_version TEXT,
profile TEXT,
features_json TEXT NOT NULL,
candle_version TEXT,
bench_version TEXT NOT NULL,
bench_sha TEXT NOT NULL,
model_id TEXT NOT NULL,
harness TEXT NOT NULL,
capabilities_json TEXT NOT NULL,
devices_json TEXT NOT NULL,
scenario_id TEXT NOT NULL,
prompt_size_approx INTEGER NOT NULL,
prompt_tokens_actual INTEGER,
max_tokens INTEGER NOT NULL,
ttft_s REAL,
decode_tps REAL,
total_s REAL,
completion_tokens INTEGER,
ok INTEGER NOT NULL,
error TEXT
);
-- The version-aware skip query keys on this tuple. scenario_id
-- encodes the prompt size (chat:<n>), so it subsumes the cell.
CREATE INDEX IF NOT EXISTS idx_runs_cell
ON runs (target_name, git_sha, model_id, scenario_id, ok);
"#,
)
.context("initialising sqlite schema")?;
Ok(())
}
/// Count successful samples already recorded for a cell. Only `ok`
/// rows count toward the per-version target so transient failures
/// don't permanently starve a cell.
pub fn count_samples(
&self,
target_name: &str,
git_sha: &str,
model_id: &str,
scenario_id: &str,
) -> Result<u32> {
let n: i64 = self.conn.query_row(
"SELECT COUNT(*) FROM runs WHERE target_name=?1 AND git_sha=?2 \
AND model_id=?3 AND scenario_id=?4 AND ok=1",
params![target_name, git_sha, model_id, scenario_id],
|row| row.get(0),
)?;
Ok(n as u32)
}
pub fn insert_run(&self, r: &RunRecord) -> Result<()> {
self.conn.execute(
"INSERT INTO runs (
ts, target_name, target_kind, endpoint,
hostname, driver_version, cuda_version, gpus_json,
git_sha, git_sha_long, package_version, git_dirty,
build_timestamp, rustc_version, profile, features_json, candle_version,
bench_version, bench_sha,
model_id, harness, capabilities_json, devices_json,
scenario_id, prompt_size_approx, prompt_tokens_actual, max_tokens,
ttft_s, decode_tps, total_s, completion_tokens,
ok, error
) VALUES (
?1, ?2, ?3, ?4,
?5, ?6, ?7, ?8,
?9, ?10, ?11, ?12,
?13, ?14, ?15, ?16, ?17,
?18, ?19,
?20, ?21, ?22, ?23,
?24, ?25, ?26, ?27,
?28, ?29, ?30, ?31,
?32, ?33
)",
params![
r.ts,
r.target_name,
r.target_kind,
r.endpoint,
r.hostname,
r.driver_version,
r.cuda_version,
r.gpus_json,
r.git_sha,
r.git_sha_long,
r.package_version,
r.git_dirty as i64,
r.build_timestamp,
r.rustc_version,
r.profile,
r.features_json,
r.candle_version,
r.bench_version,
r.bench_sha,
r.model_id,
r.harness,
r.capabilities_json,
r.devices_json,
r.scenario_id,
r.prompt_size_approx,
r.prompt_tokens_actual,
r.max_tokens,
r.ttft_s,
r.decode_tps,
r.total_s,
r.completion_tokens,
r.ok as i64,
r.error,
],
)?;
Ok(())
}
/// One reportable cell: the median metrics over the most-recently-seen
/// build SHA for each (target, model, scenario).
pub fn report_rows(&self) -> Result<Vec<ReportRow>> {
// For each (target, model, scenario), find the SHA of the latest
// successful run, then median that SHA's samples.
let mut stmt = self.conn.prepare(
"SELECT target_name, model_id, scenario_id, prompt_size_approx, git_sha,
ttft_s, decode_tps, total_s, prompt_tokens_actual
FROM runs
WHERE ok=1
ORDER BY target_name, model_id, scenario_id, id",
)?;
let rows = stmt.query_map([], |row| {
Ok(RawRow {
target_name: row.get(0)?,
model_id: row.get(1)?,
scenario_id: row.get(2)?,
prompt_size_approx: row.get(3)?,
git_sha: row.get(4)?,
ttft_s: row.get(5)?,
decode_tps: row.get(6)?,
total_s: row.get(7)?,
prompt_tokens_actual: row.get(8)?,
})
})?;
let raws: Vec<RawRow> = rows.collect::<rusqlite::Result<_>>()?;
Ok(aggregate(raws))
}
}
struct RawRow {
target_name: String,
model_id: String,
scenario_id: String,
prompt_size_approx: u32,
git_sha: String,
ttft_s: Option<f64>,
decode_tps: Option<f64>,
total_s: Option<f64>,
prompt_tokens_actual: Option<u64>,
}
/// An aggregated cell ready for the report table.
#[derive(Debug, Clone, PartialEq)]
pub struct ReportRow {
pub target_name: String,
pub model_id: String,
pub scenario_id: String,
pub prompt_size_approx: u32,
pub git_sha: String,
pub prompt_tokens: Option<u64>,
pub ttft_s_median: Option<f64>,
pub decode_tps_median: Option<f64>,
pub total_s_median: Option<f64>,
pub samples: usize,
}
/// Group by (target, model, scenario), keep only the latest SHA's rows
/// (latest = the SHA of the last-inserted row, since input is id-ordered),
/// and median each metric.
fn aggregate(raws: Vec<RawRow>) -> Vec<ReportRow> {
use std::collections::BTreeMap;
// key -> (latest_sha, rows for that sha)
let mut groups: BTreeMap<(String, String, String), Vec<RawRow>> = BTreeMap::new();
for r in raws {
groups
.entry((
r.target_name.clone(),
r.model_id.clone(),
r.scenario_id.clone(),
))
.or_default()
.push(r);
}
let mut out = Vec::new();
for ((target_name, model_id, scenario_id), rows) in groups {
// id-ordered, so the last row carries the latest SHA.
let latest_sha = rows.last().map(|r| r.git_sha.clone()).unwrap_or_default();
let cell: Vec<&RawRow> = rows.iter().filter(|r| r.git_sha == latest_sha).collect();
let prompt_size_approx = cell.first().map(|r| r.prompt_size_approx).unwrap_or(0);
out.push(ReportRow {
target_name,
model_id,
scenario_id,
prompt_size_approx,
git_sha: latest_sha,
prompt_tokens: cell.iter().find_map(|r| r.prompt_tokens_actual),
ttft_s_median: median(cell.iter().filter_map(|r| r.ttft_s)),
decode_tps_median: median(cell.iter().filter_map(|r| r.decode_tps)),
total_s_median: median(cell.iter().filter_map(|r| r.total_s)),
samples: cell.len(),
});
}
out
}
fn median(values: impl Iterator<Item = f64>) -> Option<f64> {
let mut v: Vec<f64> = values.collect();
if v.is_empty() {
return None;
}
v.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
// lo == hi for odd lengths (the middle element); they straddle the
// centre for even lengths. Avoids a `% 2` branch.
let lo = (v.len() - 1) / 2;
let hi = v.len() / 2;
Some((v[lo] + v[hi]) / 2.0)
}
#[cfg(test)]
mod tests {
use super::*;
fn rec(target: &str, sha: &str, model: &str, scenario: &str, ok: bool) -> RunRecord {
RunRecord {
ts: "2026-06-13T00:00:00Z".into(),
target_name: target.into(),
target_kind: "neuron".into(),
endpoint: "http://x:13131".into(),
hostname: Some("x".into()),
driver_version: None,
cuda_version: None,
gpus_json: None,
git_sha: sha.into(),
git_sha_long: None,
package_version: "0.1.16".into(),
git_dirty: false,
build_timestamp: None,
rustc_version: None,
profile: None,
features_json: "[]".into(),
candle_version: None,
bench_version: "0.1.16".into(),
bench_sha: "deadbee".into(),
model_id: model.into(),
harness: "candle".into(),
capabilities_json: "[]".into(),
devices_json: "[]".into(),
scenario_id: scenario.into(),
prompt_size_approx: 128,
prompt_tokens_actual: Some(130),
max_tokens: 256,
ttft_s: Some(0.1),
decode_tps: Some(50.0),
total_s: Some(1.0),
completion_tokens: Some(50),
ok,
error: if ok { None } else { Some("boom".into()) },
}
}
#[test]
fn counts_only_successful_samples() {
let s = Store::open_in_memory().unwrap();
s.insert_run(&rec("beast", "abc", "m", "chat:128", true))
.unwrap();
s.insert_run(&rec("beast", "abc", "m", "chat:128", true))
.unwrap();
s.insert_run(&rec("beast", "abc", "m", "chat:128", false))
.unwrap();
assert_eq!(s.count_samples("beast", "abc", "m", "chat:128").unwrap(), 2);
// Different SHA is a different cell.
assert_eq!(s.count_samples("beast", "xyz", "m", "chat:128").unwrap(), 0);
}
#[test]
fn report_uses_latest_sha_per_cell() {
let s = Store::open_in_memory().unwrap();
// old build
s.insert_run(&rec("beast", "old", "m", "chat:128", true))
.unwrap();
// new build, two samples
let mut r = rec("beast", "new", "m", "chat:128", true);
r.ttft_s = Some(0.2);
s.insert_run(&r).unwrap();
r.ttft_s = Some(0.4);
s.insert_run(&r).unwrap();
let rows = s.report_rows().unwrap();
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].git_sha, "new");
assert_eq!(rows[0].samples, 2);
assert!((rows[0].ttft_s_median.unwrap() - 0.3).abs() < 1e-9);
}
}

View File

@@ -0,0 +1,250 @@
//! The version-aware sweep loop.
//!
//! Each sweep visits every configured target, polls its build identity
//! and warm models, and tops up benchmark samples per
//! (target, build SHA, model, scenario) to `samples_per_version`. Cells
//! already at target are skipped — so once every neuron's current build
//! is fully sampled, sweeps cost only the cheap metadata polls until a
//! new SHA ships. Runs are recorded to SQLite with full provenance.
use crate::client::TargetClient;
use crate::config::{BenchConfig, TargetConfig, TargetKind};
use crate::scenario::{RunCtx, build_scenarios};
use crate::store::{RunRecord, Store};
use anyhow::Result;
use cortex_core::build_info::BuildInfo;
use cortex_core::discovery::DiscoveryResponse;
use cortex_core::harness::ModelInfo;
/// helexa-bench's own build version.
fn bench_version() -> String {
env!("CARGO_PKG_VERSION").to_string()
}
/// helexa-bench's own build SHA, injected by CI via `HELEXA_BUILD_SHA`
/// at compile time; `"unknown"` for ad-hoc local builds.
fn bench_sha() -> String {
option_env!("HELEXA_BUILD_SHA")
.filter(|s| !s.is_empty())
.unwrap_or("unknown")
.to_string()
}
#[derive(Debug, Default, Clone)]
pub struct SweepSummary {
pub measured: usize,
pub skipped: usize,
pub failed: usize,
pub targets_unreachable: usize,
}
pub struct Sweeper {
cfg: BenchConfig,
client: TargetClient,
store: Store,
}
impl Sweeper {
pub fn new(cfg: BenchConfig) -> Result<Self> {
let client = TargetClient::new(cfg.bench.request_timeout())?;
let store = Store::open(&cfg.bench.db_path)?;
Ok(Sweeper { cfg, client, store })
}
/// Run sweeps forever, pausing `sweep_interval` between them.
pub async fn run_forever(&self) -> ! {
loop {
match self.run_once().await {
Ok(s) => tracing::info!(
measured = s.measured,
skipped = s.skipped,
failed = s.failed,
unreachable = s.targets_unreachable,
"sweep complete"
),
Err(e) => tracing::error!(error = %format!("{e:#}"), "sweep errored"),
}
tracing::debug!(
secs = self.cfg.bench.sweep_interval_secs,
"sleeping until next sweep"
);
tokio::time::sleep(self.cfg.bench.sweep_interval()).await;
}
}
/// One full pass over all targets.
pub async fn run_once(&self) -> Result<SweepSummary> {
let mut summary = SweepSummary::default();
for target in &self.cfg.targets {
if let Err(e) = self.sweep_target(target, &mut summary).await {
summary.targets_unreachable += 1;
tracing::warn!(target = %target.name, error = %format!("{e:#}"), "target skipped");
}
}
Ok(summary)
}
async fn sweep_target(&self, target: &TargetConfig, summary: &mut SweepSummary) -> Result<()> {
let build = self.client.fetch_version(target).await?;
let discovery = self.client.fetch_discovery(target).await.unwrap_or(None);
let models = self.client.warm_models(target).await?;
tracing::info!(
target = %target.name,
sha = %build.git_sha,
warm_models = models.len(),
"sweeping target"
);
let scenarios = build_scenarios(&self.cfg.scenarios);
for model in &models {
for scenario in scenarios.iter().filter(|s| s.applies_to(model)) {
let have = self.store.count_samples(
&target.name,
&build.git_sha,
&model.id,
scenario.id(),
)?;
let need = self.cfg.bench.samples_per_version.saturating_sub(have);
if need == 0 {
summary.skipped += 1;
tracing::debug!(
target = %target.name, model = %model.id, scenario = scenario.id(),
sha = %build.git_sha, "cell already satisfied, skipping"
);
continue;
}
let ctx = RunCtx {
client: self.client.http(),
chat_url: self.client.chat_url(target),
model_id: model.id.clone(),
max_tokens: self.cfg.scenarios.max_tokens,
timeout: self.cfg.bench.request_timeout(),
};
// One unmeasured warmup when the cell is empty (matches
// bench.py — first run after a load hits cold caches).
if have == 0 {
tracing::debug!(model = %model.id, scenario = scenario.id(), "warmup run");
let _ = scenario.run(&ctx).await;
}
for i in 0..need {
match scenario.run(&ctx).await {
Ok(m) => {
let rec = self.build_record(
target,
&build,
discovery.as_ref(),
model,
scenario.id(),
scenario.prompt_size(),
Ok(&m),
);
self.store.insert_run(&rec)?;
summary.measured += 1;
tracing::info!(
target = %target.name, model = %model.id, scenario = scenario.id(),
ttft_s = m.ttft_s, decode_tps = ?m.decode_tps, total_s = m.total_s,
"{}/{} recorded", have + i + 1, self.cfg.bench.samples_per_version
);
}
Err(e) => {
let msg = format!("{e:#}");
let rec = self.build_record(
target,
&build,
discovery.as_ref(),
model,
scenario.id(),
scenario.prompt_size(),
Err(&msg),
);
self.store.insert_run(&rec)?;
summary.failed += 1;
tracing::warn!(
target = %target.name, model = %model.id, scenario = scenario.id(),
error = %msg, "iteration failed"
);
}
}
tokio::time::sleep(self.cfg.bench.iteration_pause()).await;
}
}
}
Ok(())
}
#[allow(clippy::too_many_arguments)]
fn build_record(
&self,
target: &TargetConfig,
build: &BuildInfo,
discovery: Option<&DiscoveryResponse>,
model: &ModelInfo,
scenario_id: &str,
prompt_size: u32,
result: Result<&crate::scenario::ScenarioMetrics, &str>,
) -> RunRecord {
let (ok, error, ttft, decode, total, prompt_tokens, completion) = match result {
Ok(m) => (
true,
None,
Some(m.ttft_s),
m.decode_tps,
Some(m.total_s),
m.prompt_tokens,
Some(m.completion_tokens),
),
Err(e) => (false, Some(e.to_string()), None, None, None, None, None),
};
RunRecord {
ts: chrono::Utc::now().to_rfc3339(),
target_name: target.name.clone(),
target_kind: kind_str(target.kind).to_string(),
endpoint: target.endpoint.clone(),
hostname: discovery.map(|d| d.hostname.clone()),
driver_version: discovery.and_then(|d| d.driver_version.clone()),
cuda_version: discovery.and_then(|d| d.cuda_version.clone()),
gpus_json: discovery
.map(|d| serde_json::to_string(&d.devices).unwrap_or_else(|_| "[]".to_string())),
git_sha: build.git_sha.clone(),
git_sha_long: build.git_sha_long.clone(),
package_version: build.package_version.clone(),
git_dirty: build.git_dirty,
build_timestamp: build.build_timestamp.clone(),
rustc_version: build.rustc_version.clone(),
profile: build.profile.clone(),
features_json: serde_json::to_string(&build.features)
.unwrap_or_else(|_| "[]".to_string()),
candle_version: build.candle_version.clone(),
bench_version: bench_version(),
bench_sha: bench_sha(),
model_id: model.id.clone(),
harness: model.harness.clone(),
capabilities_json: serde_json::to_string(&model.capabilities)
.unwrap_or_else(|_| "[]".to_string()),
devices_json: serde_json::to_string(&model.devices)
.unwrap_or_else(|_| "[]".to_string()),
scenario_id: scenario_id.to_string(),
prompt_size_approx: prompt_size,
prompt_tokens_actual: prompt_tokens,
max_tokens: self.cfg.scenarios.max_tokens,
ttft_s: ttft,
decode_tps: decode,
total_s: total,
completion_tokens: completion,
ok,
error,
}
}
}
fn kind_str(kind: TargetKind) -> &'static str {
match kind {
TargetKind::Neuron => "neuron",
TargetKind::Openai => "openai",
}
}

View File

@@ -0,0 +1,132 @@
//! End-to-end sweep against a mock neuron: a sweep records samples, a
//! second sweep skips the satisfied cell, and bumping the reported build
//! SHA resumes fresh sampling.
use axum::Router;
use axum::extract::State;
use axum::http::header;
use axum::response::{IntoResponse, Json};
use axum::routing::{get, post};
use helexa_bench::config::{BenchConfig, BenchSettings, ScenarioConfig, TargetConfig, TargetKind};
use helexa_bench::sweep::Sweeper;
use serde_json::json;
use std::sync::{Arc, Mutex};
#[derive(Clone)]
struct MockState {
sha: Arc<Mutex<String>>,
}
async fn version(State(s): State<MockState>) -> Json<serde_json::Value> {
let sha = s.sha.lock().unwrap().clone();
Json(json!({
"package_version": "0.1.16",
"git_sha": sha,
"git_dirty": false,
"features": ["cuda", "cudnn"],
"candle_version": "0.10.2",
}))
}
async fn discovery() -> Json<serde_json::Value> {
Json(json!({
"hostname": "mock-beast",
"os": "Linux",
"kernel": "6.19.0",
"cuda_version": "13.0",
"driver_version": "580.159",
"devices": [{"index": 0, "name": "RTX 5090", "vram_total_mb": 32614, "compute_capability": "12.0"}],
"harnesses": ["candle"],
}))
}
async fn models() -> Json<serde_json::Value> {
Json(json!([
{"id": "Qwen/Qwen3.6-27B", "harness": "candle", "status": "loaded", "devices": [0], "capabilities": ["text"]},
// A non-warm model the bench must ignore.
{"id": "Qwen/cold", "harness": "candle", "status": "recovering", "devices": [0]},
]))
}
async fn chat() -> impl IntoResponse {
let body = concat!(
"data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"Hello\"},\"finish_reason\":null}]}\n\n",
"data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\" world\"},\"finish_reason\":null}]}\n\n",
"data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":130,\"completion_tokens\":2,\"total_tokens\":132}}\n\n",
"data: [DONE]\n\n",
);
([(header::CONTENT_TYPE, "text/event-stream")], body)
}
async fn spawn_mock(sha: &str) -> (String, Arc<Mutex<String>>) {
let shared = Arc::new(Mutex::new(sha.to_string()));
let state = MockState {
sha: shared.clone(),
};
let app = Router::new()
.route("/version", get(version))
.route("/discovery", get(discovery))
.route("/models", get(models))
.route("/v1/chat/completions", post(chat))
.with_state(state);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
axum::serve(listener, app).await.unwrap();
});
(format!("http://{addr}"), shared)
}
fn config_for(endpoint: String, db_path: String) -> BenchConfig {
BenchConfig {
bench: BenchSettings {
sweep_interval_secs: 1,
samples_per_version: 2,
iteration_pause_secs: 0,
request_timeout_secs: 30,
db_path,
},
scenarios: ScenarioConfig {
prompt_sizes: vec![128], // single scenario keeps assertions simple
max_tokens: 16,
},
targets: vec![TargetConfig {
name: "mock".into(),
kind: TargetKind::Neuron,
endpoint,
label: None,
}],
}
}
#[tokio::test]
async fn sweep_records_skips_and_resumes_on_new_sha() {
let (endpoint, sha_handle) = spawn_mock("aaaaaaa").await;
// Unique db path per run (bound port is unique).
let port = endpoint.rsplit(':').next().unwrap();
let db_path = std::env::temp_dir().join(format!("helexa-bench-it-{port}.sqlite"));
let _ = std::fs::remove_file(&db_path);
let db_str = db_path.to_string_lossy().to_string();
let sweeper = Sweeper::new(config_for(endpoint, db_str)).unwrap();
// First sweep: one warm model × one scenario × 2 samples.
let s1 = sweeper.run_once().await.unwrap();
assert_eq!(s1.measured, 2, "should record samples_per_version samples");
assert_eq!(s1.skipped, 0);
assert_eq!(s1.failed, 0);
// Second sweep at same SHA: cell satisfied, nothing measured.
let s2 = sweeper.run_once().await.unwrap();
assert_eq!(s2.measured, 0, "satisfied cell must be skipped");
assert_eq!(s2.skipped, 1);
// Bump the reported build SHA: a new cell → fresh sampling resumes.
*sha_handle.lock().unwrap() = "bbbbbbb".to_string();
let s3 = sweeper.run_once().await.unwrap();
assert_eq!(s3.measured, 2, "new SHA must resume sampling");
assert_eq!(s3.skipped, 0);
let _ = std::fs::remove_file(&db_path);
}

View File

@@ -60,6 +60,11 @@ tokio-stream.workspace = true
figment.workspace = true
toml.workspace = true
# Parallel in-situ quantization (#1): fans candle's per-block k-quant
# math across the CPU pool at model-load time. Already in the tree
# transitively via candle-core.
rayon = "1"
# candle for in-process inference. CUDA support is gated behind the
# crate's `cuda` feature (default off) so the workspace builds on
# non-CUDA hosts and CI runners.

View File

@@ -1,10 +1,16 @@
//! Build script: compile the CUDA kernels in `src/cuda/*.cu` into a
//! static library and link it under the `cuda` feature.
//! Build script: capture build/version metadata for `GET /version`,
//! and (under the `cuda` feature) compile the CUDA kernels in
//! `src/cuda/*.cu` into a static library and link it.
//!
//! Patterned on `EricLBuehler/mistral.rs::mistralrs-core/build.rs` —
//! same `cudaforge::KernelBuilder` invocation, same NVCC flag set.
//! The CUDA portion is patterned on
//! `EricLBuehler/mistral.rs::mistralrs-core/build.rs` — same
//! `cudaforge::KernelBuilder` invocation, same NVCC flag set.
use std::process::Command;
fn main() {
emit_build_metadata();
#[cfg(feature = "cuda")]
{
use std::path::PathBuf;
@@ -64,3 +70,127 @@ fn main() {
}
}
}
/// Emit `cargo:rustc-env=` vars consumed by `env!()` in `src/version.rs`
/// so the daemon can report its own build identity from `GET /version`.
///
/// We re-run only when HEAD moves or the SHA override changes — not on
/// every compile — so the captured timestamp is stable for a given
/// build input rather than churning on each `cargo build`.
fn emit_build_metadata() {
println!("cargo:rerun-if-env-changed=HELEXA_BUILD_SHA");
println!("cargo:rerun-if-changed=.git/HEAD");
// A detached/normal HEAD points at a ref whose file is what actually
// changes on commit; watch the packed-refs fallback too.
println!("cargo:rerun-if-changed=.git/packed-refs");
// SHA: prefer the CI/RPM-injected override (tarball builds have no
// .git), then fall back to git, then to "unknown".
let (sha_short, sha_long, dirty) = match std::env::var("HELEXA_BUILD_SHA") {
Ok(s) if !s.trim().is_empty() => {
let s = s.trim().to_string();
let short = s.chars().take(7).collect::<String>();
(short, Some(s), false)
}
_ => {
let long = git(&["rev-parse", "HEAD"]);
let short = git(&["rev-parse", "--short", "HEAD"]);
let dirty = git(&["status", "--porcelain"])
.map(|s| !s.trim().is_empty())
.unwrap_or(false);
match short {
Some(short) => (short, long, dirty),
None => ("unknown".to_string(), None, false),
}
}
};
println!("cargo:rustc-env=HELEXA_GIT_SHA={sha_short}");
println!(
"cargo:rustc-env=HELEXA_GIT_SHA_LONG={}",
sha_long.unwrap_or_default()
);
println!("cargo:rustc-env=HELEXA_GIT_DIRTY={dirty}");
// RFC3339 build timestamp. `date` is universally present on the
// Linux hosts neuron targets; empty if it ever isn't.
let ts = Command::new("date")
.args(["-u", "+%Y-%m-%dT%H:%M:%SZ"])
.output()
.ok()
.filter(|o| o.status.success())
.map(|o| String::from_utf8_lossy(&o.stdout).trim().to_string())
.unwrap_or_default();
println!("cargo:rustc-env=HELEXA_BUILD_TIMESTAMP={ts}");
// Compiler version: cargo sets $RUSTC to the rustc it invokes.
let rustc = std::env::var("RUSTC").unwrap_or_else(|_| "rustc".to_string());
let rustc_version = Command::new(rustc)
.arg("--version")
.output()
.ok()
.filter(|o| o.status.success())
.map(|o| String::from_utf8_lossy(&o.stdout).trim().to_string())
.unwrap_or_default();
println!("cargo:rustc-env=HELEXA_RUSTC_VERSION={rustc_version}");
println!(
"cargo:rustc-env=HELEXA_BUILD_PROFILE={}",
std::env::var("PROFILE").unwrap_or_default()
);
println!(
"cargo:rustc-env=HELEXA_TARGET={}",
std::env::var("TARGET").unwrap_or_default()
);
// Enabled features: cargo exports CARGO_FEATURE_<NAME> for each.
// Reverse the mangling (uppercase, '-'→'_') best-effort for display.
let mut features: Vec<String> = std::env::vars()
.filter_map(|(k, _)| k.strip_prefix("CARGO_FEATURE_").map(|f| f.to_string()))
.map(|f| f.to_lowercase().replace('_', "-"))
// `default` is the meta-feature, not a perf-relevant flag.
.filter(|f| f != "default")
.collect();
features.sort();
println!("cargo:rustc-env=HELEXA_FEATURES={}", features.join(","));
println!(
"cargo:rustc-env=HELEXA_CANDLE_VERSION={}",
candle_version().unwrap_or_default()
);
}
fn git(args: &[&str]) -> Option<String> {
let out = Command::new("git").args(args).output().ok()?;
if !out.status.success() {
return None;
}
let s = String::from_utf8_lossy(&out.stdout).trim().to_string();
if s.is_empty() { None } else { Some(s) }
}
/// Best-effort: read the locked `candle-core` version from the workspace
/// `Cargo.lock` (two levels up from this crate). Returns `None` if the
/// lockfile is absent (e.g. some packaging flows) or the entry isn't
/// found.
fn candle_version() -> Option<String> {
let manifest = std::env::var("CARGO_MANIFEST_DIR").ok()?;
let lock = std::path::Path::new(&manifest)
.join("..")
.join("..")
.join("Cargo.lock");
println!("cargo:rerun-if-changed={}", lock.display());
let text = std::fs::read_to_string(lock).ok()?;
// Cargo.lock entries are `[[package]]\nname = "x"\nversion = "y"`.
let mut in_candle = false;
for line in text.lines() {
let line = line.trim();
if line == "[[package]]" {
in_candle = false;
} else if line == "name = \"candle-core\"" {
in_candle = true;
} else if in_candle && let Some(rest) = line.strip_prefix("version = \"") {
return Some(rest.trim_end_matches('"').to_string());
}
}
None
}

View File

@@ -41,6 +41,7 @@ pub struct NeuronState {
/// Build the neuron API router.
pub fn neuron_routes() -> Router<Arc<NeuronState>> {
Router::new()
.route("/version", get(version_handler))
.route("/discovery", get(discovery_handler))
.route("/health", get(health_handler))
.route("/models", get(list_models))
@@ -51,6 +52,14 @@ pub fn neuron_routes() -> Router<Arc<NeuronState>> {
.route("/v1/responses", post(responses))
}
/// `GET /version` — the daemon's own build identity (git SHA, enabled
/// features, rustc/candle versions). Static for the process lifetime, so
/// no state is touched. This is the canonical "which build is live"
/// probe for fleet validation and benchmark attribution.
async fn version_handler() -> Json<cortex_core::build_info::BuildInfo> {
Json(crate::version::build_info())
}
async fn discovery_handler(State(state): State<Arc<NeuronState>>) -> Json<DiscoveryResponse> {
Json(state.discovery.clone())
}
@@ -81,6 +90,21 @@ async fn load_model(
State(state): State<Arc<NeuronState>>,
Json(spec): Json<ModelSpec>,
) -> impl IntoResponse {
// Driver/library mismatch preflight (#19): every CUDA load is
// guaranteed to fail until the host reboots. Reject up front with
// the operator-actionable reason instead of letting the load die
// minutes later inside cuInit/NCCL with a cryptic error.
if let Some(reason) = &state.discovery.cuda_unavailable_reason {
tracing::warn!(model = %spec.model_id, reason = %reason, "load_model rejected: CUDA unavailable");
return (
StatusCode::SERVICE_UNAVAILABLE,
Json(json!({
"error": reason,
"code": "cuda_unavailable",
})),
)
.into_response();
}
let registry = state.registry.read().await;
match registry.load_model(&spec).await {
Ok(()) => Json(json!({"status": "loaded"})).into_response(),

View File

@@ -72,6 +72,51 @@ pub struct CandleHarnessConfig {
/// cache_dir. This keeps single-source configs ergonomic.
#[serde(default)]
pub sources: HashMap<String, SourceConfig>,
/// Prefix KV cache across requests (#11). Applies per loaded
/// model, on architectures that support cache snapshots (qwen3_5).
#[serde(default)]
pub prefix_cache: PrefixCacheConfig,
}
/// `[harness.candle.prefix_cache]` settings.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PrefixCacheConfig {
/// Master switch. On by default — set `false` to restore the
/// clear-every-request behaviour.
#[serde(default = "default_prefix_cache_enabled")]
pub enabled: bool,
/// Snapshot byte budget per loaded model, in MiB. Snapshots live
/// on the model's device, so this comes out of the same VRAM that
/// serves inference — size it against the device's headroom after
/// the model weights.
#[serde(default = "default_prefix_cache_budget_mb")]
pub budget_mb: u64,
/// Maximum live snapshots per loaded model, regardless of budget.
#[serde(default = "default_prefix_cache_max_entries")]
pub max_entries: usize,
}
impl Default for PrefixCacheConfig {
fn default() -> Self {
Self {
enabled: default_prefix_cache_enabled(),
budget_mb: default_prefix_cache_budget_mb(),
max_entries: default_prefix_cache_max_entries(),
}
}
}
fn default_prefix_cache_enabled() -> bool {
true
}
fn default_prefix_cache_budget_mb() -> u64 {
1024
}
fn default_prefix_cache_max_entries() -> usize {
8
}
/// Per-scheme source configuration. Mirrors the shape `hf_hub::ApiBuilder`

View File

@@ -100,6 +100,87 @@ pub fn parse_health_info(csv_output: &str) -> Result<Vec<DeviceHealth>> {
Ok(devices)
}
// ── Driver/library mismatch preflight (#19) ─────────────────────────
/// Classify a failed nvidia-smi invocation: is it the classic
/// "Driver/library version mismatch" (userspace libs updated, kernel
/// module not reloaded — every CUDA call on the host is dead until a
/// reboot)? Returns the userspace NVML library version when the
/// message carries one ("NVML library version: 580.159"), or
/// `Some("unknown")` for a mismatch without a parsable version.
/// `None` for any other failure — other errors (no devices, perms)
/// are NOT the mismatch and must not trigger the loud diagnosis.
pub fn classify_driver_mismatch(combined_output: &str) -> Option<String> {
if !combined_output.contains("Driver/library version mismatch") {
return None;
}
let userspace = combined_output
.lines()
.find_map(|l| l.trim().strip_prefix("NVML library version:"))
.map(|v| v.trim().to_string())
.filter(|v| !v.is_empty())
.unwrap_or_else(|| "unknown".to_string());
Some(userspace)
}
/// Extract the loaded kernel module's driver version from
/// `/proc/driver/nvidia/version` contents. Typical first line:
///
/// ```text
/// NVRM version: NVIDIA UNIX Open Kernel Module for x86_64 580.159.03 Release Build (...)
/// ```
pub fn parse_kernel_module_version(proc_contents: &str) -> Option<String> {
let is_numeric = |p: &str| !p.is_empty() && p.chars().all(|c| c.is_ascii_digit());
let line = proc_contents
.lines()
.find(|l| l.starts_with("NVRM version:"))?;
line.split_whitespace()
.find(|tok| {
let mut parts = tok.split('.');
parts.next().is_some_and(is_numeric) && parts.next().is_some_and(is_numeric)
})
.map(|s| s.to_string())
}
/// Render the operator-actionable mismatch description carried in
/// `DiscoveryResponse::cuda_unavailable_reason` and logged at startup.
pub fn mismatch_reason(userspace: &str, kernel_module: Option<&str>) -> String {
format!(
"host NVIDIA driver/library mismatch (userspace NVML {userspace} vs loaded kernel \
module {}) — reboot the host to reload the kernel module; all CUDA inference is \
unavailable until then",
kernel_module.unwrap_or("unknown")
)
}
/// Outcome of an nvidia-smi invocation, distinguishing "binary not
/// present" (CPU-only host, not an error) from "present but failing"
/// (possible driver mismatch — worth classifying).
enum SmiOutcome {
Ok(String),
Failed(String),
Absent,
}
async fn run_nvidia_smi(args: &[&str]) -> SmiOutcome {
match tokio::process::Command::new("nvidia-smi")
.args(args)
.output()
.await
{
Err(_) => SmiOutcome::Absent,
Ok(out) if out.status.success() => {
SmiOutcome::Ok(String::from_utf8_lossy(&out.stdout).to_string())
}
Ok(out) => {
let mut combined = String::from_utf8_lossy(&out.stdout).to_string();
combined.push('\n');
combined.push_str(&String::from_utf8_lossy(&out.stderr));
SmiOutcome::Failed(combined)
}
}
}
// ── Command execution wrappers ──────────────────────────────────────
async fn run_command(cmd: &str, args: &[&str]) -> Result<String> {
@@ -139,23 +220,42 @@ pub async fn discover_system() -> Result<DiscoveryResponse> {
.trim()
.to_string();
let (devices, driver_version) = match run_command_optional(
"nvidia-smi",
&[
let (devices, driver_version, cuda_unavailable_reason) = match run_nvidia_smi(&[
&format!("--query-gpu={NVIDIA_SMI_DISCOVERY_QUERY}"),
"--format=csv,noheader,nounits",
],
)
])
.await
{
Some(output) => {
SmiOutcome::Ok(output) => {
let devs = parse_gpu_info(&output).unwrap_or_default();
let driver = parse_driver_version(&output);
(devs, driver)
(devs, driver, None)
}
None => {
SmiOutcome::Absent => {
tracing::info!("nvidia-smi not found — no GPU devices discovered");
(vec![], None)
(vec![], None, None)
}
SmiOutcome::Failed(combined) => {
// nvidia-smi exists but can't talk to the driver. The case
// worth diagnosing precisely is the userspace↔kernel-module
// version skew after an un-rebooted driver update (#19) —
// every CUDA call on the host fails until a reboot, and
// without this classification it surfaces as a cryptic
// NCCL/cuInit error deep inside the first model load.
let reason = classify_driver_mismatch(&combined).map(|userspace| {
let kmod = std::fs::read_to_string("/proc/driver/nvidia/version")
.ok()
.as_deref()
.and_then(parse_kernel_module_version);
mismatch_reason(&userspace, kmod.as_deref())
});
if reason.is_none() {
tracing::warn!(
output = %combined.trim(),
"nvidia-smi present but failing — no GPU devices discovered"
);
}
(vec![], None, reason)
}
};
@@ -172,6 +272,7 @@ pub async fn discover_system() -> Result<DiscoveryResponse> {
driver_version,
devices,
harnesses: vec![], // populated by harness registry in Phase 8
cuda_unavailable_reason,
})
}
@@ -272,4 +373,63 @@ mod tests {
assert_eq!(health[1].vram_used_mb, 4096);
assert_eq!(health[1].temp_c, 58);
}
// ── #19 driver/library mismatch preflight ────────────────────────
#[test]
fn classify_driver_mismatch_detects_and_extracts_nvml_version() {
// Verbatim shape of nvidia-smi's failure output on a host
// whose userspace libs were updated without a reboot.
let out = "Failed to initialize NVML: Driver/library version mismatch\n\
NVML library version: 580.159\n";
assert_eq!(classify_driver_mismatch(out).as_deref(), Some("580.159"));
}
#[test]
fn classify_driver_mismatch_without_version_line() {
let out = "Failed to initialize NVML: Driver/library version mismatch\n";
assert_eq!(classify_driver_mismatch(out).as_deref(), Some("unknown"));
}
#[test]
fn classify_driver_mismatch_ignores_other_failures() {
// Other nvidia-smi failures must NOT be diagnosed as the
// mismatch (no false positives on healthy or odd hosts).
for out in [
"No devices were found\n",
"Failed to initialize NVML: Insufficient Permissions\n",
"NVIDIA-SMI has failed because it couldn't communicate with the NVIDIA driver.\n",
"",
] {
assert_eq!(
classify_driver_mismatch(out),
None,
"false positive on: {out:?}"
);
}
}
#[test]
fn parse_kernel_module_version_from_proc() {
let proc = "NVRM version: NVIDIA UNIX Open Kernel Module for x86_64 580.159.03 Release Build (dvs-builder@U22-I3-AE24-12-2) Tue May 12 21:03:35 UTC 2026\n\
GCC version: gcc version 15.2.1 20251022 (Red Hat 15.2.1-3) (GCC)\n";
assert_eq!(
parse_kernel_module_version(proc).as_deref(),
Some("580.159.03")
);
}
#[test]
fn parse_kernel_module_version_absent() {
assert_eq!(parse_kernel_module_version(""), None);
assert_eq!(parse_kernel_module_version("GCC version: gcc 15\n"), None);
}
#[test]
fn mismatch_reason_is_operator_actionable() {
let reason = mismatch_reason("580.159", Some("580.159.03"));
assert!(reason.contains("580.159"));
assert!(reason.contains("580.159.03"));
assert!(reason.contains("reboot"));
}
}

View File

@@ -24,6 +24,7 @@ use super::linear_attn::GatedDeltaNet;
use super::mlp::Qwen3_5MLP;
use super::rmsnorm::Qwen3_5RmsNorm;
use super::rope::RotaryEmbedding;
use super::snapshot::LayerKvSnapshot;
/// One of the two attention flavours sitting in a decoder layer's
/// attention slot. Full-attention layers need the rotary table and
@@ -93,12 +94,13 @@ impl Qwen3_5DecoderLayer {
&mut self,
x: &Tensor,
attn_mask: Option<&Tensor>,
offset: usize,
cos: &Tensor,
sin: &Tensor,
) -> candle_core::Result<Tensor> {
let h = self.input_layernorm.forward(x)?;
let attn_out = match &mut self.attention {
AttentionKind::Full(attn) => attn.forward(&h, attn_mask, offset)?,
// Linear attention ignores attn_mask + offset; its causal
AttentionKind::Full(attn) => attn.forward(&h, attn_mask, cos, sin)?,
// Linear attention ignores attn_mask + rope; its causal
// structure is baked into the recurrent state lifecycle.
AttentionKind::Linear(net) => net.forward(&h)?,
};
@@ -114,4 +116,37 @@ impl Qwen3_5DecoderLayer {
AttentionKind::Linear(net) => net.clear_kv_cache(),
}
}
/// Capture this layer's cache state for a prefix snapshot.
pub fn snapshot_kv(&self) -> candle_core::Result<LayerKvSnapshot> {
Ok(match &self.attention {
AttentionKind::Full(attn) => LayerKvSnapshot::Full(attn.snapshot_kv()),
AttentionKind::Linear(net) => {
let (conv_state, recurrent_state) = net.snapshot_state()?;
LayerKvSnapshot::Linear {
conv_state,
recurrent_state,
}
}
})
}
/// Replace this layer's cache state from a snapshot. The snapshot
/// variant must match the layer's attention kind — a mismatch
/// means the snapshot came from a different model.
pub fn restore_kv(&mut self, snap: &LayerKvSnapshot) -> candle_core::Result<()> {
match (&mut self.attention, snap) {
(AttentionKind::Full(attn), LayerKvSnapshot::Full(kv)) => attn.restore_kv(kv.as_ref()),
(
AttentionKind::Linear(net),
LayerKvSnapshot::Linear {
conv_state,
recurrent_state,
},
) => net.restore_state(conv_state.as_ref(), recurrent_state.as_ref()),
_ => candle_core::bail!(
"restore_kv: snapshot layer kind does not match this layer's attention kind"
),
}
}
}

View File

@@ -96,7 +96,8 @@ impl Qwen3_5Attention {
&mut self,
x: &Tensor,
attn_mask: Option<&Tensor>,
offset: usize,
cos: &Tensor,
sin: &Tensor,
) -> candle_core::Result<Tensor> {
let (b, l, _) = x.dims3()?;
@@ -131,8 +132,9 @@ impl Qwen3_5Attention {
.transpose(1, 2)?
.contiguous()?;
// 3. RoPE on q, k.
let (q, k) = self.rotary.apply(&q, &k, offset)?;
// 3. RoPE on q, k (cos/sin built once per forward by the model —
// interleaved M-RoPE for image tokens, plain for text).
let (q, k) = self.rotary.apply_cos_sin(&q, &k, cos, sin)?;
// 4. KV cache.
let (k, v) = self.kv_cache.append(&k, &v)?;
@@ -163,6 +165,26 @@ impl Qwen3_5Attention {
pub fn clear_kv_cache(&mut self) {
self.kv_cache.reset();
}
/// Capture the KV cache contents for a prefix snapshot. Shallow
/// clones: `ConcatKvCache::append` cats into fresh allocations and
/// never mutates stored tensors in place, so the captured tensors
/// stay valid after the live cache moves on.
pub fn snapshot_kv(&self) -> Option<(Tensor, Tensor)> {
match (self.kv_cache.k(), self.kv_cache.v()) {
(Some(k), Some(v)) => Some((k.clone(), v.clone())),
_ => None,
}
}
/// Replace the live KV cache with a previously captured snapshot.
pub fn restore_kv(&mut self, snap: Option<&(Tensor, Tensor)>) -> candle_core::Result<()> {
self.kv_cache.reset();
if let Some((k, v)) = snap {
self.kv_cache.append(k, v)?;
}
Ok(())
}
}
fn load_linear_no_bias(

View File

@@ -49,11 +49,15 @@
//!
//! ## Performance note
//!
//! This impl is the **recurrent** delta-rule for both prefill and
//! decode — i.e. the algorithm in `torch_recurrent_gated_delta_rule`.
//! Correctness-first. The chunked algorithm (chunk_size=64) in
//! `torch_chunk_gated_delta_rule` is a perf optimisation for long
//! prefill; can be added later without changing the surface.
//! Prefill (seq_len ≥ 64) runs the **chunked** delta rule (#23) — the
//! algorithm in `torch_chunk_gated_delta_rule`, reorganised into
//! per-chunk batched matmuls; see [`run_chunk_gated_delta_rule`].
//! Decode steps and short prompts keep the **recurrent** per-token
//! rule (`torch_recurrent_gated_delta_rule`): a CUDA kernel on
//! device, a pure-Rust loop on CPU. Both produce identical results
//! (pinned by the `chunked_matches_recurrent_*` parity tests);
//! `NEURON_GDN_CHUNKED=0` forces the recurrent paths for A/B
//! measurement.
use anyhow::{Context, Result};
use candle_core::{Module, Tensor};
@@ -184,6 +188,42 @@ impl GatedDeltaNet {
self.state = GatedDeltaNetState::default();
}
/// Deep-copy the recurrent state for a prefix snapshot. Must be a
/// real copy (`Tensor::copy`), not a refcount clone: the CUDA
/// delta-rule kernels write the state buffer in place, so a
/// shared-storage snapshot would be corrupted by the next forward.
pub fn snapshot_state(&self) -> candle_core::Result<(Option<Tensor>, Option<Tensor>)> {
let conv = self
.state
.conv_state
.as_ref()
.map(Tensor::copy)
.transpose()?;
let rec = self
.state
.recurrent_state
.as_ref()
.map(Tensor::copy)
.transpose()?;
Ok((conv, rec))
}
/// Replace the live recurrent state with a deep copy of a
/// previously captured snapshot. Deep copy for the same in-place
/// kernel reason as [`Self::snapshot_state`] — the snapshot must
/// survive being restored more than once.
pub fn restore_state(
&mut self,
conv_state: Option<&Tensor>,
recurrent_state: Option<&Tensor>,
) -> candle_core::Result<()> {
self.state = GatedDeltaNetState {
conv_state: conv_state.map(Tensor::copy).transpose()?,
recurrent_state: recurrent_state.map(Tensor::copy).transpose()?,
};
Ok(())
}
/// `x` shape: `(B, L, hidden_size)`. Returns the same shape.
pub fn forward(&mut self, x: &Tensor) -> candle_core::Result<Tensor> {
let (batch_size, seq_len, _) = x.dims3()?;
@@ -357,6 +397,16 @@ pub(crate) fn run_delta_rule(
head_k_dim: usize,
head_v_dim: usize,
) -> candle_core::Result<(Tensor, Tensor)> {
// Prefill takes the chunk-parallel algorithm (#23): identical
// delta-rule math reorganised into per-chunk matmuls (cuBLAS /
// tensor cores on CUDA, gemm on CPU) instead of an O(L)-sequential
// per-token recurrence. Decode steps (seq_len 1) and short
// prompts stay on the recurrent paths below. The env kill switch
// exists for A/B measurement on the fleet.
const CHUNK_ALGO_THRESHOLD: usize = 64;
if seq_len >= CHUNK_ALGO_THRESHOLD && chunked_prefill_enabled() {
return run_chunk_gated_delta_rule(q, k, v, g, beta, state);
}
#[cfg(feature = "cuda")]
{
// Only dispatch to the kernel if the inputs are on a CUDA
@@ -371,6 +421,198 @@ pub(crate) fn run_delta_rule(
run_delta_rule_rust(q, k, v, g, beta, state, seq_len)
}
/// `NEURON_GDN_CHUNKED=0` falls back to the per-token recurrent
/// paths for prefill — kept for A/B measurement on live hosts.
fn chunked_prefill_enabled() -> bool {
static ENABLED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
*ENABLED.get_or_init(|| {
std::env::var("NEURON_GDN_CHUNKED")
.map(|v| v != "0" && !v.eq_ignore_ascii_case("false"))
.unwrap_or(true)
})
}
/// Chunk-parallel gated delta rule — a faithful port of the HF
/// reference `torch_chunk_gated_delta_rule` (chunk_size = 64) in
/// `transformers/models/qwen3_5/modeling_qwen3_5.py`, minus the steps
/// our caller has already done (q/k L2-norm, q pre-scaled by
/// `1/sqrt(D_k)`, inputs already `(B, H, L, D)` f32).
///
/// Same inputs/outputs as [`run_delta_rule`]'s recurrent paths:
/// `q`/`k`: `(B, H, L, D_k)`, `v`: `(B, H, L, D_v)`, `g`/`beta`:
/// `(B, H, L)`, `state`: `(B, H, D_k, D_v)` (zeros or a restored
/// prefix snapshot's recurrent state). Returns
/// `(out: (B, H, L, D_v), state: (B, H, D_k, D_v))`, all f32.
///
/// The reference's in-place UT-transform row loop is kept as-is
/// (with rows accumulating into a fresh tensor — candle tensors are
/// immutable); see the numerical-caution note at the loop for why the
/// tempting nilpotent-squaring shortcut is wrong. The parity tests
/// pin this against the recurrent path.
pub(crate) fn run_chunk_gated_delta_rule(
q: &Tensor,
k: &Tensor,
v: &Tensor,
g: &Tensor,
beta: &Tensor,
state: Tensor,
) -> candle_core::Result<(Tensor, Tensor)> {
const C: usize = 64;
let (b, h, l, dk) = q.dims4()?;
let dv = v.dim(3)?;
let device = q.device().clone();
// Pad L up to a multiple of the chunk size. Padded positions
// carry beta = 0 (no state update) and g = 0 (no decay), so they
// are inert in the recurrence; their outputs are sliced off at
// the end.
let pad = (C - l % C) % C;
let (q, k, v, g, beta) = if pad > 0 {
(
q.pad_with_zeros(2, 0, pad)?,
k.pad_with_zeros(2, 0, pad)?,
v.pad_with_zeros(2, 0, pad)?,
g.pad_with_zeros(2, 0, pad)?,
beta.pad_with_zeros(2, 0, pad)?,
)
} else {
(q.clone(), k.clone(), v.clone(), g.clone(), beta.clone())
};
let lt = l + pad;
let n = lt / C;
let beta_e = beta.unsqueeze(3)?; // (B, H, Lt, 1)
let v_beta = v.broadcast_mul(&beta_e)?;
let k_beta = k.broadcast_mul(&beta_e)?;
// Chunk reshape, flattening (B, H, N) into one batch dim — candle's
// matmul supports at most two batch dims, so the chunk-local math
// runs rank-3 over B·H·N and reshapes back to rank-5 for the
// inter-chunk loop's per-chunk narrows.
let bhn = b * h * n;
let q3 = q.reshape((bhn, C, dk))?;
let k3 = k.reshape((bhn, C, dk))?;
let k_beta3 = k_beta.reshape((bhn, C, dk))?;
let v_beta3 = v_beta.reshape((bhn, C, dv))?;
// Within-chunk cumulative log-decay.
let g3 = g.reshape((bhn, C))?.cumsum(1)?;
// Lower-triangular masks, broadcast over the batch dim.
let tril_incl = {
let mut m = vec![0f32; C * C];
for i in 0..C {
for j in 0..=i {
m[i * C + j] = 1.0;
}
}
Tensor::from_vec(m, (C, C), &device)?
};
let tril_strict = {
let mut m = vec![0f32; C * C];
for i in 0..C {
for j in 0..i {
m[i * C + j] = 1.0;
}
}
Tensor::from_vec(m, (C, C), &device)?
};
// decay_mask[i][j] = exp(g_i - g_j) on the lower triangle
// (diagonal = 1), zero above. Mask-multiply replaces the
// reference's tril/exp/tril dance: upper entries become
// exp(0) = 1 mid-way and are re-zeroed.
let g_col = g3.unsqueeze(2)?; // (BHN, C, 1)
let g_row = g3.unsqueeze(1)?; // (BHN, 1, C)
let decay_mask3 = g_col
.broadcast_sub(&g_row)?
.broadcast_mul(&tril_incl)?
.exp()?
.broadcast_mul(&tril_incl)?
.contiguous()?;
// T = strict lower of -((k_beta k^T) ⊙ decay), then
// M = (I - T)^{-1} by forward substitution over rows — the
// reference's in-place UT-transform loop, with processed rows
// accumulating in `done` instead of mutating in place.
//
// Numerical caution: T is nilpotent (T^64 = 0), so the inverse
// also equals Π (I + T^(2^j)) — six matmuls — but that form is
// numerically unsafe: raw powers of T grow combinatorially
// (path counts up to C(62,31) ≈ 4.6e17) before nilpotency
// collapses them, destroying f32 precision on real prompts with
// correlated keys. The forward substitution's intermediates are
// the convergent M entries themselves, matching the reference's
// behaviour exactly. Pinned by `chunked_ut_transform_survives_
// correlated_keys`.
let kkt = k_beta3.matmul(&k3.transpose(1, 2)?.contiguous()?)?;
let t = kkt
.broadcast_mul(&decay_mask3)?
.broadcast_mul(&tril_strict)?
.neg()?
.contiguous()?;
let eye = Tensor::eye(C, candle_core::DType::F32, &device)?;
// Row 0 of the strict-lower T is all zeros and passes through
// unchanged, seeding the processed-rows accumulator.
let mut done = t.narrow(1, 0, 1)?.contiguous()?;
for i in 1..C {
let row = t.narrow(1, i, 1)?; // (BHN, 1, C)
let coeffs = row.narrow(2, 0, i)?.contiguous()?; // (BHN, 1, i)
let updated = (&row + coeffs.matmul(&done)?)?; // (BHN, 1, C)
done = Tensor::cat(&[&done, &updated], 1)?;
}
let m = done.broadcast_add(&eye)?.contiguous()?;
// value' = M v_beta ; k_cumdecay = M (k_beta ⊙ exp(g)).
let value_c3 = m.matmul(&v_beta3.contiguous()?)?;
let g_exp3 = g3.exp()?.unsqueeze(2)?; // (BHN, C, 1)
let k_cumdecay3 = m.matmul(&k_beta3.broadcast_mul(&g_exp3)?.contiguous()?)?;
// Rank-5 views for the per-chunk narrows below.
let q = q3.reshape((b, h, n, C, dk))?;
let k = k3.reshape((b, h, n, C, dk))?;
let value_c = value_c3.reshape((b, h, n, C, dv))?;
let k_cumdecay = k_cumdecay3.reshape((b, h, n, C, dk))?;
let decay_mask = decay_mask3.reshape((b, h, n, C, C))?;
let g = g3.reshape((b, h, n, C))?;
// Inter-chunk recurrence: a handful of matmuls per 64 tokens.
let mut state = state.to_dtype(candle_core::DType::F32)?;
let mut outs: Vec<Tensor> = Vec::with_capacity(n);
for i in 0..n {
let q_i = q.narrow(2, i, 1)?.squeeze(2)?.contiguous()?; // (B, H, C, Dk)
let k_i = k.narrow(2, i, 1)?.squeeze(2)?.contiguous()?;
let v_i = value_c.narrow(2, i, 1)?.squeeze(2)?.contiguous()?; // (B, H, C, Dv)
let dm_i = decay_mask.narrow(2, i, 1)?.squeeze(2)?; // (B, H, C, C)
let g_i = g.narrow(2, i, 1)?.squeeze(2)?; // (B, H, C)
let kcd_i = k_cumdecay.narrow(2, i, 1)?.squeeze(2)?.contiguous()?;
let attn = q_i
.matmul(&k_i.transpose(2, 3)?.contiguous()?)?
.broadcast_mul(&dm_i)?
.contiguous()?;
let v_prime = kcd_i.matmul(&state)?;
let v_new = (v_i - v_prime)?.contiguous()?;
let g_i_exp = g_i.exp()?.unsqueeze(3)?; // (B, H, C, 1)
let attn_inter = q_i.broadcast_mul(&g_i_exp)?.contiguous()?.matmul(&state)?;
let out_i = (attn_inter + attn.matmul(&v_new)?)?;
outs.push(out_i.unsqueeze(2)?);
// state ← state · exp(g_last) + (k_i ⊙ exp(g_last - g_i))^T v_new
let g_last = g_i.narrow(2, C - 1, 1)?; // (B, H, 1)
let carry = g_last.exp()?.unsqueeze(3)?; // (B, H, 1, 1)
let w = k_i.broadcast_mul(&g_last.broadcast_sub(&g_i)?.exp()?.unsqueeze(3)?)?;
state =
(state.broadcast_mul(&carry)? + w.transpose(2, 3)?.contiguous()?.matmul(&v_new)?)?;
}
let out = Tensor::cat(&outs, 2)?
.reshape((b, h, lt, dv))?
.narrow(2, 0, l)?
.contiguous()?;
Ok((out, state))
}
/// CUDA path. Flattens (B, H, ...) → (BH, ...) at the kernel boundary
/// (the kernel uses BH = batch*heads as its outer batch axis) and
/// reshapes the kernel's outputs back to (B, H, ...) for the caller.
@@ -687,6 +929,151 @@ mod tests {
use super::*;
use candle_core::{DType, Device};
/// Plausible delta-rule inputs matching `run_delta_rule`'s
/// contract: q/k L2-normed (q pre-scaled by 1/sqrt(D_k)), g a
/// negative log-decay, beta in (0, 1). All f32 on CPU.
fn delta_rule_inputs(
b: usize,
h: usize,
l: usize,
dk: usize,
dv: usize,
) -> (Tensor, Tensor, Tensor, Tensor, Tensor) {
let dev = Device::Cpu;
let scale = 1.0 / (dk as f64).sqrt();
let q = Tensor::randn(0f32, 1.0, (b, h, l, dk), &dev).unwrap();
let q = (l2norm(&q, 1e-6).unwrap() * scale).unwrap();
let k = Tensor::randn(0f32, 1.0, (b, h, l, dk), &dev).unwrap();
let k = l2norm(&k, 1e-6).unwrap();
let v = (Tensor::randn(0f32, 1.0, (b, h, l, dv), &dev).unwrap() * 0.5).unwrap();
// g in (-1, 0): a realistic per-token log-decay.
let g = (Tensor::rand(0f32, 1f32, (b, h, l), &dev).unwrap() * -1.0).unwrap();
let beta = Tensor::rand(0.05f32, 0.95f32, (b, h, l), &dev).unwrap();
(q, k, v, g, beta)
}
fn max_abs_diff(a: &Tensor, b: &Tensor) -> f32 {
(a - b)
.unwrap()
.abs()
.unwrap()
.flatten_all()
.unwrap()
.max(0)
.unwrap()
.to_scalar::<f32>()
.unwrap()
}
/// The #23 parity gate: the chunk-parallel algorithm must produce
/// the same outputs and final state as the per-token recurrence.
/// L = 130 exercises the pad-to-chunk-multiple path (130 = 2×64 + 2).
#[test]
fn chunked_matches_recurrent_with_padding() {
let (b, h, l, dk, dv) = (1, 2, 130, 16, 16);
let (q, k, v, g, beta) = delta_rule_inputs(b, h, l, dk, dv);
let zeros = || Tensor::zeros((b, h, dk, dv), DType::F32, &Device::Cpu).unwrap();
let (out_rec, state_rec) = run_delta_rule_rust(&q, &k, &v, &g, &beta, zeros(), l).unwrap();
let (out_chk, state_chk) =
run_chunk_gated_delta_rule(&q, &k, &v, &g, &beta, zeros()).unwrap();
assert_eq!(out_chk.dims(), out_rec.dims());
let d_out = max_abs_diff(&out_rec, &out_chk);
let d_state = max_abs_diff(&state_rec, &state_chk);
assert!(d_out < 2e-4, "output diverged: {d_out}");
assert!(d_state < 2e-4, "final state diverged: {d_state}");
}
/// Exact chunk multiple (no padding) continuing from a non-zero
/// initial state — the prefix-cache-restore (#11) interaction.
#[test]
fn chunked_matches_recurrent_with_initial_state() {
let (b, h, dk, dv) = (1, 2, 16, 16);
let dev = Device::Cpu;
// Build a non-trivial initial state by running the recurrent
// path over a 50-token "restored prefix".
let (pq, pk, pv, pg, pbeta) = delta_rule_inputs(b, h, 50, dk, dv);
let zeros = Tensor::zeros((b, h, dk, dv), DType::F32, &dev).unwrap();
let (_, state0) = run_delta_rule_rust(&pq, &pk, &pv, &pg, &pbeta, zeros, 50).unwrap();
let l = 128;
let (q, k, v, g, beta) = delta_rule_inputs(b, h, l, dk, dv);
let (out_rec, state_rec) =
run_delta_rule_rust(&q, &k, &v, &g, &beta, state0.clone(), l).unwrap();
let (out_chk, state_chk) =
run_chunk_gated_delta_rule(&q, &k, &v, &g, &beta, state0).unwrap();
let d_out = max_abs_diff(&out_rec, &out_chk);
let d_state = max_abs_diff(&state_rec, &state_chk);
assert!(d_out < 2e-4, "output diverged: {d_out}");
assert!(d_state < 2e-4, "final state diverged: {d_state}");
}
/// Adversarially correlated inputs: near-identical keys with
/// beta ≈ 1 and negligible decay make the UT-transform matrix T
/// maximally coherent — raw powers of T grow combinatorially
/// (≈ C(62,31) paths), which destroyed f32 precision in the
/// nilpotent-squaring formulation this test exists to forbid.
/// Real prompts hit this through repetitive text (observed live
/// on beast: NaN logits → "!!!" replies). Forward substitution
/// must stay finite and match the recurrent path.
#[test]
fn chunked_ut_transform_survives_correlated_keys() {
let (b, h, l, dk, dv) = (1, 1, 192, 16, 16);
let dev = Device::Cpu;
let scale = 1.0 / (dk as f64).sqrt();
// One base direction plus a whisper of noise: every key is
// nearly the same unit vector.
let base = Tensor::randn(0f32, 1.0, (1, 1, 1, dk), &dev).unwrap();
let noise = (Tensor::randn(0f32, 1.0, (b, h, l, dk), &dev).unwrap() * 0.01).unwrap();
let k = l2norm(&base.broadcast_add(&noise).unwrap(), 1e-6).unwrap();
let q = (l2norm(&base.broadcast_add(&noise).unwrap(), 1e-6).unwrap() * scale).unwrap();
let v = (Tensor::randn(0f32, 1.0, (b, h, l, dv), &dev).unwrap() * 0.5).unwrap();
// Almost no decay, near-unit update rate — worst case for T.
let g = (Tensor::rand(0f32, 1f32, (b, h, l), &dev).unwrap() * -1e-3).unwrap();
let beta = Tensor::rand(0.98f32, 0.999f32, (b, h, l), &dev).unwrap();
let zeros = || Tensor::zeros((b, h, dk, dv), DType::F32, &dev).unwrap();
let (out_rec, state_rec) = run_delta_rule_rust(&q, &k, &v, &g, &beta, zeros(), l).unwrap();
let (out_chk, state_chk) =
run_chunk_gated_delta_rule(&q, &k, &v, &g, &beta, zeros()).unwrap();
let finite: Vec<f32> = out_chk.flatten_all().unwrap().to_vec1().unwrap();
assert!(
finite.iter().all(|x| x.is_finite()),
"chunked output not finite on correlated inputs"
);
let d_out = max_abs_diff(&out_rec, &out_chk);
let d_state = max_abs_diff(&state_rec, &state_chk);
assert!(
d_out < 5e-3,
"output diverged on correlated inputs: {d_out}"
);
assert!(
d_state < 5e-3,
"final state diverged on correlated inputs: {d_state}"
);
}
/// A single exact chunk — the smallest input the dispatch sends to
/// the chunked path.
#[test]
fn chunked_matches_recurrent_single_chunk() {
let (b, h, l, dk, dv) = (2, 3, 64, 8, 8);
let (q, k, v, g, beta) = delta_rule_inputs(b, h, l, dk, dv);
let zeros = || Tensor::zeros((b, h, dk, dv), DType::F32, &Device::Cpu).unwrap();
let (out_rec, state_rec) = run_delta_rule_rust(&q, &k, &v, &g, &beta, zeros(), l).unwrap();
let (out_chk, state_chk) =
run_chunk_gated_delta_rule(&q, &k, &v, &g, &beta, zeros()).unwrap();
let d_out = max_abs_diff(&out_rec, &out_chk);
let d_state = max_abs_diff(&state_rec, &state_chk);
assert!(d_out < 2e-4, "output diverged: {d_out}");
assert!(d_state < 2e-4, "final state diverged: {d_state}");
}
#[test]
fn softplus_small_x() {
// softplus(0) = ln(2) ≈ 0.6931
@@ -737,6 +1124,8 @@ mod tests {
rope_theta: 10000.0,
partial_rotary_factor: 1.0,
rope_type: None,
mrope_section: Vec::new(),
mrope_interleaved: false,
},
rms_norm_eps: 1e-6,
tie_word_embeddings: false,

View File

@@ -78,6 +78,7 @@ pub mod linear_attn;
pub mod mlp;
pub mod rmsnorm;
pub mod rope;
pub mod snapshot;
pub mod vision;
use decoder::Qwen3_5DecoderLayer;
@@ -191,11 +192,12 @@ fn default_hidden_act() -> String {
}
/// Nested `rope_parameters` block from a Qwen3-Next `config.json`.
/// `mrope_section` and `mrope_interleaved` are accepted via the
/// `#[serde(default)]` flatten-tolerance below but ignored — we treat
/// MRoPE as plain RoPE for text-only inference (the three position
/// grids carry identical ids when there's no vision input, so the
/// interleaving is a no-op).
///
/// For text-only inference the three MRoPE position grids carry
/// identical ids, so the interleave is a no-op and plain RoPE applies.
/// For vision inputs `mrope_section` + `mrope_interleaved` drive the
/// per-axis (text/height/width) rotary used by image tokens — see
/// `rope.rs`.
#[derive(Debug, Clone, Deserialize)]
pub struct RopeParameters {
/// Base for the inverse-frequency computation. Qwen3.6: 10_000_000.
@@ -211,6 +213,16 @@ pub struct RopeParameters {
/// implemented here.
#[serde(default)]
pub rope_type: Option<String>,
/// MRoPE per-axis section sizes `[text, height, width]` — e.g.
/// `[11, 11, 10]` for Qwen3.6, summing to the rotary half-dim.
/// Empty for models that don't declare MRoPE (→ plain RoPE).
#[serde(default)]
pub mrope_section: Vec<usize>,
/// Whether the three MRoPE axes are interleaved per-frequency
/// (Qwen3-VL / Qwen3.6 style, `true`) rather than block-concatenated
/// (Qwen2-VL style, `false`).
#[serde(default)]
pub mrope_interleaved: bool,
}
fn default_rope_theta() -> f64 {
@@ -303,6 +315,16 @@ pub struct Qwen3_5Model {
embed_tokens: Embedding,
layers: Vec<Qwen3_5DecoderLayer>,
norm: Qwen3_5RmsNorm,
/// Shared with every full-attention layer; the model uses it to
/// build the per-forward cos/sin (interleaved M-RoPE for image
/// tokens, plain for text) once, which the layers then apply.
rotary: Arc<RotaryEmbedding>,
/// `offset + rope_delta` is the text-axis position during decode.
/// 0 for text-only; set from `get_rope_index` during a vision
/// prefill (image tokens compress the position space, so text after
/// the image resumes from a smaller counter than the sequence
/// index). Reset in `clear_kv_cache`.
rope_delta: i64,
device: Device,
dtype: DType,
}
@@ -354,6 +376,8 @@ impl Qwen3_5Model {
embed_tokens,
layers,
norm,
rotary,
rope_delta: 0,
device,
dtype,
})
@@ -367,6 +391,45 @@ impl Qwen3_5Model {
for l in &mut self.layers {
l.clear_kv_cache();
}
// New request → no image-compressed position offset until the
// next vision prefill sets one.
self.rope_delta = 0;
}
/// Capture every layer's cache state plus the rope position
/// counter as one consistent prefix snapshot (#11). Only valid at
/// a token boundary — i.e. between forward calls, which is the
/// only time the caller can reach this anyway.
pub fn snapshot_kv_cache(&self) -> candle_core::Result<snapshot::KvCacheSnapshot> {
let layers = self
.layers
.iter()
.map(|l| l.snapshot_kv())
.collect::<candle_core::Result<Vec<_>>>()?;
Ok(snapshot::KvCacheSnapshot {
layers,
rope_delta: self.rope_delta,
})
}
/// Replace the live cache state with a previously captured
/// snapshot. The snapshot stays valid for further restores.
pub fn restore_kv_cache(
&mut self,
snap: &snapshot::KvCacheSnapshot,
) -> candle_core::Result<()> {
if snap.layers.len() != self.layers.len() {
candle_core::bail!(
"restore_kv_cache: snapshot has {} layers, model has {}",
snap.layers.len(),
self.layers.len()
);
}
for (layer, layer_snap) in self.layers.iter_mut().zip(snap.layers.iter()) {
layer.restore_kv(layer_snap)?;
}
self.rope_delta = snap.rope_delta;
Ok(())
}
fn causal_mask(&self, b: usize, tgt: usize, offset: usize) -> candle_core::Result<Tensor> {
@@ -378,7 +441,34 @@ impl Qwen3_5Model {
}
pub fn forward(&mut self, input: &Tensor, offset: usize) -> candle_core::Result<Tensor> {
self.forward_inner(input, offset, None, None)
self.forward_inner(input, offset, None, None, &[], None)
}
/// Forward for a vision-prefill chunk: optional image-embedding
/// splice plus explicit interleaved-M-RoPE `position_ids` (the
/// chunk's slice of the full prompt's 3D positions). Mirrors the TP
/// `TpQwen3_5Model::forward_with_positions` — used by
/// `Qwen3_5ForCausalLM::prefill_with_images_chunked`, which computes
/// the positions once over the whole prompt and slices them per
/// chunk so the position counters stay consistent across chunk
/// boundaries (an image compresses the position space, so per-chunk
/// offset arithmetic would be wrong).
pub fn forward_with_positions(
&mut self,
input: &Tensor,
offset: usize,
position_ids: &Tensor,
image_embeds: Option<&Tensor>,
image_token_id: Option<u32>,
) -> candle_core::Result<Tensor> {
self.forward_inner(
input,
offset,
image_embeds,
image_token_id,
&[],
Some(position_ids),
)
}
/// Forward with image-embedding splice. Stage B of the vision plan.
@@ -396,43 +486,49 @@ impl Qwen3_5Model {
///
/// The splice replaces the LM's text-side embedding at each
/// `image_token_id` position with the corresponding row from
/// `image_embeds`. After the splice the decoder runs unchanged.
///
/// **MRoPE gap.** Qwen3.6's `rope_parameters` declares MRoPE
/// (interleaved text/height/width axes); Stage B applies plain
/// text-position RoPE to image tokens. The model still attends
/// to image content but loses spatial structure that MRoPE-aware
/// position encoding would preserve. Tracked under issue #15
/// (numerical validation) — quality benchmark from Stage D should
/// surface the impact, and the fix lives in `rope::RotaryEmbedding`.
/// `image_embeds`. After the splice the decoder runs the interleaved
/// M-RoPE path: `grids` carries each image's post-merge LM grid
/// `(lm_gh, lm_gw)` so `get_rope_index` assigns image tokens their 2D
/// coordinates (dynamic resolution, #14).
pub fn forward_with_vision(
&mut self,
input_ids: &Tensor,
offset: usize,
image_embeds: &Tensor,
image_token_id: u32,
grids: &[(usize, usize)],
) -> candle_core::Result<Tensor> {
self.forward_inner(input_ids, offset, Some(image_embeds), Some(image_token_id))
self.forward_inner(
input_ids,
offset,
Some(image_embeds),
Some(image_token_id),
grids,
None,
)
}
/// Shared forward. Splices image embeddings at `image_token_id`
/// positions when present, then builds the rotary cos/sin, in
/// precedence order: explicit `position_ids` (interleaved M-RoPE,
/// the chunked-vision path that slices a once-computed position
/// tensor) > internal M-RoPE from `grids` (single-shot vision) >
/// plain positions at `offset + rope_delta` (text / decode).
fn forward_inner(
&mut self,
input: &Tensor,
offset: usize,
image_embeds: Option<&Tensor>,
image_token_id: Option<u32>,
grids: &[(usize, usize)],
position_ids: Option<&Tensor>,
) -> candle_core::Result<Tensor> {
let (b, l) = input.dims2()?;
let mut h = self.embed_tokens.forward(input)?;
// Splice image embeddings at `image_token_id` positions. The
// caller pre-expanded the prompt so every patch token in the
// image_embeds tensor has a matching position in `input`. We
// index_put the rows in place.
// Splice image embeddings at `image_token_id` positions, when
// this forward carries any. Independent of how cos/sin is built.
if let (Some(img), Some(tok_id)) = (image_embeds, image_token_id) {
// Locate image-token positions in input_ids. Operate on
// CPU since the input ids are tiny (max ~10k entries
// including the patch expansion) and the comparison is
// not in the per-step hot path.
let ids: Vec<u32> = input.flatten_all()?.to_vec1()?;
let mut positions: Vec<u32> = Vec::with_capacity(img.dim(0)?);
for (idx, id) in ids.iter().enumerate() {
@@ -443,30 +539,43 @@ impl Qwen3_5Model {
let n_img_tokens = img.dim(0)?;
if positions.len() != n_img_tokens {
candle_core::bail!(
"forward_with_vision: prompt has {} image-token positions but \
image_embeds carries {} tokens — call build_prompt_for_request to \
ensure the per-image patch-count expansion has been applied",
"forward_with_vision: chunk has {} image-token positions but \
image_embeds carries {} tokens — per-image patch-count expansion \
/ chunk slicing mismatch",
positions.len(),
n_img_tokens,
);
}
if !positions.is_empty() {
// Cast image_embeds to the LM's dtype so the splice
// produces a uniform tensor for the decoder stack.
// Cast image_embeds to the LM's dtype, then splice the
// contiguous `<|image_pad|>` runs in place.
let img = img.to_dtype(self.dtype)?;
// index_select would return the rows; we want to put.
// candle's slice_assign with explicit positions ranges
// doesn't exist; use scatter via index_select + an
// accumulator: build a `(B, L, hidden)` zero tensor,
// scatter the image rows in, then add to a masked
// version of `h`. Simpler approach: walk positions
// and use `slice_assign` for contiguous runs. Since
// image_pad runs are contiguous (template emits
// `<|vision_start|><|image_pad|>×N<|vision_end|>`),
// we group positions and assign per run.
h = splice_runs(&h, &img, &positions)?;
}
}
// Build interleaved M-RoPE cos/sin so image tokens carry their
// 2D (lm_gh × lm_gw) grid coordinates. Text / decode take the
// plain-RoPE fast path — bit-for-bit the pre-M-RoPE behaviour
// when `rope_delta == 0`.
let (cos, sin) = if let Some(pos) = position_ids {
// Pre-computed positions sliced for this chunk — the splice
// above already advanced `rope_delta`'s effect into `pos`.
self.rotary.mrope_cos_sin(pos)?
} else if let Some(tok_id) = image_token_id {
// Single-shot vision: compute the whole prompt's M-RoPE here
// and stash `rope_delta` for the decode that follows.
let ids: Vec<u32> = input.flatten_all()?.to_vec1()?;
let (text, height, width, delta) = rope::get_rope_index(&ids, tok_id, grids)
.map_err(|e| candle_core::Error::Msg(format!("get_rope_index: {e}")))?;
self.rope_delta = delta;
let pos = rope::mrope_position_tensor(&text, &height, &width, &self.device)?;
self.rotary.mrope_cos_sin(&pos)?
} else {
let base = (offset as i64 + self.rope_delta).max(0) as usize;
self.rotary.plain_cos_sin(base, l)?
};
// Causal mask only needed for L > 1 prefill; full-attention
// layers consume it via broadcast_add. Linear-attention layers
// ignore the mask.
@@ -476,7 +585,7 @@ impl Qwen3_5Model {
Some(self.causal_mask(b, l, offset)?)
};
for layer in &mut self.layers {
h = layer.forward(&h, causal.as_ref(), offset)?;
h = layer.forward(&h, causal.as_ref(), &cos, &sin)?;
}
self.norm.forward(&h)
}
@@ -577,17 +686,182 @@ impl Qwen3_5ForCausalLM {
offset: usize,
image_embeds: &Tensor,
image_token_id: u32,
grids: &[(usize, usize)],
) -> candle_core::Result<Tensor> {
let (_, l) = input.dims2()?;
let hidden = self
.base
.forward_with_vision(input, offset, image_embeds, image_token_id)?;
let hidden =
self.base
.forward_with_vision(input, offset, image_embeds, image_token_id, grids)?;
hidden.i((.., l - 1.., ..))?.apply(&self.lm_head)
}
/// Forward for a vision-prefill chunk: explicit M-RoPE positions +
/// optional image splice. Mirrors `forward_with_vision` but routes
/// through `Qwen3_5Model::forward_with_positions`. Used by
/// [`Self::prefill_with_images_chunked`].
pub fn forward_with_positions(
&mut self,
input: &Tensor,
offset: usize,
position_ids: &Tensor,
image_embeds: Option<&Tensor>,
image_token_id: Option<u32>,
) -> candle_core::Result<Tensor> {
let (_, l) = input.dims2()?;
let hidden = self.base.forward_with_positions(
input,
offset,
position_ids,
image_embeds,
image_token_id,
)?;
hidden.i((.., l - 1.., ..))?.apply(&self.lm_head)
}
/// Encode every preprocessed `(C, H, W)` image once through the
/// vision tower and concatenate along the patch axis →
/// `(sum_patches, hidden)`. Done once per prefill, not per chunk.
fn encode_images_concat(&self, image_pixels: &[Tensor]) -> candle_core::Result<Tensor> {
let tower = self.vision.as_ref().ok_or_else(|| {
candle_core::Error::Msg(
"encode_images_concat: loaded without a vision tower \
(config.json::vision_config absent or weights missing)"
.into(),
)
})?;
let mut per_image = Vec::with_capacity(image_pixels.len());
for (idx, img) in image_pixels.iter().enumerate() {
let embed = tower
.forward(img)
.map_err(|e| candle_core::Error::Msg(format!("encode image[{idx}]: {e:#}")))?;
per_image.push(embed);
}
Tensor::cat(&per_image.iter().collect::<Vec<_>>(), 0)
}
/// Chunked image prefill for the single-GPU path (#18) — parity with
/// `TpQwen3_5ForCausalLM::prefill_with_images_chunked`. Encodes the
/// image(s) once, then walks the (pre-expanded) prompt in
/// `chunk_size`-token windows — exactly like the text
/// `chunked_prefill_*` paths — splicing the patch embeddings into
/// whichever chunk(s) carry `<|image_pad|>` positions. Activation
/// memory is bounded by the chunk, not the full prompt, so a long
/// vision context no longer single-shot-OOMs.
///
/// The KV cache (and GDN recurrent state) accumulate across chunks
/// via the growing offset — the same per-chunk associativity the
/// text chunked prefill and prefix cache (#11/#23) rely on. Only the
/// final chunk's last-position logits are returned; intermediate
/// chunks just populate the cache. The caller is responsible for
/// clearing the cache first.
///
/// `base_offset` is the KV position the prefill starts at (0 for a
/// fresh request). `image_pixels` are device-resident `(C, H, W)`
/// tensors; grids and the interleaved-M-RoPE position ids are
/// recomputed here so an image's position compression is consistent
/// across chunk boundaries.
pub fn prefill_with_images_chunked(
&mut self,
tokens: &[u32],
base_offset: usize,
image_pixels: &[Tensor],
image_token_id: u32,
chunk_size: usize,
) -> candle_core::Result<Tensor> {
if image_pixels.is_empty() {
candle_core::bail!("prefill_with_images_chunked: called with zero images");
}
if tokens.is_empty() {
candle_core::bail!("prefill_with_images_chunked: empty prompt");
}
let chunk_size = chunk_size.max(1);
let device = self.base.device.clone();
let image_embeds = self.encode_images_concat(image_pixels)?;
// Each image's LM grid (lm_gh, lm_gw) = (h/factor, w/factor),
// factor = patch×merge — recomputed from the pixel tensors (#14
// dynamic resolution).
let factor = self
.vision
.as_ref()
.map(|v| {
let c = v.config();
c.patch_size * c.spatial_merge_size
})
.ok_or_else(|| {
candle_core::Error::Msg(
"prefill_with_images_chunked: loaded without a vision tower".into(),
)
})?;
let grids: Vec<(usize, usize)> = image_pixels
.iter()
.map(|t| {
let (_, h, w) = t.dims3()?;
Ok::<(usize, usize), candle_core::Error>((h / factor, w / factor))
})
.collect::<candle_core::Result<Vec<_>>>()?;
// Interleaved-M-RoPE 3D positions for the whole prompt, computed
// once and sliced per chunk so image tokens get their grid
// coordinates and text after an image resumes from the
// compressed counter. `rope_delta` is stashed on the base model
// for the decode that follows this prefill.
let (text, height, width, delta) = rope::get_rope_index(tokens, image_token_id, &grids)
.map_err(|e| candle_core::Error::Msg(format!("get_rope_index: {e}")))?;
self.base.rope_delta = delta;
let full_pos = rope::mrope_position_tensor(&text, &height, &width, &device)?;
let mut last_logits: Option<Tensor> = None;
// Rows of `image_embeds` already spliced by earlier chunks. The
// `<|image_pad|>` run is contiguous, so chunks consume embedding
// rows in order.
let mut img_off = 0usize;
let mut start = 0usize;
while start < tokens.len() {
let end = (start + chunk_size).min(tokens.len());
let chunk = &tokens[start..end];
let input = Tensor::new(chunk, &device)?.unsqueeze(0)?;
let pos_slice = full_pos.narrow(1, start, end - start)?;
let n_here = chunk.iter().filter(|&&t| t == image_token_id).count();
let logits = if n_here == 0 {
self.forward_with_positions(&input, base_offset + start, &pos_slice, None, None)?
} else {
// Splice the next `n_here` patch rows at this chunk's
// local image-pad positions.
let rows = image_embeds.narrow(0, img_off, n_here)?;
img_off += n_here;
self.forward_with_positions(
&input,
base_offset + start,
&pos_slice,
Some(&rows),
Some(image_token_id),
)?
};
last_logits = Some(logits);
start = end;
}
last_logits
.ok_or_else(|| candle_core::Error::Msg("prefill_with_images_chunked: no chunks".into()))
}
pub fn clear_kv_cache(&mut self) {
self.base.clear_kv_cache();
}
/// See [`Qwen3_5Model::snapshot_kv_cache`].
pub fn snapshot_kv_cache(&self) -> candle_core::Result<snapshot::KvCacheSnapshot> {
self.base.snapshot_kv_cache()
}
/// See [`Qwen3_5Model::restore_kv_cache`].
pub fn restore_kv_cache(
&mut self,
snap: &snapshot::KvCacheSnapshot,
) -> candle_core::Result<()> {
self.base.restore_kv_cache(snap)
}
}
#[cfg(test)]

View File

@@ -1,19 +1,27 @@
//! Rotary position embedding for Qwen3-Next's full-attention layers.
//!
//! Qwen3.6 ships with MRoPE (multimodal RoPE) machinery in the
//! reference Python — three position grids interleaved per
//! `mrope_section`. For text-only inference all three grids carry the
//! same position ids and the interleave is a no-op, so this module
//! implements the plain (non-mrope) flavour: the standard inv_freq
//! cosine/sine tables driven by `rope_theta` and `head_dim`.
//! Qwen3.6 declares **interleaved M-RoPE** (multimodal RoPE): the
//! rotary half-dimension is split across three position axes —
//! `[text, height, width]` per `mrope_section` (`[11,11,10]` for
//! Qwen3.6) — interleaved per-frequency. For **text** every token's
//! three axes carry the same position id, so the interleave is a no-op
//! and this reduces exactly to plain RoPE. For **image** tokens the
//! height/width axes carry the patch's 2D grid coordinates, which is
//! how the model reads the 14×14 patch layout (without it, all patches
//! share a height position and the image reads as vertical repetition).
//!
//! Rotation flavour: **GLM-style** rotate-half (the second half of the
//! head dim is negated and swapped into the first). The reference
//! Python uses `apply_rotary_pos_emb` with `rotate_half`; candle's
//! `rope_slow` is the matching helper.
//! Two cos/sin builders feed a shared [`RotaryEmbedding::apply`]:
//! - [`RotaryEmbedding::plain_cos_sin`] narrows the precomputed tables
//! at a scalar position — the text / decode fast path.
//! - [`RotaryEmbedding::mrope_cos_sin`] builds per-token cos/sin from a
//! `(3, seq)` position-id tensor, blending the three axes' frequencies
//! at the interleave index sets — the vision-prefill path.
//!
//! Rotation flavour: **GLM-style** rotate-half (candle's `rope_slow`),
//! matching the reference Python's `apply_rotary_pos_emb` + `rotate_half`.
use anyhow::Result;
use candle_core::{DType, Device, Tensor};
use candle_core::{DType, Device, IndexOp, Tensor};
use super::TextConfig;
@@ -21,6 +29,18 @@ use super::TextConfig;
pub struct RotaryEmbedding {
sin: Tensor,
cos: Tensor,
/// Inverse frequencies, shape `(1, rotary_dim/2)`. Retained (beyond
/// the precomputed `sin`/`cos` tables) so [`Self::mrope_cos_sin`] can
/// build cos/sin from arbitrary per-axis position ids.
inv_freq: Tensor,
/// Per-axis column masks over the rotary half-dim, shape `(1, half)`,
/// f32 0/1. `mask_t + mask_h + mask_w` partitions the columns; a
/// column belongs to exactly one axis. For a non-MRoPE config
/// `mask_t` is all-ones and the others all-zero (→ plain RoPE).
mask_t: Tensor,
mask_h: Tensor,
mask_w: Tensor,
dtype: DType,
/// Number of dims at the head's leading edge that the rotation
/// covers. The remaining `head_dim - rotary_dim` dims pass through
/// unchanged. Qwen3-Next uses `partial_rotary_factor = 0.25`, so
@@ -29,6 +49,52 @@ pub struct RotaryEmbedding {
head_dim: usize,
}
/// Build the per-axis 0/1 column masks over the rotary half-dim from
/// `mrope_section`. Returns `(temporal, height, width)` each length
/// `half`. Temporal is the complement of height width, so the three
/// masks always partition `0..half` and reduce to all-temporal (plain
/// RoPE) when no usable section is given.
fn mrope_masks(
half: usize,
section: &[usize],
interleaved: bool,
) -> (Vec<f32>, Vec<f32>, Vec<f32>) {
let mut mh = vec![0f32; half];
let mut mw = vec![0f32; half];
if section.len() == 3 {
if interleaved {
// Qwen3-VL: height at columns 1,4,7,… ; width at 2,5,8,… ;
// temporal keeps 0,3,6,… — each `take`n from `mrope_section`.
for i in (1..half).step_by(3).take(section[1]) {
mh[i] = 1.0;
}
for i in (2..half).step_by(3).take(section[2]) {
mw[i] = 1.0;
}
} else {
// Qwen2-VL: contiguous blocks [text | height | width].
let h_start = section[0].min(half);
let h_end = (section[0] + section[1]).min(half);
for m in mh.iter_mut().take(h_end).skip(h_start) {
*m = 1.0;
}
for m in mw.iter_mut().take(half).skip(h_end) {
*m = 1.0;
}
}
}
let mt: Vec<f32> = (0..half)
.map(|i| {
if mh[i] == 0.0 && mw[i] == 0.0 {
1.0
} else {
0.0
}
})
.collect();
(mt, mh, mw)
}
impl RotaryEmbedding {
pub fn new(dtype: DType, cfg: &TextConfig, dev: &Device) -> Result<Self> {
let head_dim = cfg.head_dim;
@@ -52,44 +118,88 @@ impl RotaryEmbedding {
.step_by(2)
.map(|i| 1f32 / rope.rope_theta.powf(i as f64 / rotary_dim as f64) as f32)
.collect();
let n = inv_freq.len();
let inv_freq = Tensor::from_vec(inv_freq, (1, n), dev)?.to_dtype(DType::F32)?;
let half = inv_freq.len();
let inv_freq = Tensor::from_vec(inv_freq, (1, half), dev)?.to_dtype(DType::F32)?;
let t = Tensor::arange(0u32, max_seq_len as u32, dev)?
.to_dtype(DType::F32)?
.reshape((max_seq_len, 1))?;
let freqs = t.matmul(&inv_freq)?;
// MRoPE axis masks. `sum(mrope_section)` should equal `half`;
// warn-tolerant: any shortfall just stays on the temporal axis.
let (mt, mh, mw) = mrope_masks(half, &rope.mrope_section, rope.mrope_interleaved);
let mask_t = Tensor::from_vec(mt, (1, half), dev)?;
let mask_h = Tensor::from_vec(mh, (1, half), dev)?;
let mask_w = Tensor::from_vec(mw, (1, half), dev)?;
Ok(Self {
sin: freqs.sin()?.to_dtype(dtype)?,
cos: freqs.cos()?.to_dtype(dtype)?,
inv_freq,
mask_t,
mask_h,
mask_w,
dtype,
rotary_dim,
head_dim,
})
}
/// Apply RoPE to q, k.
///
/// `q`, `k` shape: `(B, H, L, head_dim)`. `offset` is the index
/// into the cached cos/sin table — the position of the first token
/// in the current step.
///
/// When `rotary_dim < head_dim` the rotation is applied only to the
/// first `rotary_dim` dims of each head; the tail passes through
/// unchanged (matches the reference Python's
/// `apply_rotary_pos_emb` with non-trivial `partial_rotary_factor`).
pub fn apply(
/// cos/sin for a contiguous run of `seq_len` positions starting at
/// `pos`, by narrowing the precomputed tables. The text / decode
/// path (all three MRoPE axes equal → plain RoPE). Shape
/// `(seq_len, rotary_dim/2)`.
pub fn plain_cos_sin(
&self,
pos: usize,
seq_len: usize,
) -> candle_core::Result<(Tensor, Tensor)> {
let cos = self.cos.narrow(0, pos, seq_len)?;
let sin = self.sin.narrow(0, pos, seq_len)?;
Ok((cos, sin))
}
/// cos/sin from explicit per-token 3D position ids, shape
/// `(3, seq_len)` (axes: text, height, width). Builds each axis's
/// frequencies and blends them at the interleave index sets, so
/// every rotary frequency slot is driven by exactly one axis.
/// Reduces exactly to [`Self::plain_cos_sin`] when the three axes are
/// equal. Returns cos/sin of shape `(seq_len, rotary_dim/2)`.
pub fn mrope_cos_sin(&self, position_ids: &Tensor) -> candle_core::Result<(Tensor, Tensor)> {
let pos = position_ids.to_dtype(DType::F32)?;
let (axes, seq_len) = pos.dims2()?;
debug_assert_eq!(axes, 3, "mrope position_ids must have 3 axes");
// Per-axis freqs: pos[a] (seq,1) @ inv_freq (1,half) → (seq,half).
let ft = pos.i(0)?.reshape((seq_len, 1))?.matmul(&self.inv_freq)?;
let fh = pos.i(1)?.reshape((seq_len, 1))?.matmul(&self.inv_freq)?;
let fw = pos.i(2)?.reshape((seq_len, 1))?.matmul(&self.inv_freq)?;
// Blend: each column belongs to exactly one axis (masks partition
// the half-dim), so this picks the right axis per frequency slot.
let blended = ft
.broadcast_mul(&self.mask_t)?
.add(&fh.broadcast_mul(&self.mask_h)?)?
.add(&fw.broadcast_mul(&self.mask_w)?)?;
let cos = blended.cos()?.to_dtype(self.dtype)?;
let sin = blended.sin()?.to_dtype(self.dtype)?;
Ok((cos, sin))
}
/// Apply rotary to `q`, `k` (shape `(B, H, L, head_dim)`) using
/// precomputed `cos`/`sin` of shape `(L, rotary_dim/2)`. Partial
/// rotary: only the first `rotary_dim` dims rotate; the tail passes
/// through unchanged.
pub fn apply_cos_sin(
&self,
q: &Tensor,
k: &Tensor,
offset: usize,
cos: &Tensor,
sin: &Tensor,
) -> candle_core::Result<(Tensor, Tensor)> {
let (_, _, seq_len, head_dim_in) = q.dims4()?;
let (_, _, _seq_len, head_dim_in) = q.dims4()?;
debug_assert_eq!(head_dim_in, self.head_dim, "q head_dim mismatch");
let cos = self.cos.narrow(0, offset, seq_len)?;
let sin = self.sin.narrow(0, offset, seq_len)?;
if self.rotary_dim == self.head_dim {
// Full rotation.
let q_embed = candle_nn::rotary_emb::rope_slow(&q.contiguous()?, &cos, &sin)?;
let k_embed = candle_nn::rotary_emb::rope_slow(&k.contiguous()?, &cos, &sin)?;
let q_embed = candle_nn::rotary_emb::rope_slow(&q.contiguous()?, cos, sin)?;
let k_embed = candle_nn::rotary_emb::rope_slow(&k.contiguous()?, cos, sin)?;
Ok((q_embed, k_embed))
} else {
// Partial rotation: narrow → rotate → cat the untouched tail.
@@ -102,8 +212,8 @@ impl RotaryEmbedding {
.narrow(candle_core::D::Minus1, 0, self.rotary_dim)?
.contiguous()?;
let k_pass = k.narrow(candle_core::D::Minus1, self.rotary_dim, tail)?;
let q_rotated = candle_nn::rotary_emb::rope_slow(&q_rot, &cos, &sin)?;
let k_rotated = candle_nn::rotary_emb::rope_slow(&k_rot, &cos, &sin)?;
let q_rotated = candle_nn::rotary_emb::rope_slow(&q_rot, cos, sin)?;
let k_rotated = candle_nn::rotary_emb::rope_slow(&k_rot, cos, sin)?;
let q_embed =
Tensor::cat(&[&q_rotated, &q_pass.contiguous()?], candle_core::D::Minus1)?;
let k_embed =
@@ -112,3 +222,358 @@ impl RotaryEmbedding {
}
}
}
/// Compute interleaved-M-RoPE 3D position ids for a full prompt that may
/// contain image-placeholder runs, plus the decode `rope_delta`.
///
/// Mirrors the reference `get_rope_index`:
/// - text tokens advance a single running counter `c`, all three axes
/// equal (`[c, c, c]`);
/// - each contiguous run of `image_token_id` is one image; its tokens get
/// `[base + t, base + h, base + w]` in row-major (t outer, h, w inner),
/// where `base` is the counter at the run's start; after the run the
/// counter resumes from `base + max(grid_t, grid_h, grid_w)`.
///
/// Returns `(text_pos, height_pos, width_pos, rope_delta)`, each pos `Vec`
/// length `input_ids.len()`. `rope_delta = final_counter - seq_len`: add it
/// to a plain decode offset so text resumes from the counter after the
/// (position-compressed) image blocks.
///
/// Whether interleaved M-RoPE for image tokens is enabled. Default
/// **on** — Qwen3.6 was trained with interleaved M-RoPE, and this
/// implementation matches the HF `apply_interleaved_mrope` /
/// `get_rope_index` reference exactly (verified column-for-column). The
/// env var is a **kill switch**: `NEURON_MROPE=0` falls back to plain
/// sequential positions for image tokens (the pre-M-RoPE behaviour).
pub(crate) fn mrope_enabled() -> bool {
std::env::var("NEURON_MROPE")
.map(|v| {
!matches!(
v.trim().to_ascii_lowercase().as_str(),
"0" | "false" | "no" | "off"
)
})
.unwrap_or(true)
}
/// Position ids for the forward path. Gated by [`mrope_enabled`]: when
/// off, returns plain sequential identity positions on all three axes
/// (`mrope_cos_sin` then reduces exactly to plain RoPE), restoring the
/// pre-M-RoPE behaviour without touching the rest of the forward.
pub(crate) fn get_rope_index(
input_ids: &[u32],
image_token_id: u32,
grids: &[(usize, usize)],
) -> Result<MRopeIndex> {
if !mrope_enabled() {
let seq: Vec<i64> = (0..input_ids.len() as i64).collect();
return Ok((seq.clone(), seq.clone(), seq, 0));
}
compute_mrope_index(input_ids, image_token_id, grids)
}
/// The real interleaved-M-RoPE position-id computation (always active in
/// unit tests; gated behind [`get_rope_index`] at runtime).
///
/// `grids` carries the post-merge LM grid `(lm_gh, lm_gw)` for each image
/// run, in prompt order — a run length alone cannot recover its
/// factorisation, so the grids must be passed (#14 dynamic resolution).
/// Each image is a still frame (`grid_t = 1`); its tokens get
/// `[base, base + hh, base + ww]` row-major and the shared counter
/// resumes at `base + max(lm_gh, lm_gw)`. Multi-image is correct because
/// the counter threads across images and interleaved text.
pub(crate) fn compute_mrope_index(
input_ids: &[u32],
image_token_id: u32,
grids: &[(usize, usize)],
) -> Result<MRopeIndex> {
let n = input_ids.len();
let mut text = Vec::with_capacity(n);
let mut height = Vec::with_capacity(n);
let mut width = Vec::with_capacity(n);
let mut counter: i64 = 0;
let mut i = 0;
let mut k = 0; // index into `grids`, one per image run
while i < n {
if input_ids[i] == image_token_id {
let start = i;
while i < n && input_ids[i] == image_token_id {
i += 1;
}
let run = i - start;
let (grid_h, grid_w) = *grids.get(k).ok_or_else(|| {
anyhow::anyhow!(
"get_rope_index: image run #{k} (len {run}) has no matching grid \
({} grids supplied)",
grids.len()
)
})?;
k += 1;
if grid_h * grid_w != run {
anyhow::bail!(
"get_rope_index: image run #{} length {run} != grid {grid_h}×{grid_w} = {}",
k - 1,
grid_h * grid_w
);
}
let base = counter;
for hh in 0..grid_h {
for ww in 0..grid_w {
text.push(base); // grid_t = 1 → temporal axis const
height.push(base + hh as i64);
width.push(base + ww as i64);
}
}
counter = base + grid_h.max(grid_w) as i64;
} else {
text.push(counter);
height.push(counter);
width.push(counter);
counter += 1;
i += 1;
}
}
if k != grids.len() {
anyhow::bail!(
"get_rope_index: prompt has {k} image run(s) but {} grid(s) were supplied",
grids.len()
);
}
let delta = counter - n as i64;
Ok((text, height, width, delta))
}
/// `(text_pos, height_pos, width_pos, rope_delta)` returned by
/// [`get_rope_index`]; the three vectors combine into the `(3, seq)`
/// MRoPE position-id tensor.
pub(crate) type MRopeIndex = (Vec<i64>, Vec<i64>, Vec<i64>, i64);
/// Build the `(3, seq)` position-id tensor consumed by
/// [`RotaryEmbedding::mrope_cos_sin`] from the three axis vectors.
///
/// Built directly as **f32** (positions are small integers, exact in
/// f32 well past any context length): the freqs matmul needs float
/// anyway, and this avoids an i64 tensor / i64→f32 cast on the GPU.
pub(crate) fn mrope_position_tensor(
text: &[i64],
height: &[i64],
width: &[i64],
dev: &Device,
) -> candle_core::Result<Tensor> {
let seq = text.len();
let mut flat = Vec::with_capacity(3 * seq);
flat.extend(text.iter().map(|&x| x as f32));
flat.extend(height.iter().map(|&x| x as f32));
flat.extend(width.iter().map(|&x| x as f32));
Tensor::from_vec(flat, (3, seq), dev)
}
#[cfg(test)]
mod tests {
use super::*;
use candle_core::IndexOp;
/// A TextConfig stub with Qwen3.6's rope params (head_dim 256,
/// partial 0.25 → rotary_dim 64 → half 32; section [11,11,10]).
fn qwen36_cfg() -> TextConfig {
serde_json::from_value(serde_json::json!({
"hidden_size": 5120,
"num_hidden_layers": 1,
"num_attention_heads": 64,
"num_key_value_heads": 8,
"head_dim": 256,
"intermediate_size": 1,
"vocab_size": 10,
"rms_norm_eps": 1e-6,
"max_position_embeddings": 64,
"layer_types": ["full_attention"],
"rope_parameters": {
"rope_theta": 10000000.0,
"partial_rotary_factor": 0.25,
"mrope_section": [11, 11, 10],
"mrope_interleaved": true
}
}))
.expect("cfg")
}
#[test]
fn mrope_masks_partition_the_half_dim() {
let (mt, mh, mw) = mrope_masks(32, &[11, 11, 10], true);
// Each column belongs to exactly one axis.
for i in 0..32 {
let s = mt[i] + mh[i] + mw[i];
assert_eq!(s, 1.0, "column {i} covered {s} times");
}
assert_eq!(mt.iter().sum::<f32>(), 11.0);
assert_eq!(mh.iter().sum::<f32>(), 11.0);
assert_eq!(mw.iter().sum::<f32>(), 10.0);
// Interleave: temporal 0,3,…; height 1,4,…; width 2,5,…
assert_eq!(mt[0], 1.0);
assert_eq!(mh[1], 1.0);
assert_eq!(mw[2], 1.0);
assert_eq!(mt[3], 1.0);
}
/// The load-bearing invariant: when all three position axes are
/// equal (text), `mrope_cos_sin` must reproduce `plain_cos_sin`
/// bit-for-bit — i.e. M-RoPE is a no-op for text, so text inference
/// is unchanged.
#[test]
fn mrope_reduces_to_plain_for_equal_axes() {
let dev = Device::Cpu;
let rope = RotaryEmbedding::new(DType::F32, &qwen36_cfg(), &dev).unwrap();
// positions 5,6,7 on all three axes.
let base: Vec<i64> = vec![5, 6, 7];
let pos =
Tensor::from_vec([base.clone(), base.clone(), base].concat(), (3, 3), &dev).unwrap();
let (mc, ms) = rope.mrope_cos_sin(&pos).unwrap();
let (pc, ps) = rope.plain_cos_sin(5, 3).unwrap();
let dcos = (mc - pc).unwrap().abs().unwrap().max_all().unwrap();
let dsin = (ms - ps).unwrap().abs().unwrap().max_all().unwrap();
assert!(
dcos.to_scalar::<f32>().unwrap() < 1e-6,
"cos mismatch {dcos:?}"
);
assert!(
dsin.to_scalar::<f32>().unwrap() < 1e-6,
"sin mismatch {dsin:?}"
);
}
/// Hand-checked interleave: a width-axis column (index 2) must track
/// the WIDTH position, while a temporal column (index 0) tracks the
/// TEXT position, even when the axes differ.
#[test]
fn mrope_blends_axes_at_interleave_columns() {
let dev = Device::Cpu;
let rope = RotaryEmbedding::new(DType::F32, &qwen36_cfg(), &dev).unwrap();
let half = rope.inv_freq.dim(1).unwrap();
let inv: Vec<f32> = rope.inv_freq.i(0).unwrap().to_vec1().unwrap();
// One token: text=10, height=3, width=7 — all distinct.
let pos = Tensor::from_vec(vec![10i64, 3, 7], (3, 1), &dev).unwrap();
let (cos, _sin) = rope.mrope_cos_sin(&pos).unwrap();
let cos_row: Vec<f32> = cos.i(0).unwrap().to_vec1().unwrap();
assert_eq!(cos_row.len(), half);
// Column 0 (temporal) → text pos 10. Column 1 (height) → 3.
// Column 2 (width) → 7.
assert!((cos_row[0] - (10.0 * inv[0]).cos()).abs() < 1e-5);
assert!((cos_row[1] - (3.0 * inv[1]).cos()).abs() < 1e-5);
assert!((cos_row[2] - (7.0 * inv[2]).cos()).abs() < 1e-5);
assert!((cos_row[3] - (10.0 * inv[3]).cos()).abs() < 1e-5);
}
#[test]
fn get_rope_index_text_only_is_sequential() {
let (t, h, w, delta) = compute_mrope_index(&[1, 2, 3, 4], 99, &[]).unwrap();
assert_eq!(t, vec![0, 1, 2, 3]);
assert_eq!(h, vec![0, 1, 2, 3]);
assert_eq!(w, vec![0, 1, 2, 3]);
assert_eq!(delta, 0, "no image → delta 0 → plain decode positions");
}
#[test]
fn get_rope_index_text_image_text() {
// [text, image(2x2 run of 4), text]. image_token = 99, grid (2,2).
let ids = [1u32, 99, 99, 99, 99, 2];
let (t, h, w, delta) = compute_mrope_index(&ids, 99, &[(2, 2)]).unwrap();
// token 0: text → 0. image base=1, grid 2x2:
// t all = 1; h = base+row = [1,1,2,2]; w = base+col = [1,2,1,2].
// resume from base + max(2,2) = 3. trailing text → 3.
assert_eq!(t, vec![0, 1, 1, 1, 1, 3]);
assert_eq!(h, vec![0, 1, 1, 2, 2, 3]);
assert_eq!(w, vec![0, 1, 2, 1, 2, 3]);
// final counter = 4, seq_len = 6 → delta = -2 (the 4 image tokens
// advanced the counter by only 2).
assert_eq!(delta, -2);
// Decode after the prompt (offset = 6) → text position 6 + (-2) = 4.
assert_eq!(6 + delta, 4);
}
#[test]
fn get_rope_index_nonsquare_single_image() {
// text + image(2 rows × 3 cols = 6 tokens). grid (2,3).
let ids = [1u32, 99, 99, 99, 99, 99, 99];
let (t, h, w, delta) = compute_mrope_index(&ids, 99, &[(2, 3)]).unwrap();
// base = 1; row-major h = [0,0,0,1,1,1]+1, w = [0,1,2,0,1,2]+1.
assert_eq!(t, vec![0, 1, 1, 1, 1, 1, 1]);
assert_eq!(h, vec![0, 1, 1, 1, 2, 2, 2]);
assert_eq!(w, vec![0, 1, 2, 3, 1, 2, 3]);
// resume from base + max(2,3) = 4; seq_len 7, counter 4 → delta -3.
assert_eq!(delta, 4 - 7);
}
#[test]
fn get_rope_index_two_images_different_grids() {
// img(2x2)=4, text, img(1x3)=3. grids [(2,2),(1,3)].
let ids = [99, 99, 99, 99, 7, 99, 99, 99];
let (t, h, w, delta) = compute_mrope_index(&ids, 99, &[(2, 2), (1, 3)]).unwrap();
// img1 base=0 → t=0, h=[0,0,1,1], w=[0,1,0,1]; resume max(2,2)=2.
// text at counter 2. img2 base=3 → t=3, h=[3,3,3], w=[3,4,5];
// resume 3+max(1,3)=6.
assert_eq!(t, vec![0, 0, 0, 0, 2, 3, 3, 3]);
assert_eq!(h, vec![0, 0, 1, 1, 2, 3, 3, 3]);
assert_eq!(w, vec![0, 1, 0, 1, 2, 3, 4, 5]);
assert_eq!(delta, 6 - 8);
}
#[test]
fn get_rope_index_on_by_default() {
// With NEURON_MROPE unset (default ON), the runtime path returns
// the real interleaved-M-RoPE positions. (NEURON_MROPE=0 would fall
// back to identity; not asserted here since it depends on env.)
let (t, h, w, _delta) = get_rope_index(&[1, 99, 99, 99, 99, 2], 99, &[(2, 2)]).unwrap();
assert_eq!(t, vec![0, 1, 1, 1, 1, 3]);
assert_eq!(h, vec![0, 1, 1, 2, 2, 3]);
assert_eq!(w, vec![0, 1, 2, 1, 2, 3]);
}
#[test]
fn get_rope_index_grid_mismatches_error() {
// run length != grid product.
assert!(compute_mrope_index(&[99u32; 6], 99, &[(2, 2)]).is_err());
// too few grids for the number of image runs.
assert!(compute_mrope_index(&[99, 99, 7, 99], 99, &[(1, 2)]).is_err());
// too many grids.
assert!(compute_mrope_index(&[99, 99], 99, &[(1, 2), (1, 1)]).is_err());
}
#[test]
fn position_tensor_round_trips_through_mrope_cos_sin() {
// get_rope_index → (3,seq) tensor → mrope_cos_sin, and confirm an
// image token's height column tracks its grid row (not the text
// counter), i.e. the end-to-end position plumbing is wired right.
let dev = Device::Cpu;
let rope = RotaryEmbedding::new(DType::F32, &qwen36_cfg(), &dev).unwrap();
let ids = [1u32, 99, 99, 99, 99]; // text + 2x2 image
let (t, h, w, _d) = compute_mrope_index(&ids, 99, &[(2, 2)]).unwrap();
let pos = mrope_position_tensor(&t, &h, &w, &dev).unwrap();
assert_eq!(pos.dims(), &[3, 5]);
let (cos, _sin) = rope.mrope_cos_sin(&pos).unwrap();
assert_eq!(cos.dims(), &[5, rope.inv_freq.dim(1).unwrap()]);
let inv: Vec<f32> = rope.inv_freq.i(0).unwrap().to_vec1().unwrap();
// Last image token (index 4): grid (h=1, w=1) → base 1 → h=2, w=2.
// Height column (index 1) must track h-position 2, not text.
let last: Vec<f32> = cos.i(4).unwrap().to_vec1().unwrap();
assert!((last[1] - (2.0 * inv[1]).cos()).abs() < 1e-5);
}
#[test]
fn get_rope_index_196_is_14x14() {
let mut ids = vec![1u32]; // one text token
ids.extend(std::iter::repeat_n(99u32, 196));
let (t, h, w, _delta) = compute_mrope_index(&ids, 99, &[(14, 14)]).unwrap();
// image base = 1. Last image token (index 196) is grid (h=13,w=13).
assert_eq!(*t.last().unwrap(), 1, "grid_t=1 → temporal const at base");
assert_eq!(h[1], 1, "first image row at base");
assert_eq!(w[1], 1, "first image col at base");
assert_eq!(h[196], 1 + 13, "last image row = base + 13");
assert_eq!(w[196], 1 + 13, "last image col = base + 13");
}
}

View File

@@ -0,0 +1,299 @@
//! Cache-state snapshots for prefix KV caching (#11).
//!
//! A snapshot captures everything `clear_kv_cache` would destroy, at
//! one consistent token boundary:
//!
//! - full-attention layers: the `ConcatKvCache` k/v tensors,
//! - linear-attention layers: the GatedDeltaNet `conv_state` +
//! `recurrent_state`,
//! - the model-level `rope_delta` position counter.
//!
//! The GatedDeltaNet recurrent state cannot be rewound to an earlier
//! token, so a snapshot is only reusable when its entire token
//! sequence is an exact prefix of an incoming prompt — matching policy
//! lives in `harness/prefix_cache.rs`; this module is just the state
//! capture.
//!
//! ## Copy semantics
//!
//! Attention k/v snapshots share storage with the live cache:
//! `ConcatKvCache::append` never mutates stored tensors in place (it
//! `cat`s into fresh allocations), so a shallow `Tensor` clone stays
//! valid after the live cache moves on. The GDN states are
//! **deep-copied** in both directions (`Tensor::copy`): the CUDA
//! delta-rule kernels update the recurrent-state buffer in place, and
//! `flatten`/`contiguous` on an already-contiguous tensor is a view —
//! a shared-storage snapshot would be corrupted by the next forward.
use candle_core::Tensor;
/// Per-layer captured state. Variant kind must match the layer's
/// `AttentionKind` on restore.
pub enum LayerKvSnapshot {
/// `ConcatKvCache` contents. `None` when the cache was empty
/// (a zero-token snapshot — valid but useless; the registry never
/// stores one).
Full(Option<(Tensor, Tensor)>),
/// GatedDeltaNet state. Either tensor is `None` before the first
/// forward touches it.
Linear {
conv_state: Option<Tensor>,
recurrent_state: Option<Tensor>,
},
}
/// One consistent cache snapshot of a `Qwen3_5Model` (or its TP
/// mirror `tp_qwen3_5::TpQwen3_5Model`, whose per-rank shard state
/// has the same shape) at a token boundary. Fields are `pub(crate)`
/// so the TP module can construct/consume the same type; holders
/// outside the harness only ever pass it back to `restore_kv_cache`.
pub struct KvCacheSnapshot {
pub(crate) layers: Vec<LayerKvSnapshot>,
pub(crate) rope_delta: i64,
}
impl KvCacheSnapshot {
/// Number of layer snapshots held (test/diagnostic helper).
pub fn layer_count(&self) -> usize {
self.layers.len()
}
/// Total bytes of tensor data held by this snapshot. Used for the
/// prefix-cache VRAM budget. Attention k/v shares storage with the
/// live cache at capture time, but the live cache is cleared or
/// replaced before the next request, so counting the full size is
/// the honest steady-state figure.
pub fn size_bytes(&self) -> u64 {
fn t_bytes(t: &Tensor) -> u64 {
(t.elem_count() * t.dtype().size_in_bytes()) as u64
}
self.layers
.iter()
.map(|l| match l {
LayerKvSnapshot::Full(Some((k, v))) => t_bytes(k) + t_bytes(v),
LayerKvSnapshot::Full(None) => 0,
LayerKvSnapshot::Linear {
conv_state,
recurrent_state,
} => {
conv_state.as_ref().map(t_bytes).unwrap_or(0)
+ recurrent_state.as_ref().map(t_bytes).unwrap_or(0)
}
})
.sum()
}
}
#[cfg(test)]
mod tests {
use super::super::{Qwen3_5Model, RopeParameters, TextConfig};
use candle_core::{DType, Device, Tensor};
use std::collections::HashMap;
/// Tiny two-layer config covering both attention kinds.
fn tiny_config() -> TextConfig {
TextConfig {
vocab_size: 32,
hidden_size: 16,
intermediate_size: 32,
num_hidden_layers: 2,
num_attention_heads: 2,
num_key_value_heads: 1,
head_dim: 8,
max_position_embeddings: 64,
rope_parameters: RopeParameters {
rope_theta: 10000.0,
partial_rotary_factor: 0.5,
rope_type: None,
mrope_section: Vec::new(),
mrope_interleaved: false,
},
rms_norm_eps: 1e-6,
tie_word_embeddings: true,
attn_output_gate: true,
layer_types: vec!["linear_attention".into(), "full_attention".into()],
full_attention_interval: Some(4),
hidden_act: "silu".into(),
linear_num_value_heads: 4,
linear_num_key_heads: 2,
linear_key_head_dim: 4,
linear_value_head_dim: 4,
linear_conv_kernel_dim: 4,
}
}
/// Build a Qwen3_5Model from random weights written to a temp
/// safetensors file — the same `ShardedVarBuilder` path the real
/// loader uses.
fn tiny_model(cfg: &TextConfig) -> Qwen3_5Model {
let dev = Device::Cpu;
let randn = |shape: &[usize]| Tensor::randn(0f32, 0.2f32, shape, &dev).unwrap();
let h = cfg.hidden_size;
let inter = cfg.intermediate_size;
let key_dim = cfg.linear_key_head_dim * cfg.linear_num_key_heads;
let value_dim = cfg.linear_value_head_dim * cfg.linear_num_value_heads;
let conv_dim = key_dim * 2 + value_dim;
let nv = cfg.linear_num_value_heads;
let hd = cfg.head_dim;
let q_out = cfg.num_attention_heads * hd * 2;
let kv_out = cfg.num_key_value_heads * hd;
let mut t: HashMap<String, Tensor> = HashMap::new();
let p = "model.language_model";
t.insert(
format!("{p}.embed_tokens.weight"),
randn(&[cfg.vocab_size, h]),
);
t.insert(format!("{p}.norm.weight"), randn(&[h]));
for (i, kind) in cfg.layer_types.iter().enumerate() {
let lp = format!("{p}.layers.{i}");
t.insert(format!("{lp}.input_layernorm.weight"), randn(&[h]));
t.insert(format!("{lp}.post_attention_layernorm.weight"), randn(&[h]));
t.insert(format!("{lp}.mlp.gate_proj.weight"), randn(&[inter, h]));
t.insert(format!("{lp}.mlp.up_proj.weight"), randn(&[inter, h]));
t.insert(format!("{lp}.mlp.down_proj.weight"), randn(&[h, inter]));
match kind.as_str() {
"linear_attention" => {
let ap = format!("{lp}.linear_attn");
t.insert(format!("{ap}.in_proj_qkv.weight"), randn(&[conv_dim, h]));
t.insert(format!("{ap}.in_proj_z.weight"), randn(&[value_dim, h]));
t.insert(format!("{ap}.in_proj_b.weight"), randn(&[nv, h]));
t.insert(format!("{ap}.in_proj_a.weight"), randn(&[nv, h]));
t.insert(format!("{ap}.out_proj.weight"), randn(&[h, value_dim]));
t.insert(
format!("{ap}.conv1d.weight"),
randn(&[conv_dim, 1, cfg.linear_conv_kernel_dim]),
);
t.insert(format!("{ap}.dt_bias"), randn(&[nv]));
t.insert(format!("{ap}.A_log"), randn(&[nv]));
t.insert(
format!("{ap}.norm.weight"),
randn(&[cfg.linear_value_head_dim]),
);
}
"full_attention" => {
let ap = format!("{lp}.self_attn");
t.insert(format!("{ap}.q_proj.weight"), randn(&[q_out, h]));
t.insert(format!("{ap}.k_proj.weight"), randn(&[kv_out, h]));
t.insert(format!("{ap}.v_proj.weight"), randn(&[kv_out, h]));
t.insert(
format!("{ap}.o_proj.weight"),
randn(&[h, cfg.num_attention_heads * hd]),
);
t.insert(format!("{ap}.q_norm.weight"), randn(&[hd]));
t.insert(format!("{ap}.k_norm.weight"), randn(&[hd]));
}
other => panic!("unexpected layer type {other}"),
}
}
let dir = tempfile::tempdir().expect("tempdir");
let path = dir.path().join("model.safetensors");
candle_core::safetensors::save(&t, &path).expect("save safetensors");
// SAFETY: mmap of a file this test just wrote and nothing else
// mutates — same justification as the real loader.
let vb = unsafe {
candle_nn::var_builder::ShardedSafeTensors::var_builder(
&[path.clone()],
DType::F32,
&dev,
)
.expect("build ShardedVarBuilder")
};
Qwen3_5Model::load(cfg, &vb).expect("load tiny qwen3_5 model")
}
fn forward_tokens(model: &mut Qwen3_5Model, tokens: &[u32], offset: usize) -> Vec<f32> {
let input = Tensor::new(tokens, &Device::Cpu)
.unwrap()
.unsqueeze(0)
.unwrap();
let hidden = model.forward(&input, offset).unwrap();
// Last-position hidden row — what the lm_head would consume.
let (_, l, _) = hidden.dims3().unwrap();
hidden
.narrow(1, l - 1, 1)
.unwrap()
.flatten_all()
.unwrap()
.to_vec1()
.unwrap()
}
fn max_abs_diff(a: &[f32], b: &[f32]) -> f32 {
assert_eq!(a.len(), b.len());
a.iter()
.zip(b)
.map(|(x, y)| (x - y).abs())
.fold(0f32, f32::max)
}
/// The gold test for #11: prefill a prefix, snapshot, perturb the
/// live state with unrelated tokens, restore, prefill only the
/// suffix — the result must match a fresh full prefill. Exercises
/// attention KV, GDN conv/recurrent state, and offset bookkeeping
/// in one pass; the perturbation step would corrupt a
/// shared-storage (non-deep-copied) GDN snapshot.
#[test]
fn restore_then_suffix_matches_full_prefill() {
let cfg = tiny_config();
let mut model = tiny_model(&cfg);
let prefix: &[u32] = &[1, 2, 3];
let suffix: &[u32] = &[4, 5, 6];
let full: Vec<u32> = prefix.iter().chain(suffix).copied().collect();
model.clear_kv_cache();
let h_full = forward_tokens(&mut model, &full, 0);
model.clear_kv_cache();
forward_tokens(&mut model, prefix, 0);
let snap = model.snapshot_kv_cache().expect("snapshot");
assert_eq!(snap.layer_count(), 2);
assert!(snap.size_bytes() > 0);
// Advance the live state past the snapshot boundary — a
// different continuation, as a subsequent request would be.
forward_tokens(&mut model, &[9, 8], prefix.len());
model.restore_kv_cache(&snap).expect("restore");
let h_restored = forward_tokens(&mut model, suffix, prefix.len());
let diff = max_abs_diff(&h_full, &h_restored);
assert!(diff < 1e-4, "restored-prefix forward diverged: {diff}");
// The snapshot must survive restore + forward cycles (deep
// copy of the in-place-mutated GDN state): restore again and
// expect the identical result.
model.restore_kv_cache(&snap).expect("second restore");
let h_again = forward_tokens(&mut model, suffix, prefix.len());
let diff = max_abs_diff(&h_restored, &h_again);
assert!(diff < 1e-6, "second restore diverged: {diff}");
}
/// Restoring must fully replace the live state, not blend with it
/// — a divergent continuation after restore equals the same
/// continuation after a fresh prefill of the prefix.
#[test]
fn restore_replaces_live_state() {
let cfg = tiny_config();
let mut model = tiny_model(&cfg);
let prefix: &[u32] = &[7, 7, 2, 5];
let cont: &[u32] = &[11, 13];
model.clear_kv_cache();
forward_tokens(&mut model, prefix, 0);
let h_fresh = forward_tokens(&mut model, cont, prefix.len());
model.clear_kv_cache();
forward_tokens(&mut model, prefix, 0);
let snap = model.snapshot_kv_cache().expect("snapshot");
forward_tokens(&mut model, &[3, 1, 4, 1, 5], prefix.len());
model.restore_kv_cache(&snap).expect("restore");
let h_restored = forward_tokens(&mut model, cont, prefix.len());
let diff = max_abs_diff(&h_fresh, &h_restored);
assert!(diff < 1e-5, "restore did not replace live state: {diff}");
}
}

View File

@@ -48,6 +48,31 @@ use candle_nn::var_builder::ShardedVarBuilder;
use candle_nn::{Conv2d, Conv2dConfig, Embedding, LayerNorm, Linear};
use serde::Deserialize;
fn env_truthy(name: &str) -> bool {
std::env::var(name)
.map(|v| {
matches!(
v.trim().to_ascii_lowercase().as_str(),
"1" | "true" | "yes" | "on"
)
})
.unwrap_or(false)
}
/// Legacy escape hatch: when set, use the original Stage-A sequential
/// `pos_embed` lookup instead of the bilinear grid interpolation.
/// Default off (interpolation on) — for A/B comparison only.
fn vision_legacy_pos() -> bool {
env_truthy("NEURON_VISION_LEGACY_POS")
}
/// Legacy escape hatch: when set, skip the 2D vision rotary in the ViT
/// attention (the original Stage-A behaviour). Default off (rotary on)
/// — for A/B comparison only.
fn vision_legacy_rope() -> bool {
env_truthy("NEURON_VISION_LEGACY_ROPE")
}
/// Qwen3.6 vision tower hyperparameters. Mirrors the `vision_config`
/// block of `config.json`. Only the fields we actually need are
/// captured; serde tolerates the rest.
@@ -118,10 +143,12 @@ impl VisionBlock {
})
}
/// `x`: `(N, hidden_size)` un-batched. Returns same shape.
fn forward(&self, x: &Tensor) -> Result<Tensor> {
/// `x`: `(N, hidden_size)` un-batched. `rotary`: optional
/// `(cos, sin)` each `(N, head_dim/2)` — the 2D vision rotary applied
/// to q/k. Returns same shape.
fn forward(&self, x: &Tensor, rotary: Option<&(Tensor, Tensor)>) -> Result<Tensor> {
let attn_in = self.norm1.forward(x)?;
let attn_out = self.attention(&attn_in)?;
let attn_out = self.attention(&attn_in, rotary)?;
let x = x.add(&attn_out)?;
let mlp_in = self.norm2.forward(&x)?;
let mlp_out = self.fc2.forward(&gelu_tanh(&self.fc1.forward(&mlp_in)?)?)?;
@@ -129,8 +156,11 @@ impl VisionBlock {
}
/// Multi-head self-attention over the patch sequence. No causal
/// mask — every patch attends to every other patch.
fn attention(&self, x: &Tensor) -> Result<Tensor> {
/// mask — every patch attends to every other patch. When `rotary` is
/// given, the 2D vision rotary (row/col position) is applied to q, k
/// before the scores, matching HF `apply_rotary_pos_emb_vision`
/// (`rope_slow` is the same rotate-half form).
fn attention(&self, x: &Tensor, rotary: Option<&(Tensor, Tensor)>) -> Result<Tensor> {
let (n, hidden) = x.dims2()?;
// qkv: (N, 3*hidden). Split into Q, K, V each (N, hidden).
let qkv = self.qkv.forward(x)?;
@@ -140,6 +170,15 @@ impl VisionBlock {
let q = qkv.i(0)?;
let k = qkv.i(1)?;
let v = qkv.i(2)?;
// 2D vision rotary on q, k (full head_dim; rotate-half form).
let (q, k) = match rotary {
Some((cos, sin)) => {
let q = candle_nn::rotary_emb::rope_slow(&q.unsqueeze(0)?, cos, sin)?.squeeze(0)?;
let k = candle_nn::rotary_emb::rope_slow(&k.unsqueeze(0)?, cos, sin)?.squeeze(0)?;
(q, k)
}
None => (q, k),
};
let scale = 1.0 / (self.head_dim as f64).sqrt();
// (num_heads, N, head_dim) @ (num_heads, head_dim, N) -> (num_heads, N, N)
let scores = q.matmul(&k.transpose(D::Minus2, D::Minus1)?)?;
@@ -210,11 +249,65 @@ impl VisionMerger {
}
}
/// 2D rotary position embedding for the vision tower. Each patch's
/// `head_dim` rotates by its `(row, col)` grid coordinates: the first
/// half of the rotary freqs are driven by the row position, the second
/// half by the column. Mirrors HF `Qwen3VLVisionRotaryEmbedding` +
/// `rot_pos_emb` (θ = 10000, `dim = head_dim/2`).
struct VisionRotaryEmbedding {
/// `(half,)` f32, `half = head_dim/4` freqs per spatial axis.
inv_freq: Vec<f32>,
}
impl VisionRotaryEmbedding {
fn new(head_dim: usize) -> Self {
// HF: Qwen3VLVisionRotaryEmbedding(head_dim // 2), theta 10000.
let dim = head_dim / 2;
let theta = 10000f32;
let inv_freq = (0..dim)
.step_by(2)
.map(|i| 1f32 / theta.powf(i as f32 / dim as f32))
.collect();
Self { inv_freq }
}
/// cos/sin for a `gh×gw` patch grid in **row-major** order. Returns
/// `(cos, sin)` each `(gh*gw, head_dim/2)`: per patch, the row-axis
/// freqs `row·inv_freq` followed by the col-axis freqs `col·inv_freq`
/// (then `rope_slow` duplicates them across the full head_dim).
fn cos_sin(
&self,
gh: usize,
gw: usize,
dev: &Device,
dtype: DType,
) -> candle_core::Result<(Tensor, Tensor)> {
let half = self.inv_freq.len();
let n = gh * gw;
let mut data = Vec::with_capacity(n * 2 * half);
for hi in 0..gh {
for wi in 0..gw {
for &f in &self.inv_freq {
data.push(hi as f32 * f);
}
for &f in &self.inv_freq {
data.push(wi as f32 * f);
}
}
}
let freqs = Tensor::from_vec(data, (n, 2 * half), dev)?;
let cos = freqs.cos()?.to_dtype(dtype)?;
let sin = freqs.sin()?.to_dtype(dtype)?;
Ok((cos, sin))
}
}
/// The vision tower itself.
pub struct VisionTower {
/// Sum-collapsed temporal kernel (Conv2d, see module doc).
patch_embed: Conv2d,
pos_embed: Embedding,
rotary: VisionRotaryEmbedding,
blocks: Vec<VisionBlock>,
merger: VisionMerger,
config: VisionConfig,
@@ -265,6 +358,7 @@ impl VisionTower {
.get((cfg.num_position_embeddings, cfg.hidden_size), "weight")
.context("load model.visual.pos_embed.weight")?;
let pos_embed = Embedding::new(pos_embed_weight, cfg.hidden_size);
let rotary = VisionRotaryEmbedding::new(cfg.hidden_size / cfg.num_heads);
let blocks_vb = vb.pp("blocks");
let mut blocks = Vec::with_capacity(cfg.depth);
@@ -279,6 +373,7 @@ impl VisionTower {
Ok(Self {
patch_embed,
pos_embed,
rotary,
blocks,
merger,
config: cfg,
@@ -302,6 +397,89 @@ impl VisionTower {
gh * gw * LM_TOKENS_PER_MERGE_GROUP
}
/// Bilinearly interpolate the learned `pos_embed` grid (a
/// `num_grid_per_side × num_grid_per_side` table, 48×48 for Qwen3.6)
/// onto the actual `gh × gw` patch grid, in **row-major** patch
/// order. Port of the HF `fast_pos_embed_interpolate`: for each patch
/// at fractional grid coord `(linspace(0, ngrid-1, gh)[hi],
/// linspace(0, ngrid-1, gw)[wi])`, blend the 4 surrounding grid
/// entries by bilinear weights. Returns `(gh*gw, hidden)` in
/// `self.dtype`.
fn interpolated_pos_embed(&self, gh: usize, gw: usize) -> Result<Tensor> {
let ngrid = (self.config.num_position_embeddings as f64).sqrt().round() as usize;
anyhow::ensure!(
ngrid * ngrid == self.config.num_position_embeddings,
"num_position_embeddings {} is not a perfect square",
self.config.num_position_embeddings
);
// Evenly-spaced fractional indices into the [0, ngrid-1] grid.
let lin = |n: usize| -> Vec<f64> {
if n <= 1 {
vec![0.0]
} else {
let step = (ngrid - 1) as f64 / (n - 1) as f64;
(0..n).map(|i| i as f64 * step).collect()
}
};
let hs = lin(gh);
let ws = lin(gw);
let n = gh * gw;
// Four corner index sets + bilinear weight sets, row-major.
let mut idx: [Vec<u32>; 4] = [
Vec::with_capacity(n),
Vec::with_capacity(n),
Vec::with_capacity(n),
Vec::with_capacity(n),
];
let mut wts: [Vec<f32>; 4] = [
Vec::with_capacity(n),
Vec::with_capacity(n),
Vec::with_capacity(n),
Vec::with_capacity(n),
];
for &hv in &hs {
let hf = hv as usize; // floor (hv >= 0)
let hc = (hf + 1).min(ngrid - 1);
let dh = (hv - hf as f64) as f32;
for &wv in &ws {
let wf = wv as usize;
let wc = (wf + 1).min(ngrid - 1);
let dw = (wv - wf as f64) as f32;
idx[0].push((hf * ngrid + wf) as u32);
wts[0].push((1.0 - dh) * (1.0 - dw));
idx[1].push((hf * ngrid + wc) as u32);
wts[1].push((1.0 - dh) * dw);
idx[2].push((hc * ngrid + wf) as u32);
wts[2].push(dh * (1.0 - dw));
idx[3].push((hc * ngrid + wc) as u32);
wts[3].push(dh * dw);
}
}
// Blend in f32 and cast once at the end — the reference keeps
// the bilinear weights f32 against bf16 table rows; rounding
// the weights to bf16 first costs a visible slice of fixture
// parity (#15).
let mut acc: Option<Tensor> = None;
for corner in 0..4 {
let idx_t = Tensor::from_vec(std::mem::take(&mut idx[corner]), (n,), &self.device)?;
let emb = self
.pos_embed
.forward(&idx_t)?
.to_dtype(candle_core::DType::F32)?; // (n, hidden)
let wt = Tensor::from_vec(std::mem::take(&mut wts[corner]), (n, 1), &self.device)?;
let term = emb.broadcast_mul(&wt)?;
acc = Some(match acc {
Some(a) => a.add(&term)?,
None => term,
});
}
acc.expect("4 corners accumulated")
.to_dtype(self.dtype)
.map_err(Into::into)
}
/// Encode one image.
///
/// `image`: row-major `(3, H, W)` f32 tensor on `self.device`,
@@ -339,16 +517,34 @@ impl VisionTower {
let x = x.permute((1, 2, 0))?.contiguous()?;
let x = x.reshape((n_patches, self.config.hidden_size))?;
// Add learned positional embeddings (sequential indices for
// Stage A's fixed-resolution path; full 2D positional logic
// lands with variable resolution, issue #14).
// Learned absolute position embeddings. The `pos_embed` table is
// a `num_position_embeddings = num_grid_per_side²` learned grid
// (48×48 for Qwen3.6); for a `gh×gw` patch grid the reference
// (`fast_pos_embed_interpolate`) bilinearly interpolates that
// grid to `gh×gw`. The legacy path (a naive sequential lookup of
// the first `n_patches` rows) mis-maps the grid stride and
// scrambles spatial structure — kept only behind
// `NEURON_VISION_LEGACY_POS=1` for A/B comparison.
let pos = if vision_legacy_pos() {
let positions = Tensor::arange(0u32, n_patches as u32, &self.device)?;
let pos = self.pos_embed.forward(&positions)?;
self.pos_embed.forward(&positions)?
} else {
self.interpolated_pos_embed(gh, gw)?
};
let mut x = x.add(&pos)?;
// 2D vision rotary (row/col per patch), computed once and applied
// in every block's attention. Legacy escape hatch skips it.
let rotary = if vision_legacy_rope() {
None
} else {
Some(self.rotary.cos_sin(gh, gw, &self.device, self.dtype)?)
};
let rotary_ref = rotary.as_ref();
for (i, block) in self.blocks.iter().enumerate() {
x = block
.forward(&x)
.forward(&x, rotary_ref)
.with_context(|| format!("vision block {i}"))?;
}
@@ -516,9 +712,11 @@ mod tests {
spatial_merge_size: cfg.spatial_merge_size,
};
let rotary = VisionRotaryEmbedding::new(cfg.hidden_size / cfg.num_heads);
VisionTower {
patch_embed,
pos_embed,
rotary,
blocks,
merger,
config: cfg.clone(),
@@ -548,6 +746,51 @@ mod tests {
);
}
#[test]
fn interpolated_pos_embed_reduces_to_sequential_at_native_grid() {
// When the patch grid equals the pos_embed grid (gh=gw=ngrid),
// linspace(0,ngrid-1,ngrid) is the integer ladder, so every patch
// lands exactly on a grid node (dh=dw=0, corner-0 weight 1) and
// the bilinear result is the raw pos_embed rows in row-major
// order — i.e. identical to the legacy sequential lookup.
let cfg = tiny_config();
let tower = tiny_tower(&cfg);
let ngrid = (cfg.num_position_embeddings as f64).sqrt() as usize; // 8
let interp = tower.interpolated_pos_embed(ngrid, ngrid).unwrap();
let seq = tower
.pos_embed
.forward(&Tensor::arange(0u32, (ngrid * ngrid) as u32, &Device::Cpu).unwrap())
.unwrap();
let a: Vec<f32> = interp.flatten_all().unwrap().to_vec1().unwrap();
let b: Vec<f32> = seq.flatten_all().unwrap().to_vec1().unwrap();
assert_eq!(a.len(), b.len());
for (x, y) in a.iter().zip(b.iter()) {
assert!((x - y).abs() < 1e-5, "interp {x} vs seq {y}");
}
}
#[test]
fn vision_rotary_row_col_structure() {
// head_dim 8 → rotary dim 4 → inv_freq over [0,2] → 2 freqs/axis.
let rot = VisionRotaryEmbedding::new(8);
assert_eq!(rot.inv_freq.len(), 2);
let (cos, sin) = rot.cos_sin(2, 2, &Device::Cpu, DType::F32).unwrap();
assert_eq!(cos.dims(), &[4, 4]); // 4 patches, head_dim/2 = 4 cols
// Patch (0,0): all freqs 0 → cos 1, sin 0.
let s0: Vec<f32> = sin.i(0).unwrap().to_vec1().unwrap();
assert!(s0.iter().all(|&s| s.abs() < 1e-6));
// Patch index 2 = grid (1,0): row=1 drives the first half, col=0
// leaves the second half at zero.
let s2: Vec<f32> = sin.i(2).unwrap().to_vec1().unwrap();
assert!(s2[0].abs() > 1e-6, "row half must be non-zero");
assert!(
s2[2].abs() < 1e-6 && s2[3].abs() < 1e-6,
"col half must be zero"
);
}
#[test]
fn lm_token_count_matches_grid() {
let cfg = tiny_config();

File diff suppressed because it is too large Load Diff

View File

@@ -13,10 +13,11 @@
//! ARCH model state in this state slab will gain a companion
//! `tp_models: HashMap<TpHandle, Box<TpLeaderModel>>`.
use crate::harness::arch::qwen3_5::snapshot::KvCacheSnapshot;
use crate::harness::candle::ModelArch;
#[cfg(feature = "cuda")]
use crate::harness::device_worker::jobs::TpHandle;
use crate::harness::device_worker::jobs::{ArchHandle, ImageInput, Job};
use crate::harness::device_worker::jobs::{ArchHandle, ImageInput, Job, KvSnapshotId};
#[cfg(feature = "cuda")]
use crate::harness::tp::TpLeaderModel;
use crate::harness::tp::nccl_state::NcclState;
@@ -46,6 +47,14 @@ struct DeviceWorkerState {
/// increments and returns the new value. Wraps at u64::MAX after
/// ~10^19 model loads — not a practical concern.
next_handle: u64,
/// Prefix-cache snapshots (#11), keyed by the owning model's
/// handle plus a per-worker snapshot counter. Kept beside the
/// model slab (not inside it) so every existing `get_mut` on
/// `models` stays untouched; `DropArch` retains this map down so
/// snapshot tensors drop on this thread alongside the model's.
kv_snapshots: HashMap<(ArchHandle, u64), KvCacheSnapshot>,
/// Counter for minting fresh `KvSnapshotId`s.
next_kv_snapshot_id: u64,
/// Leader's NCCL state. Populated by `Job::NcclInit`; the
/// underlying `Comm`'s libnccl handle lives bound to this thread
/// for its entire lifetime. Subprocess workers maintain their own
@@ -60,6 +69,12 @@ struct DeviceWorkerState {
/// Counter for minting fresh `TpHandle`s.
#[cfg(feature = "cuda")]
next_tp_handle: u64,
/// Leader-side TP prefix snapshots (#11), keyed by the owning TP
/// handle plus the **pool-minted** snapshot id (no local counter —
/// the id must match what the subprocess ranks stored). `DropTp`
/// retains this map down with the model.
#[cfg(feature = "cuda")]
tp_kv_snapshots: HashMap<(TpHandle, u64), KvCacheSnapshot>,
#[cfg(feature = "cuda")]
#[allow(dead_code)]
/// `None` only if `CudaContext::new()` failed — in that case the
@@ -124,6 +139,10 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
Job::DropArch { handle, reply } => {
let removed = state.models.remove(&handle);
let was_present = removed.is_some();
// Prefix snapshots are scoped to the model: drop them
// here (on this thread) so a stale async-side id can
// never resurrect tensors from an unloaded model.
state.kv_snapshots.retain(|(h, _), _| *h != handle);
// Explicit drop on this thread — runs the Box<ModelArch>
// Drop with the CUDA context bound here, which frees
// all device tensors on the right context. The Drop is
@@ -150,6 +169,76 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
}
let _ = reply.send(result);
}
Job::SnapshotKv { handle, reply } => {
let result = match state.models.get(&handle) {
Some(arch) => arch.snapshot_kv_cache().map(|snap| {
let id = KvSnapshotId(state.next_kv_snapshot_id);
state.next_kv_snapshot_id = state.next_kv_snapshot_id.wrapping_add(1);
let bytes = snap.size_bytes();
state.kv_snapshots.insert((handle, id.0), snap);
tracing::debug!(
device_index,
handle = handle.0,
snapshot = id.0,
bytes,
stored = state.kv_snapshots.len(),
"device worker: kv snapshot captured"
);
(id, bytes)
}),
None => Err(anyhow::anyhow!(
"SnapshotKv: no model for handle {}",
handle.0
)),
};
let _ = reply.send(result);
}
Job::RestoreKv {
handle,
snapshot,
reply,
} => {
let result = match (
state.models.get_mut(&handle),
state.kv_snapshots.get(&(handle, snapshot.0)),
) {
(Some(arch), Some(snap)) => arch.restore_kv_cache(snap),
(None, _) => Err(anyhow::anyhow!(
"RestoreKv: no model for handle {}",
handle.0
)),
(_, None) => Err(anyhow::anyhow!(
"RestoreKv: no snapshot {} for handle {}",
snapshot.0,
handle.0
)),
};
// The replaced live cache state just freed its
// tensors — same release-to-driver point as ClearKv.
if result.is_ok() {
trim_device_pool(&state);
}
let _ = reply.send(result);
}
Job::DropKvSnapshot {
handle,
snapshot,
reply,
} => {
let was_present = state.kv_snapshots.remove(&(handle, snapshot.0)).is_some();
if was_present {
trim_device_pool(&state);
}
tracing::debug!(
device_index,
handle = handle.0,
snapshot = snapshot.0,
was_present,
stored = state.kv_snapshots.len(),
"device worker: kv snapshot dropped"
);
let _ = reply.send(());
}
Job::ForwardLogits {
handle,
tokens,
@@ -201,6 +290,16 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
let _ = reply.send(resp);
}
#[cfg(feature = "cuda")]
Job::GetLeaderComm { reply } => {
// Clone the leader's Arc<Comm> out for the async-side
// watchdog. `None` before NcclInit. (#17 Stage 2)
let comm = state
.nccl
.comm()
.map(crate::harness::tp::nccl_state::SendComm);
let _ = reply.send(comm);
}
#[cfg(feature = "cuda")]
Job::TpLoadShard {
model_id,
config_json,
@@ -226,6 +325,7 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
let removed = state.tp_models.remove(&handle);
let was_present = removed.is_some();
drop(removed);
state.tp_kv_snapshots.retain(|(h, _), _| *h != handle);
tracing::debug!(
device_index,
tp_handle = handle.0,
@@ -253,6 +353,89 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
let _ = reply.send(result);
}
#[cfg(feature = "cuda")]
Job::TpSnapshotKv {
handle,
snapshot_id,
reply,
} => {
let result = match state.tp_models.get(&handle) {
Some(model) => {
model
.snapshot_kv_cache()
.map_err(anyhow::Error::from)
.map(|snap| {
let bytes = snap.size_bytes();
state.tp_kv_snapshots.insert((handle, snapshot_id), snap);
tracing::debug!(
device_index,
tp_handle = handle.0,
snapshot_id,
bytes,
stored = state.tp_kv_snapshots.len(),
"device worker: TP kv snapshot captured"
);
bytes
})
}
None => Err(anyhow::anyhow!(
"TpSnapshotKv: no TP model for handle {}",
handle.0
)),
};
let _ = reply.send(result);
}
#[cfg(feature = "cuda")]
Job::TpRestoreKv {
handle,
snapshot_id,
reply,
} => {
let result = match (
state.tp_models.get_mut(&handle),
state.tp_kv_snapshots.get(&(handle, snapshot_id)),
) {
(Some(model), Some(snap)) => {
model.restore_kv_cache(snap).map_err(anyhow::Error::from)
}
(None, _) => Err(anyhow::anyhow!(
"TpRestoreKv: no TP model for handle {}",
handle.0
)),
(_, None) => Err(anyhow::anyhow!(
"TpRestoreKv: no snapshot {} for handle {}",
snapshot_id,
handle.0
)),
};
if result.is_ok() {
trim_device_pool(&state);
}
let _ = reply.send(result);
}
#[cfg(feature = "cuda")]
Job::TpDropKvSnapshot {
handle,
snapshot_id,
reply,
} => {
let was_present = state
.tp_kv_snapshots
.remove(&(handle, snapshot_id))
.is_some();
if was_present {
trim_device_pool(&state);
}
tracing::debug!(
device_index,
tp_handle = handle.0,
snapshot_id,
was_present,
stored = state.tp_kv_snapshots.len(),
"device worker: TP kv snapshot dropped"
);
let _ = reply.send(());
}
#[cfg(feature = "cuda")]
Job::TpForwardLogits {
handle,
tokens,
@@ -269,6 +452,7 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
offset,
image_token_id,
image_data_uris,
chunk_size,
reply,
} => {
let result = tp_forward_logits_with_images(
@@ -278,6 +462,7 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
offset,
image_token_id,
&image_data_uris,
chunk_size,
);
let _ = reply.send(result);
}
@@ -351,9 +536,12 @@ fn init_state(device_index: u32) -> DeviceWorkerState {
device,
models: HashMap::new(),
next_handle: 1,
kv_snapshots: HashMap::new(),
next_kv_snapshot_id: 1,
nccl: NcclState::new(),
tp_models: HashMap::new(),
next_tp_handle: 1,
tp_kv_snapshots: HashMap::new(),
ctx,
}
}
@@ -364,6 +552,8 @@ fn init_state(device_index: u32) -> DeviceWorkerState {
device: candle_core::Device::Cpu,
models: HashMap::new(),
next_handle: 1,
kv_snapshots: HashMap::new(),
next_kv_snapshot_id: 1,
nccl: NcclState::new(),
}
}
@@ -768,6 +958,7 @@ fn tp_forward_logits_with_images(
offset: usize,
image_token_id: u32,
image_data_uris: &[String],
chunk_size: usize,
) -> anyhow::Result<Vec<f32>> {
use crate::harness::preprocess::{PreprocessProfile, preprocess_data_uri};
use candle_core::{DType, Tensor};
@@ -776,24 +967,20 @@ fn tp_forward_logits_with_images(
anyhow::bail!("TpForwardLogitsWithImages dispatched with zero images");
}
// Preprocess every image into a device-resident (C, H, W) tensor.
// Same fixed-resolution profile + decode path the subprocess workers
// run, so the encoded embeddings match across ranks bit-for-bit.
// Preprocess every image into a device-resident (C, H, W) tensor at
// its native-aspect resized dims (#14). Same `smart_resize` + decode
// path the subprocess workers run, so the encoded embeddings — and
// the per-image grids derived from these dims — match across ranks
// bit-for-bit.
let profile = PreprocessProfile::qwen3_6();
let (h, w) = (
profile.target_height as usize,
profile.target_width as usize,
);
let mut pixels: Vec<Tensor> = Vec::with_capacity(image_data_uris.len());
for (idx, uri) in image_data_uris.iter().enumerate() {
let px = preprocess_data_uri(uri, &profile)
let (px, h, w) = preprocess_data_uri(uri, &profile)
.with_context(|| format!("preprocess image[{idx}] (TP leader)"))?;
let t = Tensor::from_vec(px, (3, h, w), &state.device)?;
let t = Tensor::from_vec(px, (3, h as usize, w as usize), &state.device)?;
pixels.push(t);
}
let input = Tensor::new(tokens, &state.device)?.unsqueeze(0)?;
let model = state.tp_models.get_mut(&handle).ok_or_else(|| {
anyhow::anyhow!(
"TpForwardLogitsWithImages: no model for handle {}",
@@ -801,7 +988,10 @@ fn tp_forward_logits_with_images(
)
})?;
let logits = model.forward_with_images(&input, offset, &pixels, image_token_id)?;
// Chunked prefill (encode once, splice per chunk) — bounded
// activation, in lockstep with the subprocess ranks.
let logits =
model.prefill_with_images_chunked(tokens, offset, &pixels, image_token_id, chunk_size)?;
let logits = logits.squeeze(0)?.squeeze(0)?;
let logits = logits.to_dtype(DType::F32)?.flatten_all()?;
let values = logits.to_vec1::<f32>()?;
@@ -869,13 +1059,10 @@ fn forward_logits_with_images(
anyhow::bail!("ForwardLogitsWithImages dispatched with zero images");
}
let arch = state.models.get_mut(&handle).ok_or_else(|| {
anyhow::anyhow!("ForwardLogitsWithImages: no model for handle {}", handle.0)
})?;
// Encode every image on the worker's device, collecting per-image
// post-merger embeddings as device-resident tensors.
let mut per_image: Vec<Tensor> = Vec::with_capacity(images.len());
// Reconstruct the preprocessed pixels into device-resident
// `(C, H, W)` tensors first (immutable `state.device` borrow), then
// take the `&mut` model borrow for the chunked prefill below.
let mut image_pixels: Vec<Tensor> = Vec::with_capacity(images.len());
for (idx, img) in images.into_iter().enumerate() {
anyhow::ensure!(
img.pixels.len() == img.c * img.h * img.w,
@@ -885,19 +1072,26 @@ fn forward_logits_with_images(
img.h,
img.w,
);
let image = Tensor::from_vec(img.pixels, (img.c, img.h, img.w), &state.device)?;
let embed = arch
.encode_image(&image)
.with_context(|| format!("encode image[{idx}]"))?;
per_image.push(embed);
image_pixels.push(Tensor::from_vec(
img.pixels,
(img.c, img.h, img.w),
&state.device,
)?);
}
// Concatenate per-image embeddings along the patch axis →
// (sum_of_patches, hidden). `Tensor::cat` keeps the result
// device-resident.
let image_embeds = Tensor::cat(&per_image.iter().collect::<Vec<_>>(), 0)?;
let input = Tensor::new(tokens, &state.device)?.unsqueeze(0)?;
let logits = arch.forward_with_vision(&input, offset, &image_embeds, image_token_id)?;
let chunk_size = crate::harness::candle::prefill_chunk_tokens();
let arch = state.models.get_mut(&handle).ok_or_else(|| {
anyhow::anyhow!("ForwardLogitsWithImages: no model for handle {}", handle.0)
})?;
// Chunked image prefill (#18): encode once, walk the prompt in
// `chunk_size` windows splicing per-chunk image-pad rows — parity
// with the TP path so a long single-GPU vision context serves
// instead of single-shot OOMing. Returns the final chunk's
// `[vocab]` logits.
let logits = arch
.prefill_with_images_chunked(tokens, offset, &image_pixels, image_token_id, chunk_size)
.context("chunked vision prefill")?;
let values = logits
.to_dtype(DType::F32)?
.flatten_all()?
@@ -978,6 +1172,18 @@ fn drain_poisoned(job: Job, device_index: u32) {
Job::ClearKv { reply, .. } => {
let _ = reply.send(Err(err()));
}
Job::SnapshotKv { reply, .. } => {
let _ = reply.send(Err(err()));
}
Job::RestoreKv { reply, .. } => {
let _ = reply.send(Err(err()));
}
Job::DropKvSnapshot { reply, .. } => {
// Same shape as DropArch: unit reply so the caller's await
// resolves; the snapshot leaks with the rest of the slab
// per the poisoned-thread design.
let _ = reply.send(());
}
Job::ForwardLogits { reply, .. } => {
let _ = reply.send(Err(err()));
}
@@ -993,6 +1199,10 @@ fn drain_poisoned(job: Job, device_index: u32) {
message: format!("device worker {device_index} poisoned"),
});
}
#[cfg(feature = "cuda")]
Job::GetLeaderComm { reply } => {
let _ = reply.send(None);
}
Job::NcclSanity { reply } => {
let _ = reply.send(crate::harness::tp::rpc::WorkerResponse::Error {
kind: "device_worker_poisoned".into(),
@@ -1012,6 +1222,20 @@ fn drain_poisoned(job: Job, device_index: u32) {
let _ = reply.send(Err(err()));
}
#[cfg(feature = "cuda")]
Job::TpSnapshotKv { reply, .. } => {
let _ = reply.send(Err(err()));
}
#[cfg(feature = "cuda")]
Job::TpRestoreKv { reply, .. } => {
let _ = reply.send(Err(err()));
}
#[cfg(feature = "cuda")]
Job::TpDropKvSnapshot { reply, .. } => {
// Bookkeeping-only — unit reply so eviction never wedges
// on a poisoned worker (same shape as DropKvSnapshot).
let _ = reply.send(());
}
#[cfg(feature = "cuda")]
Job::TpForwardLogits { reply, .. } => {
let _ = reply.send(Err(err()));
}

View File

@@ -28,6 +28,14 @@ pub struct ArchHandle(pub u64);
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct TpHandle(pub u64);
/// Opaque handle to a prefix-cache snapshot (#11) stored worker-side
/// next to the model slab. Scoped to the `ArchHandle` it was captured
/// from — `Job::DropArch` drops every snapshot under its handle. The
/// snapshot's tensors never leave the worker thread; the async side
/// holds only this id plus the token sequence it covers.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct KvSnapshotId(pub u64);
/// One image payload for `Job::ForwardLogitsWithImages` /
/// `Job::EncodeImage`. Pixels are row-major `(c, h, w)` f32 — the
/// shape `harness::preprocess::preprocess` produces. Carries the
@@ -36,8 +44,13 @@ pub struct TpHandle(pub u64);
/// `Clone` so the vision-aware dispatch in `chat_completion` can
/// match `&vision_route` (carrying borrowed images) and still hand
/// owned `Vec<ImageInput>` to the worker job. The clone cost is one
/// pixel-buffer memcpy per image — fine at fixed-resolution sizes
/// (3 × 448 × 448 × 4 bytes = ~2.4 MiB per image).
/// pixel-buffer memcpy per image — now variable with dynamic resolution
/// (#14): `3 × h × w × 4` bytes, up to ~6.3 MiB at the default 1024²
/// `max_pixels` budget.
///
/// `h`/`w` are the **resized** dims (factor-aligned), so the per-image LM
/// grid is `(h/factor, w/factor)` — derived downstream for the splice
/// and the interleaved-M-RoPE position ids.
#[derive(Clone)]
pub struct ImageInput {
pub pixels: Vec<f32>,
@@ -100,6 +113,30 @@ pub enum Job {
handle: ArchHandle,
reply: oneshot::Sender<Result<()>>,
},
/// Capture the model's live cache state (attention KV + GDN
/// recurrent state + position counters) as a prefix snapshot
/// (#11). The snapshot stays in the worker's state, keyed by the
/// returned id; the reply carries `(id, bytes)` so the async side
/// can do budget accounting without touching tensors. Errors on
/// archs without snapshot support.
SnapshotKv {
handle: ArchHandle,
reply: oneshot::Sender<Result<(KvSnapshotId, u64)>>,
},
/// Replace the model's live cache state with a stored snapshot,
/// instead of `ClearKv`, so prefill can resume at the snapshot's
/// token boundary. The snapshot remains stored (restorable again).
RestoreKv {
handle: ArchHandle,
snapshot: KvSnapshotId,
reply: oneshot::Sender<Result<()>>,
},
/// Drop one stored snapshot (prefix-cache eviction). Idempotent.
DropKvSnapshot {
handle: ArchHandle,
snapshot: KvSnapshotId,
reply: oneshot::Sender<()>,
},
/// Run one forward step and copy the resulting `[vocab]` logits to
/// CPU. The caller takes the returned `Vec<f32>`, wraps it in a
/// CPU `Tensor`, and runs `apply_repeat_penalty` + sampling
@@ -187,6 +224,17 @@ pub enum Job {
NcclSanity {
reply: oneshot::Sender<crate::harness::tp::rpc::WorkerResponse>,
},
/// Hand a clonable handle to the leader's NCCL `Comm` back to the
/// async side, so the TP step watchdog can call `ncclCommAbort` on
/// it from a *different* thread to unblock a wedged collective
/// (#17 Stage 2). Fetched once at init while the worker thread is
/// still responsive — a thread already wedged in a collective can't
/// service this job, which is exactly why the handle is cached
/// up front. Replies `None` before `NcclInit` has run.
#[cfg(feature = "cuda")]
GetLeaderComm {
reply: oneshot::Sender<Option<crate::harness::tp::nccl_state::SendComm>>,
},
/// Load the leader's TP shard on the worker thread. The dispatch
/// handler reads `state.nccl.comm()` directly (no cross-thread
/// `Arc<Comm>` transfer, no `SendComm` wrapper) and builds the
@@ -219,6 +267,31 @@ pub enum Job {
handle: TpHandle,
reply: oneshot::Sender<Result<()>>,
},
/// Capture the leader's TP cache state as a prefix snapshot (#11),
/// stored worker-side under the pool-minted `snapshot_id` (shared
/// with the subprocess ranks, so all ranks key the same snapshot
/// identically). Replies with the leader shard's snapshot bytes.
#[cfg(feature = "cuda")]
TpSnapshotKv {
handle: TpHandle,
snapshot_id: u64,
reply: oneshot::Sender<Result<u64>>,
},
/// Replace the leader's live TP cache state with a stored
/// snapshot. Mirrors `RestoreKv` for single-GPU.
#[cfg(feature = "cuda")]
TpRestoreKv {
handle: TpHandle,
snapshot_id: u64,
reply: oneshot::Sender<Result<()>>,
},
/// Drop one stored leader TP snapshot (eviction). Idempotent.
#[cfg(feature = "cuda")]
TpDropKvSnapshot {
handle: TpHandle,
snapshot_id: u64,
reply: oneshot::Sender<()>,
},
/// Run one TP forward step on the leader's shard. Returns CPU-
/// side logits as a `Vec<f32>` so the async caller can sample
/// without holding a device tensor. The caller is also
@@ -246,6 +319,7 @@ pub enum Job {
offset: usize,
image_token_id: u32,
image_data_uris: Vec<String>,
chunk_size: usize,
reply: oneshot::Sender<Result<Vec<f32>>>,
},
/// Tell the worker to break its dispatch loop and exit. Any jobs

View File

@@ -51,7 +51,7 @@ use tokio::sync::oneshot;
#[cfg(feature = "cuda")]
pub use jobs::TpHandle;
pub use jobs::{ArchHandle, Job};
pub use jobs::{ArchHandle, Job, KvSnapshotId};
/// Errors returned by `DeviceWorkerHandle` submit methods.
#[derive(Debug, thiserror::Error)]
@@ -161,6 +161,27 @@ impl DeviceWorkerHandle {
}
}
/// Fetch a clonable handle to the leader's NCCL `Comm` (#17 Stage 2).
/// The TP step watchdog caches this at init so it can call
/// `ncclCommAbort` from the async thread to unblock a wedged
/// collective. Returns `None` if uninitialised, poisoned, or gone —
/// the caller treats a missing handle as "can't abort" and logs it.
#[cfg(feature = "cuda")]
pub async fn get_leader_comm(&self) -> Option<crate::harness::tp::nccl_state::SendComm> {
if self.poisoned.load(Ordering::Acquire) {
return None;
}
let (reply_tx, reply_rx) = oneshot::channel();
if self
.tx
.send(Job::GetLeaderComm { reply: reply_tx })
.is_err()
{
return None;
}
reply_rx.await.ok().flatten()
}
/// Load a GGUF (pre-quantized) single-GPU model on the worker
/// thread. The hf-hub resolution happens on the async caller; the
/// resolved local `gguf_path` plus the spec's model_id are sent
@@ -279,6 +300,92 @@ impl DeviceWorkerHandle {
}
}
/// Capture the model's live cache state as a worker-side prefix
/// snapshot (#11). Returns the snapshot id plus its byte size for
/// the async-side budget accounting. Tensors stay on the worker.
pub async fn snapshot_kv(
&self,
handle: ArchHandle,
) -> Result<(jobs::KvSnapshotId, u64), WorkerError> {
if self.poisoned.load(Ordering::Acquire) {
return Err(WorkerError::Poisoned {
device_index: self.device_index,
});
}
let (reply_tx, reply_rx) = oneshot::channel();
self.tx
.send(Job::SnapshotKv {
handle,
reply: reply_tx,
})
.map_err(|_| WorkerError::Gone {
device_index: self.device_index,
})?;
match reply_rx.await {
Ok(result) => result.map_err(WorkerError::from),
Err(_) => Err(WorkerError::Gone {
device_index: self.device_index,
}),
}
}
/// Replace the model's live cache state with a stored snapshot —
/// called instead of [`Self::clear_kv_cache`] on a prefix-cache
/// hit. The snapshot remains stored and restorable again.
pub async fn restore_kv(
&self,
handle: ArchHandle,
snapshot: jobs::KvSnapshotId,
) -> Result<(), WorkerError> {
if self.poisoned.load(Ordering::Acquire) {
return Err(WorkerError::Poisoned {
device_index: self.device_index,
});
}
let (reply_tx, reply_rx) = oneshot::channel();
self.tx
.send(Job::RestoreKv {
handle,
snapshot,
reply: reply_tx,
})
.map_err(|_| WorkerError::Gone {
device_index: self.device_index,
})?;
match reply_rx.await {
Ok(result) => result.map_err(WorkerError::from),
Err(_) => Err(WorkerError::Gone {
device_index: self.device_index,
}),
}
}
/// Drop one stored prefix snapshot (eviction). Mirrors
/// [`Self::drop_arch`]'s poison-tolerant unit-reply shape so
/// bookkeeping always unblocks.
pub async fn drop_kv_snapshot(
&self,
handle: ArchHandle,
snapshot: jobs::KvSnapshotId,
) -> Result<(), WorkerError> {
let (reply_tx, reply_rx) = oneshot::channel();
self.tx
.send(Job::DropKvSnapshot {
handle,
snapshot,
reply: reply_tx,
})
.map_err(|_| WorkerError::Gone {
device_index: self.device_index,
})?;
match reply_rx.await {
Ok(()) => Ok(()),
Err(_) => Err(WorkerError::Gone {
device_index: self.device_index,
}),
}
}
/// Run one forward step and return the resulting `[vocab]` logits
/// as a CPU-side `Vec<f32>`. The caller then samples on a CPU
/// candle Tensor without ever binding the device context on its
@@ -537,6 +644,96 @@ impl DeviceWorkerHandle {
}
}
/// Capture the leader's TP cache state as a prefix snapshot (#11)
/// stored under the pool-minted `snapshot_id`. Returns the leader
/// shard's snapshot bytes.
#[cfg(feature = "cuda")]
pub async fn tp_snapshot_kv(
&self,
handle: TpHandle,
snapshot_id: u64,
) -> Result<u64, WorkerError> {
if self.poisoned.load(Ordering::Acquire) {
return Err(WorkerError::Poisoned {
device_index: self.device_index,
});
}
let (reply_tx, reply_rx) = oneshot::channel();
self.tx
.send(Job::TpSnapshotKv {
handle,
snapshot_id,
reply: reply_tx,
})
.map_err(|_| WorkerError::Gone {
device_index: self.device_index,
})?;
match reply_rx.await {
Ok(result) => result.map_err(WorkerError::from),
Err(_) => Err(WorkerError::Gone {
device_index: self.device_index,
}),
}
}
/// Replace the leader's live TP cache state with a stored
/// snapshot — called instead of [`Self::tp_clear_kv`] on a
/// prefix-cache hit.
#[cfg(feature = "cuda")]
pub async fn tp_restore_kv(
&self,
handle: TpHandle,
snapshot_id: u64,
) -> Result<(), WorkerError> {
if self.poisoned.load(Ordering::Acquire) {
return Err(WorkerError::Poisoned {
device_index: self.device_index,
});
}
let (reply_tx, reply_rx) = oneshot::channel();
self.tx
.send(Job::TpRestoreKv {
handle,
snapshot_id,
reply: reply_tx,
})
.map_err(|_| WorkerError::Gone {
device_index: self.device_index,
})?;
match reply_rx.await {
Ok(result) => result.map_err(WorkerError::from),
Err(_) => Err(WorkerError::Gone {
device_index: self.device_index,
}),
}
}
/// Drop one stored leader TP snapshot (eviction). Poison-tolerant
/// unit reply, same shape as [`Self::drop_kv_snapshot`].
#[cfg(feature = "cuda")]
pub async fn tp_drop_kv_snapshot(
&self,
handle: TpHandle,
snapshot_id: u64,
) -> Result<(), WorkerError> {
let (reply_tx, reply_rx) = oneshot::channel();
self.tx
.send(Job::TpDropKvSnapshot {
handle,
snapshot_id,
reply: reply_tx,
})
.map_err(|_| WorkerError::Gone {
device_index: self.device_index,
})?;
match reply_rx.await {
Ok(()) => Ok(()),
Err(_) => Err(WorkerError::Gone {
device_index: self.device_index,
}),
}
}
/// Run one TP forward step on the leader's shard. Returns CPU-side
/// logits as `Vec<f32>` ready for sampling. The caller is
/// responsible for fan-out / drain of the subprocess workers
@@ -579,6 +776,7 @@ impl DeviceWorkerHandle {
/// matching `GenerateStepWithImages` out to subprocess ranks so the
/// row-parallel collectives complete.
#[cfg(feature = "cuda")]
#[allow(clippy::too_many_arguments)]
pub async fn tp_forward_logits_with_images(
&self,
handle: TpHandle,
@@ -586,6 +784,7 @@ impl DeviceWorkerHandle {
offset: usize,
image_token_id: u32,
image_data_uris: Vec<String>,
chunk_size: usize,
) -> Result<Vec<f32>, WorkerError> {
if self.poisoned.load(Ordering::Acquire) {
return Err(WorkerError::Poisoned {
@@ -600,6 +799,7 @@ impl DeviceWorkerHandle {
offset,
image_token_id,
image_data_uris,
chunk_size,
reply: reply_tx,
})
.map_err(|_| WorkerError::Gone {

View File

@@ -4,8 +4,10 @@ pub mod arch;
pub mod candle;
pub mod chat_template;
pub mod device_worker;
pub mod prefix_cache;
pub mod preflight;
pub mod preprocess;
pub mod speculative;
pub mod tp;
use anyhow::Result;
@@ -114,10 +116,8 @@ impl HarnessRegistry {
for config in configs {
match config.name.as_str() {
"candle" => {
let harness = Arc::new(candle::CandleHarness::new(
bind_url.to_string(),
&settings.candle,
));
let harness =
candle::CandleHarness::new(bind_url.to_string(), &settings.candle);
registry.candle = Some(Arc::clone(&harness));
registry.harnesses.insert("candle".into(), harness);
}

View File

@@ -0,0 +1,266 @@
//! Prefix-cache registry: which cache snapshots exist for a loaded
//! model, which one matches an incoming prompt, and which to evict.
//!
//! Pure bookkeeping — no tensors live here. Each entry pairs the exact
//! token sequence a snapshot was captured at with an opaque snapshot
//! reference `R` (a worker-side snapshot id for CUDA loads, the
//! snapshot itself for CPU loads) and its byte size for the VRAM
//! budget. The caller owns actually dropping evicted snapshots.
//!
//! ## Matching policy
//!
//! A snapshot is reusable only when its **entire** token sequence is a
//! strict prefix of the incoming prompt (`entry.len() < prompt.len()`
//! — at least one suffix token must be forwarded to produce logits).
//! The GatedDeltaNet recurrent state cannot be rewound, so partial
//! matches are unusable; see `arch/qwen3_5/snapshot.rs`.
//!
//! ## Insertion policy
//!
//! Inserting an entry drops existing entries that are strict prefixes
//! of it: the append-only agent loop (turn N+1 = turn N + new text)
//! keeps exactly one entry per conversation thread that way, instead
//! of one per turn. Eviction beyond that is LRU over total bytes
//! against the configured budget, plus a max-entries cap.
/// One cached snapshot: the token sequence it was captured at, the
/// opaque snapshot reference, and bookkeeping for eviction.
struct Entry<R> {
tokens: Vec<u32>,
snapshot: R,
bytes: u64,
last_used: u64,
}
/// A match returned by [`PrefixCache::longest_match`].
pub struct PrefixMatch<R> {
/// Clone of the matched snapshot reference.
pub snapshot: R,
/// Number of prompt tokens the snapshot covers (the entry's full
/// token count). Prefill resumes at this offset.
pub tokens: usize,
}
/// LRU prefix-snapshot registry for one loaded model.
pub struct PrefixCache<R> {
entries: Vec<Entry<R>>,
budget_bytes: u64,
max_entries: usize,
/// Monotonic access clock for LRU ordering.
seq: u64,
}
impl<R: Clone> PrefixCache<R> {
pub fn new(budget_bytes: u64, max_entries: usize) -> Self {
Self {
entries: Vec::new(),
budget_bytes,
max_entries,
seq: 0,
}
}
fn tick(&mut self) -> u64 {
self.seq += 1;
self.seq
}
fn used_bytes(&self) -> u64 {
self.entries.iter().map(|e| e.bytes).sum()
}
/// Longest entry whose token sequence is a strict prefix of
/// `prompt`. Touches the entry's LRU clock on hit.
pub fn longest_match(&mut self, prompt: &[u32]) -> Option<PrefixMatch<R>> {
let idx = self
.entries
.iter()
.enumerate()
.filter(|(_, e)| e.tokens.len() < prompt.len() && prompt.starts_with(&e.tokens))
.max_by_key(|(_, e)| e.tokens.len())
.map(|(i, _)| i)?;
let now = self.tick();
let entry = &mut self.entries[idx];
entry.last_used = now;
Some(PrefixMatch {
snapshot: entry.snapshot.clone(),
tokens: entry.tokens.len(),
})
}
/// Remove the entry whose tokens exactly prefix-match what
/// `longest_match` just returned. Called when restoring its
/// snapshot failed; returns the reference so the caller can drop
/// the underlying snapshot.
pub fn remove_covering(&mut self, prompt: &[u32], tokens: usize) -> Option<R> {
let idx = self
.entries
.iter()
.position(|e| e.tokens.len() == tokens && prompt.starts_with(&e.tokens))?;
Some(self.entries.swap_remove(idx).snapshot)
}
/// Insert a fresh snapshot captured at exactly `tokens`. Returns
/// every snapshot reference the caller must now drop: replaced
/// duplicates, strict prefixes of the new entry, LRU evictions to
/// fit the byte budget / entry cap — and the new snapshot itself
/// when it alone exceeds the budget (in which case it is not
/// inserted).
pub fn insert(&mut self, tokens: Vec<u32>, snapshot: R, bytes: u64) -> Vec<R> {
let mut dropped = Vec::new();
if bytes > self.budget_bytes || self.max_entries == 0 || tokens.is_empty() {
dropped.push(snapshot);
return dropped;
}
// Drop entries the new one supersedes: exact duplicates and
// strict prefixes (the conversation they belong to has moved
// on; the new entry matches everything they would have).
let mut i = 0;
while i < self.entries.len() {
if tokens.starts_with(&self.entries[i].tokens) {
dropped.push(self.entries.swap_remove(i).snapshot);
} else {
i += 1;
}
}
let now = self.tick();
self.entries.push(Entry {
tokens,
snapshot,
bytes,
last_used: now,
});
// LRU-evict to budget and cap. The just-inserted entry has the
// freshest clock, so it is only evicted if it is the last one
// standing — and it fits the budget alone (checked above).
while self.used_bytes() > self.budget_bytes || self.entries.len() > self.max_entries {
let lru = self
.entries
.iter()
.enumerate()
.min_by_key(|(_, e)| e.last_used)
.map(|(i, _)| i)
.expect("eviction loop runs only while entries is non-empty");
dropped.push(self.entries.swap_remove(lru).snapshot);
}
dropped
}
/// Number of live entries (test/log helper).
pub fn len(&self) -> usize {
self.entries.len()
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn cache(budget: u64, max: usize) -> PrefixCache<u64> {
PrefixCache::new(budget, max)
}
#[test]
fn longest_strict_prefix_wins() {
let mut c = cache(1000, 8);
assert!(c.insert(vec![1, 2], 10, 1).is_empty());
// [1,2,3] is NOT a prefix of [1,2] superseding chain — diverge
// it so both stay live.
assert!(c.insert(vec![1, 9, 9, 9], 11, 1).is_empty());
let m = c.longest_match(&[1, 2, 3, 4]).expect("hit");
assert_eq!(m.snapshot, 10);
assert_eq!(m.tokens, 2);
}
#[test]
fn exact_length_match_is_rejected() {
// A snapshot covering the whole prompt leaves no suffix token
// to forward — must not match.
let mut c = cache(1000, 8);
c.insert(vec![1, 2, 3], 10, 1);
assert!(c.longest_match(&[1, 2, 3]).is_none());
assert!(c.longest_match(&[1, 2, 3, 4]).is_some());
}
#[test]
fn divergent_prompt_misses() {
let mut c = cache(1000, 8);
c.insert(vec![1, 2, 3], 10, 1);
assert!(c.longest_match(&[1, 2, 4, 5]).is_none());
}
#[test]
fn insert_supersedes_prefix_entries() {
let mut c = cache(1000, 8);
c.insert(vec![1, 2], 10, 1);
let dropped = c.insert(vec![1, 2, 3, 4], 11, 1);
assert_eq!(dropped, vec![10]);
assert_eq!(c.len(), 1);
// The longer entry still matches its own continuations.
assert_eq!(c.longest_match(&[1, 2, 3, 4, 5]).unwrap().snapshot, 11);
}
#[test]
fn insert_replaces_exact_duplicate() {
let mut c = cache(1000, 8);
c.insert(vec![1, 2], 10, 1);
let dropped = c.insert(vec![1, 2], 11, 1);
assert_eq!(dropped, vec![10]);
assert_eq!(c.len(), 1);
}
#[test]
fn byte_budget_evicts_lru() {
let mut c = cache(10, 8);
c.insert(vec![1], 10, 4);
c.insert(vec![2], 11, 4);
// Touch [1] so [2] becomes LRU.
assert!(c.longest_match(&[1, 5]).is_some());
let dropped = c.insert(vec![3], 12, 4);
assert_eq!(dropped, vec![11]);
assert_eq!(c.len(), 2);
assert!(c.longest_match(&[1, 5]).is_some());
assert!(c.longest_match(&[2, 5]).is_none());
}
#[test]
fn max_entries_cap_evicts_lru() {
let mut c = cache(1000, 2);
c.insert(vec![1], 10, 1);
c.insert(vec![2], 11, 1);
let dropped = c.insert(vec![3], 12, 1);
assert_eq!(dropped, vec![10]);
assert_eq!(c.len(), 2);
}
#[test]
fn oversized_snapshot_is_rejected_back() {
let mut c = cache(10, 8);
let dropped = c.insert(vec![1, 2], 10, 11);
assert_eq!(dropped, vec![10]);
assert!(c.is_empty());
}
#[test]
fn remove_covering_drops_the_matched_entry() {
let mut c = cache(1000, 8);
c.insert(vec![1, 2], 10, 1);
let m = c.longest_match(&[1, 2, 3]).unwrap();
let removed = c.remove_covering(&[1, 2, 3], m.tokens);
assert_eq!(removed, Some(10));
assert!(c.is_empty());
assert_eq!(c.remove_covering(&[1, 2, 3], m.tokens), None);
}
#[test]
fn empty_tokens_never_stored() {
let mut c = cache(1000, 8);
let dropped = c.insert(vec![], 10, 1);
assert_eq!(dropped, vec![10]);
assert!(c.is_empty());
}
}

View File

@@ -2,11 +2,11 @@
//!
//! Decodes `data:image/...;base64,...` URIs from OpenAI-style
//! `image_url` content parts into the patch tensors a candle vision
//! tower expects. Stage A ships **fixed resolution** — every image
//! is resized to the same target dimensions (default 448×448 for
//! Qwen3.6, configurable per-call) so the patch count is constant
//! per image. Variable resolution per [Qwen2VL convention] is tracked
//! as issue #14.
//! tower expects. Resolution is **dynamic** (#14): each image is
//! resized to its native aspect via Qwen `smart_resize` — a
//! factor-aligned `(h, w)` whose pixel count lands in the profile's
//! `[min_pixels, max_pixels]` budget — so the LM token count varies per
//! image (`(h/factor) × (w/factor)`).
//!
//! Spec reference: `doc/vision-qwen3_6-spec.md` — preprocessor
//! section.
@@ -21,7 +21,7 @@
//! Pipeline (per image):
//! 1. data: URI → base64 decode → bytes
//! 2. bytes → image::DynamicImage (PNG/JPEG/WebP/etc)
//! 3. resize_exact to target H×W (pixel space)
//! 3. smart_resize to a native-aspect, factor-aligned H×W (pixel space)
//! 4. RGB→f32, normalise per mean/std
//! 5. layout to (C, H, W) tensor
//!
@@ -34,39 +34,126 @@ use base64::Engine;
use image::DynamicImage;
use image::imageops::FilterType;
/// Preprocessing target. Captures the resize dimensions and the
/// channel-wise normalisation constants from the model's
/// `preprocessor_config.json`. Stage A ships a single `qwen3_6()`
/// constructor for fixed-resolution Qwen3.6 preprocessing; other
/// models can ship their own profile when added.
/// Preprocessing target. Captures the resize policy (Qwen `smart_resize`
/// factor + pixel budget) and the channel-wise normalisation constants
/// from the model's `preprocessor_config.json`. Images are resized to
/// their **native aspect** — a factor-aligned `(h, w)` whose pixel count
/// lands in `[min_pixels, max_pixels]` — not a fixed square (#14).
#[derive(Debug, Clone)]
pub struct PreprocessProfile {
pub target_height: u32,
pub target_width: u32,
/// Both output dims are multiples of this. For Qwen3.6 it is
/// `patch_size(16) × spatial_merge_size(2) = 32`, so the post-merge
/// LM grid is exactly `(h/factor, w/factor)`.
pub factor: u32,
/// Lower pixel bound — tiny images are upscaled to at least this.
pub min_pixels: u32,
/// Upper pixel bound — large images are downscaled to at most this.
/// Caps per-image LM tokens (`max_pixels / factor²`) and the
/// O(patches²) ViT attention cost.
pub max_pixels: u32,
pub image_mean: [f32; 3],
pub image_std: [f32; 3],
}
/// The Qwen3.6 vision tower rejects any image whose **patch** count
/// exceeds its learned pos-embed budget (`num_position_embeddings =
/// 2304 = 48²`; see `vision.rs`). At `patch_size = 16` that is
/// `2304 × 16² = 589_824` source pixels. `max_pixels` is hard-capped to
/// this so `smart_resize` can never produce an over-budget grid — a
/// per-rank "patch count exceeds pos_embed budget" error mid-TP-forward
/// would otherwise poison the device context. The pos-embed grid is the
/// resolution Qwen3.6 was trained at, so this cap is principled, not just
/// defensive.
const QWEN3_6_MAX_PIXELS_CAP: u32 = 2304 * 16 * 16; // 589_824 → ≤ 2304 patches → ≤ 576 LM tokens
/// Default pixel budget for Qwen3.6: `256²` (64 LM tokens) up to the
/// pos-embed cap (576 LM tokens). Generous for documents/OCR, bounded
/// for serving. Operators lower it with `NEURON_VISION_MIN_PIXELS` /
/// `NEURON_VISION_MAX_PIXELS` (the upper bound is still clamped to the
/// cap above — raising it past the budget would poison the model).
const QWEN3_6_MIN_PIXELS: u32 = 65_536;
fn env_pixels(name: &str, default: u32) -> u32 {
std::env::var(name)
.ok()
.and_then(|v| v.trim().parse::<u32>().ok())
.unwrap_or(default)
}
impl PreprocessProfile {
/// Stage A profile for Qwen3.6. Resize to 448×448, normalise to
/// `[-1, 1]` via mean=std=0.5. Fits within the model's
/// `num_position_embeddings=2304` budget at 28×28 = 784 patches
/// before merging.
/// Profile for Qwen3.6. Native-aspect `smart_resize` (factor 32),
/// normalise to `[-1, 1]` via mean=std=0.5. Pixel budget defaults to
/// [`QWEN3_6_MIN_PIXELS`]…[`QWEN3_6_MAX_PIXELS_CAP`], overridable via
/// `NEURON_VISION_MIN_PIXELS` / `NEURON_VISION_MAX_PIXELS`. Clamped
/// sane: `factor² ≤ min ≤ max`, and `max ≤` the pos-embed cap (so the
/// vision tower never rejects a resized image and poisons the context).
pub fn qwen3_6() -> Self {
let factor = 32u32;
let f2 = factor * factor;
let min_pixels = env_pixels("NEURON_VISION_MIN_PIXELS", QWEN3_6_MIN_PIXELS)
.max(f2)
.min(QWEN3_6_MAX_PIXELS_CAP);
let max_pixels = env_pixels("NEURON_VISION_MAX_PIXELS", QWEN3_6_MAX_PIXELS_CAP)
.min(QWEN3_6_MAX_PIXELS_CAP)
.max(min_pixels);
Self {
target_height: 448,
target_width: 448,
factor,
min_pixels,
max_pixels,
image_mean: [0.5, 0.5, 0.5],
image_std: [0.5, 0.5, 0.5],
}
}
/// Per-channel CHW tensor length: 3 * H * W.
pub fn pixels_chw(&self) -> usize {
3 * (self.target_height as usize) * (self.target_width as usize)
/// The factor-aligned `(h, w)` this profile would resize a source
/// `src_h × src_w` image to. Pure integer policy — no pixel work.
pub fn resized_dims(&self, src_h: u32, src_w: u32) -> Result<(u32, u32)> {
smart_resize(src_h, src_w, self.factor, self.min_pixels, self.max_pixels)
}
}
/// Qwen `smart_resize`: the smallest `factor`-aligned `(h_bar, w_bar)`
/// that preserves aspect ratio as closely as possible while keeping the
/// pixel count within `[min_pixels, max_pixels]`. Direct port of the
/// canonical Qwen2-VL / Qwen3-VL image-processor function (so neuron's
/// grid matches what the model was trained on).
///
/// Returns `(height, width)`. Errors if the aspect ratio exceeds 200:1
/// (degenerate input — a 1-pixel-tall strip), matching upstream.
pub fn smart_resize(
height: u32,
width: u32,
factor: u32,
min_pixels: u32,
max_pixels: u32,
) -> Result<(u32, u32)> {
let h = height.max(1) as f64;
let w = width.max(1) as f64;
let ratio = h.max(w) / h.min(w);
if ratio > 200.0 {
anyhow::bail!(
"image aspect ratio {ratio:.1}:1 exceeds the 200:1 limit ({height}×{width}); \
refusing to resize"
);
}
let f = factor as f64;
let (minp, maxp) = (min_pixels as f64, max_pixels as f64);
// round-to-nearest-factor (may be 0 for sub-factor inputs; the
// min-pixels branch below grows it back up).
let mut h_bar = (h / f).round() * f;
let mut w_bar = (w / f).round() * f;
if h_bar * w_bar > maxp {
let beta = (h * w / maxp).sqrt();
h_bar = f.max((h / beta / f).floor() * f);
w_bar = f.max((w / beta / f).floor() * f);
} else if h_bar * w_bar < minp {
let beta = (minp / (h * w)).sqrt();
h_bar = (h * beta / f).ceil() * f;
w_bar = (w * beta / f).ceil() * f;
}
Ok((h_bar as u32, w_bar as u32))
}
/// Decode a `data:image/...;base64,...` URI into an in-memory image.
///
/// Accepts the OpenAI Chat Completions `image_url` shape — a string
@@ -106,16 +193,13 @@ pub fn decode_data_uri(uri: &str) -> Result<DynamicImage> {
/// faster on CPU. Quality difference is marginal for downstream
/// vision-encoder consumption. The numerical-validation issue (#15)
/// will quantify any discrepancy.
pub fn preprocess(img: &DynamicImage, profile: &PreprocessProfile) -> Vec<f32> {
pub fn preprocess(img: &DynamicImage, profile: &PreprocessProfile) -> Result<(Vec<f32>, u32, u32)> {
let (h_bar, w_bar) = profile.resized_dims(img.height(), img.width())?;
let rgb = img
.resize_exact(
profile.target_width,
profile.target_height,
FilterType::Triangle,
)
.resize_exact(w_bar, h_bar, FilterType::Triangle)
.to_rgb8();
let h = profile.target_height as usize;
let w = profile.target_width as usize;
let h = h_bar as usize;
let w = w_bar as usize;
let mut out = vec![0.0_f32; 3 * h * w];
// Row-major (C, H, W). Candle's Conv2d expects NCHW, so this is
// the natural layout — the caller stacks `n` of these along the
@@ -131,16 +215,27 @@ pub fn preprocess(img: &DynamicImage, profile: &PreprocessProfile) -> Vec<f32> {
}
}
}
out
Ok((out, h_bar, w_bar))
}
/// Combined helper: decode + preprocess in one call. Most call
/// sites just want the final tensor; the two-step path exists for
/// callers (tests, future video preprocessing) that need the
/// Combined helper: decode + preprocess in one call. Returns the
/// `(3, h, w)` row-major pixels plus the resized `(h, w)` — the caller
/// needs the dims to build the tensor and to derive the LM token grid
/// `(h/factor, w/factor)`. Most call sites use this; the two-step path
/// exists for callers (tests, future video preprocessing) that need the
/// intermediate `DynamicImage`.
pub fn preprocess_data_uri(uri: &str, profile: &PreprocessProfile) -> Result<Vec<f32>> {
pub fn preprocess_data_uri(uri: &str, profile: &PreprocessProfile) -> Result<(Vec<f32>, u32, u32)> {
let img = decode_data_uri(uri)?;
Ok(preprocess(&img, profile))
preprocess(&img, profile)
}
/// Resized `(h, w)` for a data-URI image **without** running the pixel
/// normalisation — decode header + `smart_resize` only. Lets a caller
/// that just needs the LM token count (e.g. the TP leader expanding the
/// prompt) avoid materialising the full pixel tensor twice.
pub fn resized_dims_for_uri(uri: &str, profile: &PreprocessProfile) -> Result<(u32, u32)> {
let img = decode_data_uri(uri)?;
profile.resized_dims(img.height(), img.width())
}
#[cfg(test)]
@@ -205,13 +300,17 @@ mod tests {
// decoding so this test isolates the resize+normalise path.
let img: ImageBuffer<Rgb<u8>, Vec<u8>> = ImageBuffer::from_pixel(2, 2, Rgb([255, 0, 0]));
let dyn_img = DynamicImage::ImageRgb8(img);
let out = preprocess(&dyn_img, &profile);
let (out, h_bar, w_bar) = preprocess(&dyn_img, &profile).expect("preprocess");
assert_eq!(out.len(), profile.pixels_chw());
let h = h_bar as usize;
let w = w_bar as usize;
assert_eq!(out.len(), 3 * h * w);
// Dims are factor-aligned and at least the min-pixel floor.
assert_eq!(h_bar % profile.factor, 0);
assert_eq!(w_bar % profile.factor, 0);
assert!(h * w >= profile.min_pixels as usize);
// After mean=0.5, std=0.5: red channel (255/255=1.0) → (1.0 - 0.5)/0.5 = 1.0
// green/blue (0.0) → (0.0 - 0.5)/0.5 = -1.0
let h = profile.target_height as usize;
let w = profile.target_width as usize;
assert!(
(out[0] - 1.0).abs() < 1e-5,
"R[0] should be 1.0, got {}",
@@ -229,9 +328,12 @@ mod tests {
#[test]
fn preprocess_data_uri_end_to_end() {
let profile = PreprocessProfile::qwen3_6();
let out = preprocess_data_uri(&red_png_uri(), &profile).expect("e2e preprocess");
assert_eq!(out.len(), profile.pixels_chw());
let (out, h, w) = preprocess_data_uri(&red_png_uri(), &profile).expect("e2e preprocess");
assert_eq!(out.len(), 3 * h as usize * w as usize);
assert!(out.iter().all(|v| v.is_finite()));
// resized_dims_for_uri agrees with the full preprocess.
let (h2, w2) = resized_dims_for_uri(&red_png_uri(), &profile).expect("dims");
assert_eq!((h, w), (h2, w2));
}
#[test]
@@ -240,10 +342,10 @@ mod tests {
// 1x1 grayscale = 200 → after conversion to RGB, all three
// channels equal 200, normalised → (200/255 - 0.5)/0.5 ≈ 0.569
let gray = DynamicImage::ImageLuma8(ImageBuffer::from_pixel(1, 1, image::Luma([200])));
let out = preprocess(&gray, &profile);
let (out, h_bar, w_bar) = preprocess(&gray, &profile).expect("preprocess");
let expected = ((200.0 / 255.0) - 0.5) / 0.5;
let h = profile.target_height as usize;
let w = profile.target_width as usize;
let h = h_bar as usize;
let w = w_bar as usize;
for c in 0..3 {
let v = out[c * h * w];
assert!(
@@ -252,4 +354,88 @@ mod tests {
);
}
}
#[test]
fn smart_resize_keeps_factor_aligned_square_in_budget() {
// 448×448 sits inside [65536, 1048576] and is factor-aligned →
// unchanged. (Regression guard for the old fixed-res sweet spot.)
let (h, w) = smart_resize(448, 448, 32, 65_536, 1_048_576).unwrap();
assert_eq!((h, w), (448, 448));
}
#[test]
fn smart_resize_preserves_aspect_and_caps_at_max() {
// 3000×4000 (landscape) → downscaled under max_pixels, aspect kept.
let (h, w) = smart_resize(3000, 4000, 32, 65_536, 1_048_576).unwrap();
assert_eq!(h % 32, 0);
assert_eq!(w % 32, 0);
assert!(
(h as u64) * (w as u64) <= 1_048_576,
"must respect max_pixels"
);
assert!(w > h, "landscape orientation preserved");
// aspect ≈ 4000/3000 = 1.333; allow a factor-rounding tolerance.
let ar = w as f64 / h as f64;
assert!((ar - 4.0 / 3.0).abs() < 0.15, "aspect ~4:3, got {ar:.3}");
}
#[test]
fn smart_resize_floors_tiny_image_at_min() {
// 16×16 → upscaled to at least min_pixels, factor-aligned.
let (h, w) = smart_resize(16, 16, 32, 65_536, 1_048_576).unwrap();
assert_eq!(h % 32, 0);
assert_eq!(w % 32, 0);
assert!((h as u64) * (w as u64) >= 65_536, "must respect min_pixels");
}
#[test]
fn smart_resize_tall_nonsquare_stays_nonsquare() {
// A tall screenshot keeps portrait orientation.
let (h, w) = smart_resize(2000, 500, 32, 65_536, 1_048_576).unwrap();
assert!(h > w, "portrait orientation preserved");
assert_eq!(h % 32, 0);
assert_eq!(w % 32, 0);
}
#[test]
fn smart_resize_rejects_extreme_aspect() {
let err = smart_resize(1, 500, 32, 65_536, 1_048_576).unwrap_err();
assert!(format!("{err:#}").contains("200:1"));
}
#[test]
fn qwen3_6_never_exceeds_pos_embed_patch_budget() {
// The pos-embed cap must hold for huge, tall, wide, and extreme
// images — exceeding 2304 patches errors mid-tower and poisons
// the device context, so this invariant is load-bearing.
let p = PreprocessProfile::qwen3_6();
for (sh, sw) in [
(8000u32, 6000u32),
(808, 1600),
(4000, 400),
(1, 199),
(16, 16),
] {
let (h, w) = p.resized_dims(sh, sw).unwrap();
let patches = (h / 16) * (w / 16);
assert!(
patches <= 2304,
"{sh}x{sw} → {h}x{w} = {patches} patches exceeds the 2304 budget"
);
}
}
#[test]
fn qwen3_6_default_budget_bounds_lm_tokens() {
// A huge source image caps at max_pixels → the per-image LM token
// count stays within budget (so it can't blow NEURON_MAX_PROMPT_TOKENS).
let p = PreprocessProfile::qwen3_6();
let (h, w) = p.resized_dims(8000, 6000).unwrap();
let lm_tokens = (h / p.factor) * (w / p.factor);
let budget = p.max_pixels / (p.factor * p.factor);
assert!(
lm_tokens <= budget,
"max-res image LM tokens {lm_tokens} must stay within budget {budget}"
);
}
}

View File

@@ -0,0 +1,234 @@
//! Speculative decoding (#25) — a small same-family drafter proposes
//! tokens that the large target verifies in one forward pass.
//!
//! batch-1 decode is exactly the regime where speculation wins, and
//! the regime helexa lives in. A cheap drafter (Qwen3.5-0.8B) proposes
//! K tokens for the 27B target; the target verifies all K in a single
//! forward and the longest agreeing prefix is committed for free.
//!
//! ## What lives here
//!
//! This module is the **acceptance core** plus config — the pure,
//! state-free heart of the algorithm, where off-by-ones live. The
//! draft/verify loop and the GDN-state rollback (which reuses #11's
//! snapshot/restore — see the issue) wire this into the generation
//! path in later phases.
//!
//! ## Greedy acceptance
//!
//! Per round, with the target's greedy token already known at the
//! committed boundary and at each speculative position, the longest
//! drafter-matching prefix is accepted and one **bonus** token is
//! always committed on top (the target's own token at the first
//! mismatch, or a free extra token when every draft matched). So a
//! round commits between 1 and K+1 tokens — never zero, which
//! guarantees forward progress even when the drafter is useless.
//!
//! Greedy (argmax) acceptance is exact for temperature-0 sampling —
//! the fleet's probe + #22 bench regime. Stochastic acceptance that
//! preserves the target distribution at temperature > 0 is a later
//! phase.
use serde::{Deserialize, Serialize};
/// Per-target speculative-decoding settings.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SpeculativeConfig {
/// Drafter model id. MUST share the target's tokenizer/vocabulary
/// (e.g. `Qwen/Qwen3.5-0.8B` for a `Qwen/Qwen3.6-27B` target — both
/// `qwen3_5`, byte-identical tokenizer). `None` disables
/// speculation for the target.
#[serde(default)]
pub drafter: Option<String>,
/// Tokens the drafter proposes per round (K). Larger K wins more
/// when acceptance is high and loses more (wasted target compute on
/// rejected tail) when it's low. 4 is a conservative default.
#[serde(default = "default_draft_len")]
pub draft_len: usize,
}
fn default_draft_len() -> usize {
4
}
impl Default for SpeculativeConfig {
fn default() -> Self {
Self {
drafter: None,
draft_len: default_draft_len(),
}
}
}
impl SpeculativeConfig {
/// Speculation is active only when a drafter is named and K ≥ 1.
pub fn is_enabled(&self) -> bool {
self.drafter.is_some() && self.draft_len >= 1
}
}
/// Outcome of verifying one speculative round.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SpecAccept {
/// Number of drafter-proposed tokens accepted (the matching
/// prefix length), in `0..=draft.len()`.
pub accepted: usize,
/// The target's own next token, always committed after the
/// accepted prefix — the correction at the first mismatch, or a
/// free extra token when the whole draft matched.
pub bonus: u32,
}
impl SpecAccept {
/// The tokens this round commits: the accepted draft prefix
/// followed by the bonus. Always non-empty (≥ the bonus).
pub fn committed(&self, draft: &[u32]) -> Vec<u32> {
let mut out = draft[..self.accepted].to_vec();
out.push(self.bonus);
out
}
}
/// Greedy speculative acceptance.
///
/// - `draft`: the K tokens the drafter proposed this round.
/// - `target_greedy`: the target's greedy (argmax) token at each of
/// the K+1 positions — `target_greedy[j]` is what the target would
/// emit given the committed prefix plus `draft[..j]`. So
/// `target_greedy[0]` is checked against `draft[0]`, and
/// `target_greedy[K]` is the free bonus available when the whole
/// draft is accepted.
///
/// Accepts the longest prefix where the target agrees with the drafter
/// and returns the bonus token at the boundary. `target_greedy` must
/// have exactly `draft.len() + 1` entries.
pub fn greedy_accept(draft: &[u32], target_greedy: &[u32]) -> SpecAccept {
debug_assert_eq!(
target_greedy.len(),
draft.len() + 1,
"target_greedy must carry one distribution per draft position plus the bonus"
);
let mut accepted = 0;
while accepted < draft.len() && target_greedy[accepted] == draft[accepted] {
accepted += 1;
}
// `accepted` is in 0..=draft.len(), and target_greedy has
// draft.len()+1 entries, so this index is always in bounds: it's
// the target's correction at the first mismatch, or the free token
// past the end when everything matched.
SpecAccept {
accepted,
bonus: target_greedy[accepted],
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn full_accept_commits_k_plus_one() {
// Target agrees with every draft; the K+1-th greedy token is a
// free bonus.
let draft = [10, 11, 12, 13];
let target = [10, 11, 12, 13, 99];
let a = greedy_accept(&draft, &target);
assert_eq!(
a,
SpecAccept {
accepted: 4,
bonus: 99
}
);
assert_eq!(a.committed(&draft), vec![10, 11, 12, 13, 99]);
}
#[test]
fn partial_accept_takes_prefix_plus_correction() {
// Drafter right for two tokens, wrong on the third; commit the
// two + the target's correction, drop the rest of the draft.
let draft = [10, 11, 12, 13];
let target = [10, 11, 7, 13, 99];
let a = greedy_accept(&draft, &target);
assert_eq!(
a,
SpecAccept {
accepted: 2,
bonus: 7
}
);
assert_eq!(a.committed(&draft), vec![10, 11, 7]);
}
#[test]
fn zero_accept_still_commits_the_target_token() {
// First draft already wrong → accept nothing, but the target's
// own token is committed, so the round always makes progress
// (degrades to one plain decode step, never a stall).
let draft = [10, 11, 12, 13];
let target = [42, 11, 12, 13, 99];
let a = greedy_accept(&draft, &target);
assert_eq!(
a,
SpecAccept {
accepted: 0,
bonus: 42
}
);
assert_eq!(a.committed(&draft), vec![42]);
}
#[test]
fn mismatch_at_last_position() {
let draft = [10, 11, 12, 13];
let target = [10, 11, 12, 8, 99];
let a = greedy_accept(&draft, &target);
assert_eq!(
a,
SpecAccept {
accepted: 3,
bonus: 8
}
);
assert_eq!(a.committed(&draft), vec![10, 11, 12, 8]);
}
#[test]
fn single_token_draft() {
let draft = [10];
assert_eq!(
greedy_accept(&draft, &[10, 55]),
SpecAccept {
accepted: 1,
bonus: 55
}
);
assert_eq!(
greedy_accept(&draft, &[9, 55]),
SpecAccept {
accepted: 0,
bonus: 9
}
);
}
#[test]
fn config_enabled_gating() {
assert!(!SpeculativeConfig::default().is_enabled());
assert!(
!SpeculativeConfig {
drafter: Some("d".into()),
draft_len: 0,
}
.is_enabled()
);
assert!(
SpeculativeConfig {
drafter: Some("d".into()),
draft_len: 4,
}
.is_enabled()
);
}
}

View File

@@ -0,0 +1,171 @@
//! Parallel in-situ quantization (#1).
//!
//! `candle_core::quantized::QTensor::quantize` processes a tensor's
//! quantization blocks strictly sequentially on one CPU core (its
//! CUDA storage round-trips through the same CPU path), which made
//! Q6K ISQ the dominant phase of the Qwen3.6-27B TP cold load —
//! minutes of single-threaded block math per rank while 31 cores
//! idled.
//!
//! Each block is independent, so this module re-implements the same
//! quantization through candle's public per-block API
//! (`k_quants::GgmlType::from_float`) with rayon fanning the blocks
//! across the CPU pool, producing **byte-identical** output to
//! candle's sequential path (pinned by the parity tests below).
//!
//! Threading discipline: the device-to-host read and the final
//! device upload (`QStorage::from_data`) run on the *calling* thread
//! — the device worker / subprocess main thread that owns the CUDA
//! context. The rayon workers only ever touch host memory.
use anyhow::{Context, Result};
use candle_core::Tensor;
use candle_core::quantized::k_quants::{
BlockQ2K, BlockQ3K, BlockQ4_0, BlockQ4_1, BlockQ4K, BlockQ5_0, BlockQ5_1, BlockQ5K, BlockQ6K,
BlockQ8_0, BlockQ8K, GgmlType,
};
use candle_core::quantized::{GgmlDType, QStorage, QTensor};
use rayon::prelude::*;
use std::borrow::Cow;
/// Quantization blocks per rayon task. Blocks are 32256 elements; 32
/// of them per task keeps scheduling overhead negligible while a 27B
/// shard's largest tensors still split into thousands of tasks.
const BLOCKS_PER_TASK: usize = 32;
/// Drop-in replacement for `QTensor::quantize` that parallelises the
/// per-block work. Dtypes without a plain block encoding (the f32 /
/// f16 / bf16 casts, Q8_1) fall through to candle's implementation.
pub(crate) fn quantize_parallel(weight: &Tensor, dtype: GgmlDType) -> Result<QTensor> {
match dtype {
GgmlDType::Q2K => quantize_blocks::<BlockQ2K>(weight),
GgmlDType::Q3K => quantize_blocks::<BlockQ3K>(weight),
GgmlDType::Q4K => quantize_blocks::<BlockQ4K>(weight),
GgmlDType::Q5K => quantize_blocks::<BlockQ5K>(weight),
GgmlDType::Q6K => quantize_blocks::<BlockQ6K>(weight),
GgmlDType::Q8K => quantize_blocks::<BlockQ8K>(weight),
GgmlDType::Q4_0 => quantize_blocks::<BlockQ4_0>(weight),
GgmlDType::Q4_1 => quantize_blocks::<BlockQ4_1>(weight),
GgmlDType::Q5_0 => quantize_blocks::<BlockQ5_0>(weight),
GgmlDType::Q5_1 => quantize_blocks::<BlockQ5_1>(weight),
GgmlDType::Q8_0 => quantize_blocks::<BlockQ8_0>(weight),
_ => QTensor::quantize(weight, dtype)
.with_context(|| format!("QTensor::quantize fallback for {dtype:?}")),
}
}
fn quantize_blocks<T: GgmlType + Send + Sync>(weight: &Tensor) -> Result<QTensor> {
let shape = weight.shape().clone();
let block_size = T::BLCK_SIZE;
// Same constraint QTensor::quantize enforces: the last dim must
// tile into whole blocks so a block never spans two rows.
let last_dim = shape.dims().last().copied().unwrap_or(0);
if last_dim == 0 || !last_dim.is_multiple_of(block_size) {
anyhow::bail!(
"quantize_parallel: last dim of {shape:?} is not divisible by the {:?} block size {block_size}",
T::DTYPE
);
}
// Device→host read + f32 cast on the calling thread (the one
// that owns the CUDA context, when there is one).
let host: Vec<f32> = weight
.to_dtype(candle_core::DType::F32)?
.flatten_all()?
.to_vec1()
.context("copy weight to host for quantization")?;
let n_blocks = host.len() / block_size;
// Zero-initialised block buffer. The block structs have no public
// constructor, but every dispatch above is a plain `repr(C)`
// bundle of integers and (half-)floats, for which the all-zero
// bit pattern is a valid value — and `from_float` overwrites
// every block in full.
let mut blocks: Vec<T> = Vec::with_capacity(n_blocks);
// SAFETY: the buffer was allocated with capacity `n_blocks`;
// `write_bytes` zero-initialises exactly that many elements
// before `set_len` exposes them, and all-zero is a valid bit
// pattern for these POD block types (no references, no enums,
// no padding-sensitive invariants).
unsafe {
std::ptr::write_bytes(blocks.as_mut_ptr(), 0, n_blocks);
blocks.set_len(n_blocks);
}
blocks
.par_chunks_mut(BLOCKS_PER_TASK)
.zip(host.par_chunks(BLOCKS_PER_TASK * block_size))
.for_each(|(bs, xs)| T::from_float(xs, bs));
// SAFETY: a `repr(C)` slice viewed as its raw bytes; the length
// is exactly the allocation's initialised extent. `from_data`
// copies the bytes (host-side for CPU, `memcpy_htod` for CUDA)
// before this view is dropped.
let bytes = unsafe {
std::slice::from_raw_parts(
blocks.as_ptr() as *const u8,
n_blocks * std::mem::size_of::<T>(),
)
};
let storage = QStorage::from_data(Cow::Borrowed(bytes), weight.device(), T::DTYPE)
.context("upload quantized blocks")?;
Ok(QTensor::new(storage, shape)?)
}
#[cfg(test)]
mod tests {
use super::*;
use candle_core::Device;
/// The parity gate: parallel quantization must be byte-identical
/// to candle's sequential path — same per-block math, different
/// scheduling only.
fn assert_byte_parity(dtype: GgmlDType) {
let dev = Device::Cpu;
let weight = Tensor::randn(0f32, 1.0, (8, 512), &dev).unwrap();
let seq = QTensor::quantize(&weight, dtype).unwrap();
let par = quantize_parallel(&weight, dtype).unwrap();
assert_eq!(par.dtype(), seq.dtype());
assert_eq!(par.shape(), seq.shape());
let a = seq.data().unwrap();
let b = par.data().unwrap();
assert_eq!(a.as_ref(), b.as_ref(), "byte mismatch for {dtype:?}");
}
#[test]
fn parity_q6k() {
assert_byte_parity(GgmlDType::Q6K);
}
#[test]
fn parity_q4k() {
assert_byte_parity(GgmlDType::Q4K);
}
#[test]
fn parity_q5k() {
assert_byte_parity(GgmlDType::Q5K);
}
#[test]
fn parity_q8_0() {
assert_byte_parity(GgmlDType::Q8_0);
}
#[test]
fn rejects_non_divisible_last_dim() {
let dev = Device::Cpu;
// 100 is not a multiple of the 256-element k-quant block.
let weight = Tensor::randn(0f32, 1.0, (4, 100), &dev).unwrap();
assert!(quantize_parallel(&weight, GgmlDType::Q6K).is_err());
}
/// Fallback dtypes still produce a usable QTensor.
#[test]
fn fallback_f16_roundtrips() {
let dev = Device::Cpu;
let weight = Tensor::randn(0f32, 1.0, (4, 64), &dev).unwrap();
let qt = quantize_parallel(&weight, GgmlDType::F16).unwrap();
assert_eq!(qt.dtype(), GgmlDType::F16);
}
}

View File

@@ -19,6 +19,7 @@
pub mod all_reduce;
pub mod fused_load;
pub mod isq;
pub mod nccl_state;
pub mod rpc;
pub mod tp_linear;
@@ -62,21 +63,26 @@ impl TpLeaderModel {
}
}
/// Image-bearing forward on rank 0. Only the vision-capable
/// Chunked image prefill on rank 0. Only the vision-capable
/// `qwen3_5` arch supports it; the dense `qwen3` arch has no tower.
pub fn forward_with_images(
pub fn prefill_with_images_chunked(
&mut self,
input: &candle_core::Tensor,
offset: usize,
tokens: &[u32],
base_offset: usize,
image_pixels: &[candle_core::Tensor],
image_token_id: u32,
chunk_size: usize,
) -> candle_core::Result<candle_core::Tensor> {
match self {
TpLeaderModel::Qwen3_5(m) => {
m.forward_with_images(input, offset, image_pixels, image_token_id)
}
TpLeaderModel::Qwen3_5(m) => m.prefill_with_images_chunked(
tokens,
base_offset,
image_pixels,
image_token_id,
chunk_size,
),
TpLeaderModel::Qwen3(_) => {
candle_core::bail!("forward_with_images: qwen3 (dense) has no vision tower")
candle_core::bail!("prefill_with_images_chunked: qwen3 (dense) has no vision tower")
}
}
}
@@ -88,6 +94,37 @@ impl TpLeaderModel {
}
}
/// Whether this arch supports prefix snapshots (#11). Gates the
/// pool fan-out so unsupported archs never even ask the ranks.
pub fn supports_kv_snapshot(&self) -> bool {
matches!(self, TpLeaderModel::Qwen3_5(_))
}
/// Capture rank 0's cache state for a prefix snapshot (#11).
pub fn snapshot_kv_cache(
&self,
) -> candle_core::Result<crate::harness::arch::qwen3_5::snapshot::KvCacheSnapshot> {
match self {
TpLeaderModel::Qwen3(_) => {
candle_core::bail!("snapshot_kv_cache: qwen3 (dense) has no snapshot support")
}
TpLeaderModel::Qwen3_5(m) => m.snapshot_kv_cache(),
}
}
/// Replace rank 0's live cache state with a stored snapshot.
pub fn restore_kv_cache(
&mut self,
snap: &crate::harness::arch::qwen3_5::snapshot::KvCacheSnapshot,
) -> candle_core::Result<()> {
match self {
TpLeaderModel::Qwen3(_) => {
candle_core::bail!("restore_kv_cache: qwen3 (dense) has no snapshot support")
}
TpLeaderModel::Qwen3_5(m) => m.restore_kv_cache(snap),
}
}
pub fn device(&self) -> &candle_core::Device {
match self {
TpLeaderModel::Qwen3(m) => m.device(),
@@ -240,9 +277,67 @@ pub struct WorkerPool {
/// Phase 4 the load itself moves onto the worker and that bridge
/// goes away.
pub(crate) leader_worker: std::sync::Arc<super::device_worker::DeviceWorkerHandle>,
/// Cached handle to the leader's NCCL `Comm`, fetched at `init_nccl`
/// while the worker thread is responsive. The TP step watchdog uses
/// it to `ncclCommAbort` a wedged collective from the async thread —
/// the one NCCL op allowed concurrently with an in-flight collective,
/// and the only way to unblock the in-process leader thread so
/// recovery's `unload` doesn't itself hang (#17 Stage 2). `None` if
/// init couldn't cache it; the watchdog then logs that it can't abort.
#[cfg(feature = "cuda")]
leader_comm: Option<nccl_state::SendComm>,
}
/// Per-step deadline for a TP forward (#17 Stage 2). A healthy decode
/// step or chunked prefill completes in well under a second; a wedged
/// NCCL collective never returns. Generous default so no legitimate step
/// trips it; overridable via `NEURON_TP_STEP_TIMEOUT_S` (seconds).
#[cfg(feature = "cuda")]
fn tp_step_timeout() -> std::time::Duration {
let secs = std::env::var("NEURON_TP_STEP_TIMEOUT_S")
.ok()
.and_then(|v| v.trim().parse::<u64>().ok())
.filter(|&s| s > 0)
.unwrap_or(120);
std::time::Duration::from_secs(secs)
}
impl WorkerPool {
/// Abort the leader's NCCL comm to unblock a collective the watchdog
/// found wedged (#17 Stage 2). Logs the whole sequence loudly so a
/// real-world hang leaves a greppable forensic trail
/// (`tp watchdog:` / `ncclCommAbort`). Calling abort from this async
/// thread while the worker thread is blocked inside the collective is
/// the one concurrent NCCL op the library sanctions — it is how a
/// stuck/failed collective is unblocked.
#[cfg(feature = "cuda")]
fn watchdog_abort_leader_comm(&self, model_id: &str, secs: u64) {
tracing::error!(
model = %model_id,
timeout_s = secs,
"tp watchdog: leader forward exceeded deadline — NCCL collective wedged; \
aborting comm to unblock the leader thread for auto-recovery"
);
match &self.leader_comm {
Some(c) => match c.0.abort() {
Ok(()) => tracing::error!(
model = %model_id,
"tp watchdog: ncclCommAbort succeeded — wedged collective unblocked; \
failing the step so the model auto-recovers (unload+reload)"
),
Err(e) => tracing::error!(
model = %model_id, error = ?e,
"tp watchdog: ncclCommAbort failed — recovery may stall until a process restart"
),
},
None => tracing::error!(
model = %model_id,
"tp watchdog: no cached leader comm handle — cannot abort; recovery will rely \
on a process restart"
),
}
}
/// Spawn `world_size - 1` worker subprocesses. Rank 0 is the
/// leader (in-process) and is *not* spawned here — the leader
/// holds rank 0's NCCL Comm and shard in its own address space.
@@ -319,6 +414,8 @@ impl WorkerPool {
workers,
exe,
leader_worker,
#[cfg(feature = "cuda")]
leader_comm: None,
})
}
@@ -399,6 +496,23 @@ impl WorkerPool {
world_size = self.world_size,
"NCCL communicator established across all ranks"
);
// Cache the leader's Comm handle now, while the worker thread is
// responsive, so the TP step watchdog can abort a wedged
// collective later (it can't fetch it then — the thread is stuck).
// (#17 Stage 2.)
#[cfg(feature = "cuda")]
{
self.leader_comm = self.leader_worker.get_leader_comm().await;
if self.leader_comm.is_some() {
tracing::debug!("cached leader NCCL comm handle for the TP step watchdog");
} else {
tracing::warn!(
"could not cache leader NCCL comm handle; the TP step watchdog will be \
unable to abort a wedged collective (a hang would need a process restart)"
);
}
}
Ok(())
}
@@ -623,10 +737,27 @@ impl WorkerPool {
// that's the invariant the whole refactor exists to
// preserve.
let leader_start = std::time::Instant::now();
let leader_result = self
let timeout = tp_step_timeout();
let leader_fut = self
.leader_worker
.tp_forward_logits(leader_handle, tokens, offset)
.await;
.tp_forward_logits(leader_handle, tokens, offset);
let leader_result = match tokio::time::timeout(timeout, leader_fut).await {
Ok(r) => r,
Err(_elapsed) => {
// Watchdog (#17 Stage 2): the NCCL collective is wedged.
// Abort the leader comm to unblock its thread, then fail
// the step WITHOUT draining (the subprocess workers are
// wedged too; recovery's unload kills them). The error
// poisons the model → auto-recovery, which no longer hangs
// because the leader thread is now responsive.
self.watchdog_abort_leader_comm(model_id, timeout.as_secs());
anyhow::bail!(
"tp watchdog: leader forward exceeded {}s deadline; aborted wedged NCCL \
comm — model will auto-recover",
timeout.as_secs()
);
}
};
let leader_ok = leader_result.is_ok();
let leader_ms = leader_start.elapsed().as_millis();
// Surface the leader's own error at WARN before draining
@@ -722,6 +853,7 @@ impl WorkerPool {
/// embedding broadcast. Only used for prefill; decode reuses
/// `generate_step`.
#[cfg(feature = "cuda")]
#[allow(clippy::too_many_arguments)]
pub async fn generate_step_with_images(
&mut self,
model_id: &str,
@@ -730,6 +862,7 @@ impl WorkerPool {
offset: usize,
image_token_id: u32,
image_data_uris: Vec<String>,
chunk_size: usize,
) -> Result<Vec<f32>> {
let step_start = std::time::Instant::now();
let tokens_len = tokens.len();
@@ -738,6 +871,7 @@ impl WorkerPool {
tokens = tokens_len,
offset,
images = image_data_uris.len(),
chunk_size,
"WorkerPool::generate_step_with_images: fan-out"
);
@@ -749,6 +883,7 @@ impl WorkerPool {
offset,
image_token_id,
image_data_uris: image_data_uris.clone(),
chunk_size,
})
.await?;
}
@@ -758,16 +893,29 @@ impl WorkerPool {
// matching collective; CPU-side logits keep the device tensor
// from escaping the worker thread.
let leader_start = std::time::Instant::now();
let leader_result = self
.leader_worker
.tp_forward_logits_with_images(
let timeout = tp_step_timeout();
let leader_fut = self.leader_worker.tp_forward_logits_with_images(
leader_handle,
tokens,
offset,
image_token_id,
image_data_uris,
)
.await;
chunk_size,
);
let leader_result = match tokio::time::timeout(timeout, leader_fut).await {
Ok(r) => r,
Err(_elapsed) => {
// Watchdog (#17 Stage 2) — see generate_step. Vision
// prefill is still well under the deadline on healthy
// hardware; a timeout means a wedged collective.
self.watchdog_abort_leader_comm(model_id, timeout.as_secs());
anyhow::bail!(
"tp watchdog: leader image forward exceeded {}s deadline; aborted wedged \
NCCL comm — model will auto-recover",
timeout.as_secs()
);
}
};
let leader_ok = leader_result.is_ok();
let leader_ms = leader_start.elapsed().as_millis();
if !leader_ok {
@@ -877,6 +1025,123 @@ impl WorkerPool {
Ok(())
}
/// Capture every rank's cache state as one prefix snapshot (#11)
/// stored under `snapshot_id` (minted by the caller). All ranks
/// are at the same token boundary — step fan-out is synchronous —
/// so the per-rank snapshots are mutually consistent. Returns the
/// total snapshot bytes across all ranks (for budget accounting).
/// On any rank failing, the caller must `drop_kv_snapshot` the id
/// to clean up the ranks that did store.
#[cfg(feature = "cuda")]
pub async fn snapshot_kv_cache(
&mut self,
model_id: &str,
leader_handle: super::device_worker::TpHandle,
snapshot_id: u64,
) -> Result<u64> {
for w in &mut self.workers {
w.send_only(&WorkerRequest::SnapshotKvCache {
model_id: model_id.to_string(),
snapshot_id,
})
.await?;
}
let leader_result = self
.leader_worker
.tp_snapshot_kv(leader_handle, snapshot_id)
.await;
let worker_errors = drain_workers(&mut self.workers, |r| match r {
WorkerResponse::KvSnapshotStored { .. } => Ok(()),
WorkerResponse::Error { kind, message } => Err(format!("[{kind}]: {message}")),
other => Err(format!("expected KvSnapshotStored, got {other:?}")),
})
.await;
let leader_bytes = match leader_result {
Ok(b) => b,
Err(e) => anyhow::bail!("leader TP snapshot via device worker: {e}"),
};
if !worker_errors.is_empty() {
anyhow::bail!("SnapshotKvCache: {}", worker_errors.join("; "));
}
// Shards are equal-sized by construction, so the fleet total
// is the leader's bytes times the rank count.
Ok(leader_bytes.saturating_mul(self.workers.len() as u64 + 1))
}
/// Restore the snapshot `snapshot_id` on every rank, instead of
/// `clear_kv_cache`, so prefill resumes at the snapshot's token
/// boundary. On failure the caller must fall back to
/// `clear_kv_cache` + full prefill (and drop the snapshot).
#[cfg(feature = "cuda")]
pub async fn restore_kv_cache(
&mut self,
model_id: &str,
leader_handle: super::device_worker::TpHandle,
snapshot_id: u64,
) -> Result<()> {
for w in &mut self.workers {
w.send_only(&WorkerRequest::RestoreKvCache {
model_id: model_id.to_string(),
snapshot_id,
})
.await?;
}
let leader_result = self
.leader_worker
.tp_restore_kv(leader_handle, snapshot_id)
.await;
let worker_errors = drain_workers(&mut self.workers, |r| match r {
WorkerResponse::KvCacheRestored => Ok(()),
WorkerResponse::Error { kind, message } => Err(format!("[{kind}]: {message}")),
other => Err(format!("expected KvCacheRestored, got {other:?}")),
})
.await;
if let Err(e) = leader_result {
anyhow::bail!("leader TP restore via device worker: {e}");
}
if !worker_errors.is_empty() {
anyhow::bail!("RestoreKvCache: {}", worker_errors.join("; "));
}
Ok(())
}
/// Drop the snapshot `snapshot_id` on every rank (prefix-cache
/// eviction / failed-snapshot cleanup). Best-effort and
/// idempotent — errors are collected, not fatal to the caller's
/// request path, but surfaced for logging.
#[cfg(feature = "cuda")]
pub async fn drop_kv_snapshot(
&mut self,
model_id: &str,
leader_handle: super::device_worker::TpHandle,
snapshot_id: u64,
) -> Result<()> {
for w in &mut self.workers {
w.send_only(&WorkerRequest::DropKvSnapshot {
model_id: model_id.to_string(),
snapshot_id,
})
.await?;
}
let leader_result = self
.leader_worker
.tp_drop_kv_snapshot(leader_handle, snapshot_id)
.await;
let worker_errors = drain_workers(&mut self.workers, |r| match r {
WorkerResponse::KvSnapshotDropped => Ok(()),
WorkerResponse::Error { kind, message } => Err(format!("[{kind}]: {message}")),
other => Err(format!("expected KvSnapshotDropped, got {other:?}")),
})
.await;
if let Err(e) = leader_result {
anyhow::bail!("leader TP drop snapshot via device worker: {e}");
}
if !worker_errors.is_empty() {
anyhow::bail!("DropKvSnapshot: {}", worker_errors.join("; "));
}
Ok(())
}
/// Drop this model's shards on every rank. The leader's shard is
/// expected to have been dropped by the caller (its `Arc` was held
/// in the TpLoadedModel and goes away when that's removed).

View File

@@ -119,40 +119,25 @@ mod cuda_impl {
}
}
/// `Arc<Comm>` doesn't impl `Send` because `Comm` wraps a raw
/// `ncclComm_t` pointer. The NCCL contract is "operations against a
/// given comm must be serialised", not "the handle must stay on the
/// thread that created it" — so it's safe to move an `Arc<Comm>`
/// across threads as long as no concurrent ops are issued. The
/// pool's outer Mutex serialises us into `spawn_blocking`, so this
/// wrapper at the move boundary is the only thing missing.
/// Thin newtype over `Arc<Comm>`, kept for call-site clarity — it marks
/// the points where a comm handle is intentionally moved across threads
/// (e.g. cached async-side for the TP step watchdog's `ncclCommAbort`).
///
/// `Sync` is also marked safe because the `Arc<Comm>` clones held
/// by the row-parallel layers are only used from the
/// `spawn_blocking` thread driving the forward pass; concurrent
/// access from another thread would still be a bug.
/// `Send`/`Sync` are provided upstream by `cudarc`'s `Comm` (which
/// asserts the NCCL thread-safety invariant, including aborting from a
/// different thread than one inside a collective), so this type derives
/// them automatically — no manual `unsafe impl` here.
pub struct SendComm(pub Arc<Comm>);
// SAFETY: see the doc-comment above; the invariant is enforced at
// the call site (pool Mutex + single spawn_blocking thread), not at
// the type level.
unsafe impl Send for SendComm {}
unsafe impl Sync for SendComm {}
impl SendComm {
pub fn into_inner(self) -> Arc<Comm> {
self.0
}
}
// SAFETY: `cudarc::nccl::Comm` contains a raw `ncclComm_t` pointer
// (libnccl-allocated state). NCCL requires that operations against
// one Comm be issued one at a time; we serialise access by storing
// NcclState behind a Mutex in `WorkerPool`. The Comm itself is
// move-safe — NCCL doesn't track the calling OS thread, only the
// stream the operations are dispatched against.
unsafe impl Send for NcclState {}
unsafe impl Sync for NcclState {}
// `NcclState`'s `Send`/`Sync` are auto-derived: its `Arc<Comm>` and
// `Arc<CudaContext>` fields are now `Send`/`Sync` (cudarc asserts the
// comm thread-safety invariant), so no manual `unsafe impl` is needed.
/// Generate a fresh NCCL `Id` and return it hex-encoded. Used by
/// the leader to mint the shared communicator id which is then

View File

@@ -109,6 +109,10 @@ pub enum WorkerRequest {
/// image in prompt order. Each rank decodes + preprocesses these
/// identically; tens of KB each, so cheap over the stdin pipe.
image_data_uris: Vec<String>,
/// Prefill chunk size (tokens). Sent explicitly so every rank
/// walks the prompt in identical windows and the per-chunk
/// row-parallel collectives stay paired across ranks.
chunk_size: usize,
},
/// Reset the KV cache for this model on this rank. Sent at the
@@ -116,6 +120,24 @@ pub enum WorkerRequest {
/// attend over the previous one's tokens.
ClearKvCache { model_id: String },
/// Capture this rank's live cache state as a prefix snapshot
/// (#11), stored in-process under `snapshot_id`. The id is minted
/// by the leader's pool and broadcast so every rank keys the same
/// snapshot identically; all ranks are at the same token boundary
/// because step fan-out is synchronous. Worker replies
/// `KvSnapshotStored { bytes }` with this rank's snapshot size.
SnapshotKvCache { model_id: String, snapshot_id: u64 },
/// Replace this rank's live cache state with the stored snapshot,
/// instead of `ClearKvCache`, so prefill resumes at the snapshot's
/// token boundary. The snapshot remains stored.
RestoreKvCache { model_id: String, snapshot_id: u64 },
/// Drop one stored snapshot on this rank (prefix-cache eviction).
/// Idempotent — replies `KvSnapshotDropped` whether or not the id
/// was present.
DropKvSnapshot { model_id: String, snapshot_id: u64 },
/// Drop this rank's shard for the given model. Releases the VRAM
/// the shard's weights occupied; subsequent `GenerateStep` calls
/// against the same `model_id` return an `Error`.
@@ -164,6 +186,18 @@ pub enum WorkerResponse {
/// Reply to `ClearKvCache`. Empty payload.
KvCacheCleared,
/// Reply to `SnapshotKvCache`. Carries this rank's snapshot size
/// in bytes so the leader can budget-account the whole fleet's
/// footprint (shards are symmetric, so leader bytes × world_size
/// is also a fine estimate; the explicit number keeps it honest).
KvSnapshotStored { bytes: u64 },
/// Reply to `RestoreKvCache`. Empty payload.
KvCacheRestored,
/// Reply to `DropKvSnapshot`. Empty payload.
KvSnapshotDropped,
/// Reply to `UnloadModel`. Empty payload. The named model is no
/// longer present on this rank.
Unloaded,
@@ -222,6 +256,7 @@ mod tests {
offset: 0,
image_token_id: 248056,
image_data_uris: vec!["data:image/png;base64,AAA=".into()],
chunk_size: 512,
};
let wire = serde_json::to_string(&req).unwrap();
assert!(wire.contains(r#""op":"generate_step_with_images""#));

View File

@@ -24,7 +24,7 @@
//! sum carries it exactly once.
use anyhow::{Context, Result};
use candle_core::quantized::{GgmlDType, QMatMul, QTensor};
use candle_core::quantized::{GgmlDType, QMatMul};
use candle_core::{Module, Tensor};
use candle_nn::Linear;
use candle_nn::var_builder::{Shard, ShardedVarBuilder};
@@ -56,9 +56,11 @@ impl MaybeQuantLinear {
pub fn from_weight(weight: Tensor, quant: Option<GgmlDType>) -> Result<Self> {
match quant {
Some(dtype) => {
let qt = QTensor::quantize(&weight, dtype).with_context(|| {
// Parallel ISQ (#1): same bytes as QTensor::quantize,
// blocks fanned across the rayon pool.
let qt = super::isq::quantize_parallel(&weight, dtype).with_context(|| {
format!(
"QTensor::quantize to {dtype:?} for shape {:?}",
"quantize_parallel to {dtype:?} for shape {:?}",
weight.shape()
)
})?;

View File

@@ -46,6 +46,7 @@ use super::tp_linear::{ColumnParallelLinear, RowParallelLinear};
use crate::harness::arch::qwen3_5::linear_attn::repeat_interleave;
use crate::harness::arch::qwen3_5::rmsnorm::{Qwen3_5RmsNorm, Qwen3_5RmsNormGated, l2norm};
use crate::harness::arch::qwen3_5::rope::RotaryEmbedding;
use crate::harness::arch::qwen3_5::snapshot::{KvCacheSnapshot, LayerKvSnapshot};
use crate::harness::arch::qwen3_5::splice_runs;
use crate::harness::arch::qwen3_5::vision::VisionTower;
pub use crate::harness::arch::qwen3_5::{Config, TextConfig};
@@ -258,6 +259,39 @@ impl TpQwen3_5GatedDeltaNet {
self.state = TpGatedDeltaNetState::default();
}
/// Deep-copy this rank's recurrent state for a prefix snapshot.
/// Same in-place-kernel rationale as the single-GPU
/// `GatedDeltaNet::snapshot_state`.
pub fn snapshot_state(&self) -> candle_core::Result<(Option<Tensor>, Option<Tensor>)> {
let conv = self
.state
.conv_state
.as_ref()
.map(Tensor::copy)
.transpose()?;
let rec = self
.state
.recurrent_state
.as_ref()
.map(Tensor::copy)
.transpose()?;
Ok((conv, rec))
}
/// Replace this rank's live recurrent state with a deep copy of a
/// snapshot. See the single-GPU `GatedDeltaNet::restore_state`.
pub fn restore_state(
&mut self,
conv_state: Option<&Tensor>,
recurrent_state: Option<&Tensor>,
) -> candle_core::Result<()> {
self.state = TpGatedDeltaNetState {
conv_state: conv_state.map(Tensor::copy).transpose()?,
recurrent_state: recurrent_state.map(Tensor::copy).transpose()?,
};
Ok(())
}
/// `x` shape: `(B, L, hidden_size)`. Returns `(B, L, hidden_size)`
/// after the row-parallel AllReduce.
pub fn forward(&mut self, x: &Tensor) -> candle_core::Result<Tensor> {
@@ -526,7 +560,8 @@ impl TpQwen3_5Attention {
&mut self,
x: &Tensor,
attn_mask: Option<&Tensor>,
offset: usize,
cos: &Tensor,
sin: &Tensor,
) -> candle_core::Result<Tensor> {
let (b, l, _) = x.dims3()?;
@@ -559,7 +594,7 @@ impl TpQwen3_5Attention {
.transpose(1, 2)?
.contiguous()?;
let (q, k) = self.rotary.apply(&q, &k, offset)?;
let (q, k) = self.rotary.apply_cos_sin(&q, &k, cos, sin)?;
let (k, v) = self.kv_cache.append(&k, &v)?;
let k = repeat_kv(k, self.num_kv_groups)?.contiguous()?;
let v = repeat_kv(v, self.num_kv_groups)?.contiguous()?;
@@ -584,6 +619,25 @@ impl TpQwen3_5Attention {
pub fn clear_kv_cache(&mut self) {
self.kv_cache.reset();
}
/// Capture this rank's KV cache for a prefix snapshot. Shallow
/// clones are safe — see the single-GPU
/// `Qwen3_5Attention::snapshot_kv`.
pub fn snapshot_kv(&self) -> Option<(Tensor, Tensor)> {
match (self.kv_cache.k(), self.kv_cache.v()) {
(Some(k), Some(v)) => Some((k.clone(), v.clone())),
_ => None,
}
}
/// Replace this rank's live KV cache with a snapshot.
pub fn restore_kv(&mut self, snap: Option<&(Tensor, Tensor)>) -> candle_core::Result<()> {
self.kv_cache.reset();
if let Some((k, v)) = snap {
self.kv_cache.append(k, v)?;
}
Ok(())
}
}
// ─── MLP ────────────────────────────────────────────────────────────
@@ -807,11 +861,12 @@ impl TpQwen3_5DecoderLayer {
&mut self,
x: &Tensor,
attn_mask: Option<&Tensor>,
offset: usize,
cos: &Tensor,
sin: &Tensor,
) -> candle_core::Result<Tensor> {
let h = self.input_layernorm.forward(x)?;
let attn_out = match &mut self.attention {
TpAttentionKind::Full(attn) => attn.forward(&h, attn_mask, offset)?,
TpAttentionKind::Full(attn) => attn.forward(&h, attn_mask, cos, sin)?,
TpAttentionKind::Linear(net) => net.forward(&h)?,
};
let x = (x + attn_out)?;
@@ -826,6 +881,39 @@ impl TpQwen3_5DecoderLayer {
TpAttentionKind::Linear(n) => n.clear_kv_cache(),
}
}
/// Capture this layer's per-rank cache state for a prefix
/// snapshot. Reuses the single-GPU snapshot types — the shard
/// state has the same shape, just sharded head dims.
pub fn snapshot_kv(&self) -> candle_core::Result<LayerKvSnapshot> {
Ok(match &self.attention {
TpAttentionKind::Full(a) => LayerKvSnapshot::Full(a.snapshot_kv()),
TpAttentionKind::Linear(n) => {
let (conv_state, recurrent_state) = n.snapshot_state()?;
LayerKvSnapshot::Linear {
conv_state,
recurrent_state,
}
}
})
}
/// Replace this layer's per-rank cache state from a snapshot.
pub fn restore_kv(&mut self, snap: &LayerKvSnapshot) -> candle_core::Result<()> {
match (&mut self.attention, snap) {
(TpAttentionKind::Full(a), LayerKvSnapshot::Full(kv)) => a.restore_kv(kv.as_ref()),
(
TpAttentionKind::Linear(n),
LayerKvSnapshot::Linear {
conv_state,
recurrent_state,
},
) => n.restore_state(conv_state.as_ref(), recurrent_state.as_ref()),
_ => candle_core::bail!(
"restore_kv: snapshot layer kind does not match this layer's attention kind"
),
}
}
}
// ─── base Model ─────────────────────────────────────────────────────
@@ -834,6 +922,15 @@ pub struct TpQwen3_5Model {
embed_tokens: Embedding,
layers: Vec<TpQwen3_5DecoderLayer>,
norm: Qwen3_5RmsNorm,
/// Replicated rotary, shared with every full-attention layer. The
/// model builds the per-forward cos/sin (interleaved M-RoPE for image
/// tokens, plain for text) once and the layers apply it. Identical on
/// every rank, so per-rank position ids stay consistent.
rotary: Arc<RotaryEmbedding>,
/// `offset + rope_delta` is the text-axis decode position; set from
/// `get_rope_index` during a vision prefill, reset in `clear_kv_cache`.
/// See `Qwen3_5Model::rope_delta`.
rope_delta: i64,
device: Device,
dtype: DType,
}
@@ -874,7 +971,12 @@ impl TpQwen3_5Model {
let vb_l = text_vb.pp("layers");
let mut layers = Vec::with_capacity(cfg.num_hidden_layers);
log_vram(&device, rank, "before layer 0");
// Per-phase timing (#1): the layer loop is where ISQ cost
// concentrates; the per-layer line is debug, the loop total
// info, so journalctl always shows where a cold load went.
let layers_start = std::time::Instant::now();
for i in 0..cfg.num_hidden_layers {
let layer_start = std::time::Instant::now();
let layer = TpQwen3_5DecoderLayer::load(
cfg,
rotary.clone(),
@@ -891,8 +993,20 @@ impl TpQwen3_5Model {
format!("load layer {i} (rank {rank}): free={free_mb}MB / total={total_mb}MB")
})?;
layers.push(layer);
tracing::debug!(
rank,
layer = i,
elapsed_ms = layer_start.elapsed().as_millis() as u64,
"TP layer loaded"
);
log_vram(&device, rank, &format!("after layer {i}"));
}
tracing::info!(
rank,
layers = cfg.num_hidden_layers,
elapsed_ms = layers_start.elapsed().as_millis() as u64,
"TP layer loop complete"
);
let norm = Qwen3_5RmsNorm::load(&text_vb.pp("norm"), cfg.hidden_size, cfg.rms_norm_eps)?;
@@ -900,6 +1014,8 @@ impl TpQwen3_5Model {
embed_tokens,
layers,
norm,
rotary,
rope_delta: 0,
device,
dtype,
})
@@ -937,6 +1053,7 @@ impl TpQwen3_5Model {
let vb_l = text_vb.pp("layers");
let mut layers = Vec::with_capacity(cfg.num_hidden_layers);
let layers_start = std::time::Instant::now();
for i in 0..cfg.num_hidden_layers {
layers.push(TpQwen3_5DecoderLayer::load(
cfg,
@@ -949,6 +1066,12 @@ impl TpQwen3_5Model {
quant,
)?);
}
tracing::info!(
rank,
layers = cfg.num_hidden_layers,
elapsed_ms = layers_start.elapsed().as_millis() as u64,
"TP layer loop complete"
);
let norm = Qwen3_5RmsNorm::load(&text_vb.pp("norm"), cfg.hidden_size, cfg.rms_norm_eps)?;
@@ -956,6 +1079,8 @@ impl TpQwen3_5Model {
embed_tokens,
layers,
norm,
rotary,
rope_delta: 0,
device,
dtype,
})
@@ -969,6 +1094,46 @@ impl TpQwen3_5Model {
for l in &mut self.layers {
l.clear_kv_cache();
}
self.rope_delta = 0;
}
/// Capture this rank's per-layer cache state plus the rope
/// position counter as one consistent prefix snapshot (#11).
/// Mirrors `Qwen3_5Model::snapshot_kv_cache`.
pub fn snapshot_kv_cache(&self) -> candle_core::Result<KvCacheSnapshot> {
let layers = self
.layers
.iter()
.map(|l| l.snapshot_kv())
.collect::<candle_core::Result<Vec<_>>>()?;
Ok(KvCacheSnapshot {
layers,
rope_delta: self.rope_delta,
})
}
/// Replace this rank's live cache state with a snapshot. The
/// snapshot stays valid for further restores.
pub fn restore_kv_cache(&mut self, snap: &KvCacheSnapshot) -> candle_core::Result<()> {
if snap.layers.len() != self.layers.len() {
candle_core::bail!(
"restore_kv_cache: snapshot has {} layers, model has {}",
snap.layers.len(),
self.layers.len()
);
}
for (layer, layer_snap) in self.layers.iter_mut().zip(snap.layers.iter()) {
layer.restore_kv(layer_snap)?;
}
self.rope_delta = snap.rope_delta;
Ok(())
}
/// Set the decode `rope_delta` computed by `get_rope_index` during a
/// vision prefill, so decode after the image resumes text positions
/// from the image-compressed counter.
pub fn set_rope_delta(&mut self, delta: i64) {
self.rope_delta = delta;
}
fn causal_mask(&self, b: usize, tgt: usize, offset: usize) -> candle_core::Result<Tensor> {
@@ -980,64 +1145,80 @@ impl TpQwen3_5Model {
}
pub fn forward(&mut self, input: &Tensor, offset: usize) -> candle_core::Result<Tensor> {
let (b, l) = input.dims2()?;
let mut h = self.embed_tokens.forward(input)?;
let causal = if l == 1 {
None
} else {
Some(self.causal_mask(b, l, offset)?)
};
for layer in &mut self.layers {
h = layer.forward(&h, causal.as_ref(), offset)?;
}
self.norm.forward(&h)
self.forward_inner(input, offset, None, None, None)
}
/// Forward with image-embedding splice (TP, replicated tower).
///
/// Mirrors the single-GPU `Qwen3_5Model::forward_inner` splice:
/// embed locally, replace the rows at `image_token_id` positions
/// with the image patch embeddings, then run the sharded decoder
/// stack. The TP invariant is that every rank holds an identical
/// hidden state (only the attention/MLP matmuls shard, with a
/// trailing `AllReduce`). That holds here because every rank
/// encodes the *same* pixels through its *replicated* vision tower
/// and so produces identical `image_embeds` — no broadcast needed.
pub fn forward_with_vision(
/// Forward for a vision-prefill chunk: optional image-embedding
/// splice plus explicit interleaved-M-RoPE `position_ids` (the
/// chunk's slice of the full prompt's 3D positions). Used by
/// `TpQwen3_5ForCausalLM::prefill_with_images_chunked`, which
/// computes the positions once over the whole prompt and slices them
/// per chunk so every rank steps in lockstep.
pub fn forward_with_positions(
&mut self,
input: &Tensor,
offset: usize,
image_embeds: &Tensor,
image_token_id: u32,
position_ids: &Tensor,
image_embeds: Option<&Tensor>,
image_token_id: Option<u32>,
) -> candle_core::Result<Tensor> {
self.forward_inner(
input,
offset,
image_embeds,
image_token_id,
Some(position_ids),
)
}
/// Shared forward. Splices image embeddings at `image_token_id`
/// positions when present, then builds the rotary cos/sin — from the
/// explicit `position_ids` (interleaved M-RoPE, vision) when given,
/// else plain positions at `offset + rope_delta` (text / decode) —
/// and runs the sharded decoder stack. The TP replicated-hidden-state
/// invariant holds because every rank encodes the same pixels and
/// computes the same positions.
fn forward_inner(
&mut self,
input: &Tensor,
offset: usize,
image_embeds: Option<&Tensor>,
image_token_id: Option<u32>,
position_ids: Option<&Tensor>,
) -> candle_core::Result<Tensor> {
let (b, l) = input.dims2()?;
let mut h = self.embed_tokens.forward(input)?;
// Locate the image-token positions in the (pre-expanded) input
// ids and splice the patch rows in. Same CPU-side scan as the
// single-GPU path; the count must match the patch dimension or
// the prompt expansion is wrong.
if let (Some(img), Some(tok_id)) = (image_embeds, image_token_id) {
let ids: Vec<u32> = input.flatten_all()?.to_vec1()?;
let mut positions: Vec<u32> = Vec::with_capacity(image_embeds.dim(0)?);
let mut positions: Vec<u32> = Vec::with_capacity(img.dim(0)?);
for (idx, id) in ids.iter().enumerate() {
if *id == image_token_id {
if *id == tok_id {
positions.push(idx as u32);
}
}
let n_img_tokens = image_embeds.dim(0)?;
let n_img_tokens = img.dim(0)?;
if positions.len() != n_img_tokens {
candle_core::bail!(
"TP forward_with_vision: prompt has {} image-token positions but \
image_embeds carries {} tokens — ensure the per-image patch-count \
expansion has been applied",
"TP forward: chunk has {} image-token positions but image_embeds carries \
{} tokens — patch-count expansion / chunk slicing mismatch",
positions.len(),
n_img_tokens,
);
}
if !positions.is_empty() {
let img = image_embeds.to_dtype(self.dtype)?;
let img = img.to_dtype(self.dtype)?;
h = splice_runs(&h, &img, &positions)?;
}
}
let (cos, sin) = match position_ids {
Some(pos) => self.rotary.mrope_cos_sin(pos)?,
None => {
let base = (offset as i64 + self.rope_delta).max(0) as usize;
self.rotary.plain_cos_sin(base, l)?
}
};
let causal = if l == 1 {
None
@@ -1045,7 +1226,7 @@ impl TpQwen3_5Model {
Some(self.causal_mask(b, l, offset)?)
};
for layer in &mut self.layers {
h = layer.forward(&h, causal.as_ref(), offset)?;
h = layer.forward(&h, causal.as_ref(), &cos, &sin)?;
}
self.norm.forward(&h)
}
@@ -1174,21 +1355,25 @@ impl TpQwen3_5ForCausalLM {
hidden.i((.., l - 1.., ..))?.apply(&self.lm_head)
}
/// Forward with image-embedding splice (TP). Mirrors `forward` but
/// routes through `TpQwen3_5Model::forward_with_vision` so the
/// per-rank input embeddings get the image patches spliced in at
/// `image_token_id` positions before the sharded decoder stack.
pub fn forward_with_vision(
/// Forward for a vision-prefill chunk (optional image splice +
/// explicit interleaved-M-RoPE `position_ids`). Mirrors `forward`
/// but routes through `TpQwen3_5Model::forward_with_positions`.
pub fn forward_with_positions(
&mut self,
input: &Tensor,
offset: usize,
image_embeds: &Tensor,
image_token_id: u32,
position_ids: &Tensor,
image_embeds: Option<&Tensor>,
image_token_id: Option<u32>,
) -> candle_core::Result<Tensor> {
let (_, l) = input.dims2()?;
let hidden = self
.base
.forward_with_vision(input, offset, image_embeds, image_token_id)?;
let hidden = self.base.forward_with_positions(
input,
offset,
position_ids,
image_embeds,
image_token_id,
)?;
hidden.i((.., l - 1.., ..))?.apply(&self.lm_head)
}
@@ -1200,19 +1385,10 @@ impl TpQwen3_5ForCausalLM {
/// identical encode → splice → forward and keeps the replicated
/// hidden state in lockstep. Returns last-position logits
/// `(B, 1, vocab)`, same contract as `forward`.
pub fn forward_with_images(
&mut self,
input: &Tensor,
offset: usize,
image_pixels: &[Tensor],
image_token_id: u32,
) -> candle_core::Result<Tensor> {
if image_pixels.is_empty() {
candle_core::bail!("forward_with_images: called with zero images");
}
// Encode each image (immutable borrows of the tower) before the
// mutable forward below; the borrows end as each owned embedding
// is pushed.
/// Encode every preprocessed `(C,H,W)` image once through this
/// rank's replicated tower and concatenate along the patch axis →
/// `(sum_patches, hidden)`. Done once per prefill, not per chunk.
fn encode_images_concat(&self, image_pixels: &[Tensor]) -> candle_core::Result<Tensor> {
let mut per_image = Vec::with_capacity(image_pixels.len());
for (idx, img) in image_pixels.iter().enumerate() {
let embed = self
@@ -1220,14 +1396,127 @@ impl TpQwen3_5ForCausalLM {
.map_err(|e| candle_core::Error::Msg(format!("encode image[{idx}]: {e:#}")))?;
per_image.push(embed);
}
let image_embeds = Tensor::cat(&per_image.iter().collect::<Vec<_>>(), 0)?;
self.forward_with_vision(input, offset, &image_embeds, image_token_id)
Tensor::cat(&per_image.iter().collect::<Vec<_>>(), 0)
}
/// Chunked image prefill on one rank. Encodes the image(s) once,
/// then walks the (pre-expanded) prompt in `chunk_size`-token
/// windows — exactly like the text `chunked_prefill_tp` — splicing
/// the patch embeddings into whichever chunk(s) carry `<|image_pad|>`
/// positions. Activation memory is bounded by the chunk, not the
/// full prompt, so a long vision context no longer single-shot-OOMs.
///
/// Every rank runs the identical chunk sequence (same `tokens.len()`
/// and `chunk_size`), so the row-parallel `AllReduce`s pair up
/// chunk-by-chunk across ranks with no extra synchronisation. The KV
/// cache accumulates across chunks via the growing offset; only the
/// final chunk's last-position logits are returned (intermediate
/// chunks just populate the cache, same as the text path).
pub fn prefill_with_images_chunked(
&mut self,
tokens: &[u32],
base_offset: usize,
image_pixels: &[Tensor],
image_token_id: u32,
chunk_size: usize,
) -> candle_core::Result<Tensor> {
if image_pixels.is_empty() {
candle_core::bail!("prefill_with_images_chunked: called with zero images");
}
if tokens.is_empty() {
candle_core::bail!("prefill_with_images_chunked: empty prompt");
}
let chunk_size = chunk_size.max(1);
let device = self.device().clone();
let image_embeds = self.encode_images_concat(image_pixels)?;
// Each image's LM grid (lm_gh, lm_gw) = (h/factor, w/factor),
// factor = patch×merge. Recomputed per rank from this rank's own
// pixel tensors — deterministic, so every rank's grids (and hence
// M-RoPE positions) match without crossing the RPC (#14).
let factor = self
.vision
.as_ref()
.map(|v| {
let c = v.config();
c.patch_size * c.spatial_merge_size
})
.ok_or_else(|| {
candle_core::Error::Msg(
"prefill_with_images_chunked: loaded without a vision tower".into(),
)
})?;
let grids: Vec<(usize, usize)> = image_pixels
.iter()
.map(|t| {
let (_, h, w) = t.dims3()?;
Ok::<(usize, usize), candle_core::Error>((h / factor, w / factor))
})
.collect::<candle_core::Result<Vec<_>>>()?;
// Interleaved-M-RoPE 3D position ids for the whole prompt,
// computed once and sliced per chunk so every rank assigns image
// tokens their grid coordinates (and text after an image resumes
// from the compressed counter). `rope_delta` is stored on the base
// model for the decode that follows this prefill. Every chunk —
// text or image — uses the M-RoPE slice, because each image shifts
// the positions of the text around it.
let (text, height, width, delta) =
crate::harness::arch::qwen3_5::rope::get_rope_index(tokens, image_token_id, &grids)
.map_err(|e| candle_core::Error::Msg(format!("get_rope_index: {e}")))?;
self.base.set_rope_delta(delta);
let full_pos = crate::harness::arch::qwen3_5::rope::mrope_position_tensor(
&text, &height, &width, &device,
)?;
let mut last_logits: Option<Tensor> = None;
// Rows of `image_embeds` already spliced by earlier chunks. The
// `<|image_pad|>` run is contiguous, so chunks consume embedding
// rows in order.
let mut img_off = 0usize;
let mut start = 0usize;
while start < tokens.len() {
let end = (start + chunk_size).min(tokens.len());
let chunk = &tokens[start..end];
let input = Tensor::new(chunk, &device)?.unsqueeze(0)?;
let pos_slice = full_pos.narrow(1, start, end - start)?;
let n_here = chunk.iter().filter(|&&t| t == image_token_id).count();
let logits = if n_here == 0 {
self.forward_with_positions(&input, base_offset + start, &pos_slice, None, None)?
} else {
// Splice the next `n_here` patch rows at this chunk's
// local image-pad positions.
let rows = image_embeds.narrow(0, img_off, n_here)?;
img_off += n_here;
self.forward_with_positions(
&input,
base_offset + start,
&pos_slice,
Some(&rows),
Some(image_token_id),
)?
};
last_logits = Some(logits);
start = end;
}
last_logits
.ok_or_else(|| candle_core::Error::Msg("prefill_with_images_chunked: no chunks".into()))
}
pub fn clear_kv_cache(&mut self) {
self.base.clear_kv_cache();
}
/// See [`TpQwen3_5Model::snapshot_kv_cache`].
pub fn snapshot_kv_cache(&self) -> candle_core::Result<KvCacheSnapshot> {
self.base.snapshot_kv_cache()
}
/// See [`TpQwen3_5Model::restore_kv_cache`].
pub fn restore_kv_cache(&mut self, snap: &KvCacheSnapshot) -> candle_core::Result<()> {
self.base.restore_kv_cache(snap)
}
pub fn device(&self) -> &Device {
&self.base.device
}
@@ -1250,12 +1539,20 @@ fn build_lm_head(
} else {
// lm_head sits at the top level (sibling of `model.*`), NOT
// under `model.language_model`.
let lm_head_start = std::time::Instant::now();
let weight = load_replicated(
&vb.pp("lm_head"),
(cfg.vocab_size, cfg.hidden_size),
"weight",
)?;
super::tp_linear::MaybeQuantLinear::from_weight(weight, quant).context("wrap lm_head")
let head = super::tp_linear::MaybeQuantLinear::from_weight(weight, quant)
.context("wrap lm_head")?;
tracing::info!(
elapsed_ms = lm_head_start.elapsed().as_millis() as u64,
quantized = quant.is_some(),
"lm_head loaded"
);
Ok(head)
}
}

View File

@@ -47,24 +47,30 @@ impl WorkerModel {
}
}
/// Image-bearing forward on this rank. Only the vision-capable
/// Chunked image prefill on this rank. Only the vision-capable
/// `qwen3_5` arch has a replicated tower; the dense `qwen3` arch
/// errors. The returned logits are discarded by the caller (the
/// leader samples from its own rank-0 copy) — the value is the NCCL
/// collectives the forward issues.
fn forward_with_images(
/// collectives the forward issues, chunk by chunk in lockstep with
/// the leader.
fn prefill_with_images_chunked(
&mut self,
input: &candle_core::Tensor,
offset: usize,
tokens: &[u32],
base_offset: usize,
image_pixels: &[candle_core::Tensor],
image_token_id: u32,
chunk_size: usize,
) -> candle_core::Result<candle_core::Tensor> {
match self {
WorkerModel::Qwen3_5(m) => {
m.forward_with_images(input, offset, image_pixels, image_token_id)
}
WorkerModel::Qwen3_5(m) => m.prefill_with_images_chunked(
tokens,
base_offset,
image_pixels,
image_token_id,
chunk_size,
),
WorkerModel::Qwen3(_) => {
candle_core::bail!("forward_with_images: qwen3 (dense) has no vision tower")
candle_core::bail!("prefill_with_images_chunked: qwen3 (dense) has no vision tower")
}
}
}
@@ -76,6 +82,33 @@ impl WorkerModel {
}
}
/// Capture this rank's cache state for a prefix snapshot (#11).
/// Only qwen3_5 exposes its state; the dense qwen3 arch errors —
/// the leader never asks, because it gates on its own
/// `TpLeaderModel` support first.
fn snapshot_kv_cache(
&self,
) -> candle_core::Result<crate::harness::arch::qwen3_5::snapshot::KvCacheSnapshot> {
match self {
WorkerModel::Qwen3(_) => {
candle_core::bail!("snapshot_kv_cache: qwen3 (dense) has no snapshot support")
}
WorkerModel::Qwen3_5(m) => m.snapshot_kv_cache(),
}
}
fn restore_kv_cache(
&mut self,
snap: &crate::harness::arch::qwen3_5::snapshot::KvCacheSnapshot,
) -> candle_core::Result<()> {
match self {
WorkerModel::Qwen3(_) => {
candle_core::bail!("restore_kv_cache: qwen3 (dense) has no snapshot support")
}
WorkerModel::Qwen3_5(m) => m.restore_kv_cache(snap),
}
}
fn device(&self) -> &candle_core::Device {
match self {
WorkerModel::Qwen3(m) => m.device(),
@@ -158,6 +191,16 @@ struct WorkerState {
#[cfg(not(feature = "cuda"))]
#[allow(dead_code)]
models: HashMap<String, ()>,
/// Prefix-cache snapshots (#11) for this rank's shards, keyed by
/// `(model_id, snapshot_id)` — the id is minted by the leader's
/// pool and shared across all ranks. Dropped with the shard on
/// `UnloadModel`.
#[cfg(feature = "cuda")]
kv_snapshots: HashMap<(String, u64), crate::harness::arch::qwen3_5::snapshot::KvCacheSnapshot>,
/// Placeholder mirroring `models` on the non-cuda build.
#[cfg(not(feature = "cuda"))]
#[allow(dead_code)]
kv_snapshots: HashMap<(String, u64), ()>,
}
impl WorkerState {
@@ -166,6 +209,7 @@ impl WorkerState {
config,
nccl: NcclState::new(),
models: HashMap::new(),
kv_snapshots: HashMap::new(),
}
}
@@ -195,14 +239,28 @@ impl WorkerState {
offset,
image_token_id,
image_data_uris,
chunk_size,
} => self.handle_generate_step_with_images(
&model_id,
tokens,
offset,
image_token_id,
image_data_uris,
chunk_size,
),
WorkerRequest::ClearKvCache { model_id } => self.handle_clear_kv_cache(&model_id),
WorkerRequest::SnapshotKvCache {
model_id,
snapshot_id,
} => self.handle_snapshot_kv_cache(&model_id, snapshot_id),
WorkerRequest::RestoreKvCache {
model_id,
snapshot_id,
} => self.handle_restore_kv_cache(&model_id, snapshot_id),
WorkerRequest::DropKvSnapshot {
model_id,
snapshot_id,
} => self.handle_drop_kv_snapshot(&model_id, snapshot_id),
WorkerRequest::UnloadModel { model_id } => self.handle_unload_model(&model_id),
WorkerRequest::Shutdown => WorkerResponse::Bye,
}
@@ -466,6 +524,7 @@ impl WorkerState {
offset: usize,
image_token_id: u32,
image_data_uris: Vec<String>,
chunk_size: usize,
) -> WorkerResponse {
use crate::harness::preprocess::{PreprocessProfile, preprocess_data_uri};
use candle_core::Tensor;
@@ -485,16 +544,13 @@ impl WorkerState {
let device = model.device().clone();
// Preprocess each image identically to the leader so the encoded
// embeddings — and thus the spliced hidden state — match across
// ranks. Fixed 448×448 profile.
// embeddings — and thus the spliced hidden state and per-image
// grids — match across ranks. Native-aspect `smart_resize` (#14);
// deterministic, so each rank derives the same dims.
let profile = PreprocessProfile::qwen3_6();
let (h, w) = (
profile.target_height as usize,
profile.target_width as usize,
);
let mut pixels: Vec<Tensor> = Vec::with_capacity(image_data_uris.len());
for (idx, uri) in image_data_uris.iter().enumerate() {
let px = match preprocess_data_uri(uri, &profile) {
let (px, h, w) = match preprocess_data_uri(uri, &profile) {
Ok(p) => p,
Err(e) => {
return WorkerResponse::Error {
@@ -503,7 +559,7 @@ impl WorkerState {
};
}
};
match Tensor::from_vec(px, (3, h, w), &device) {
match Tensor::from_vec(px, (3, h as usize, w as usize), &device) {
Ok(t) => pixels.push(t),
Err(e) => {
return WorkerResponse::Error {
@@ -514,16 +570,6 @@ impl WorkerState {
}
}
let input = match Tensor::new(tokens.as_slice(), &device).and_then(|t| t.unsqueeze(0)) {
Ok(t) => t,
Err(e) => {
return WorkerResponse::Error {
kind: "forward_failed".into(),
message: format!("build input tensor: {e}"),
};
}
};
let start = std::time::Instant::now();
tracing::debug!(
rank = self.config.rank,
@@ -531,10 +577,14 @@ impl WorkerState {
tokens = tokens.len(),
offset,
images = pixels.len(),
"worker GenerateStepWithImages: forward starting"
chunk_size,
"worker GenerateStepWithImages: chunked prefill starting"
);
// Drop the logits — the leader samples from its own rank-0 copy.
if let Err(e) = model.forward_with_images(&input, offset, &pixels, image_token_id) {
// The chunked prefill builds its own per-chunk input tensors.
if let Err(e) =
model.prefill_with_images_chunked(&tokens, offset, &pixels, image_token_id, chunk_size)
{
tracing::warn!(
rank = self.config.rank,
model = %model_id,
@@ -564,6 +614,7 @@ impl WorkerState {
_offset: usize,
_image_token_id: u32,
_image_data_uris: Vec<String>,
_chunk_size: usize,
) -> WorkerResponse {
WorkerResponse::Error {
kind: "cuda_feature_not_enabled".into(),
@@ -591,6 +642,99 @@ impl WorkerState {
}
}
#[cfg(feature = "cuda")]
fn handle_snapshot_kv_cache(&mut self, model_id: &str, snapshot_id: u64) -> WorkerResponse {
let Some(model) = self.models.get(model_id) else {
return WorkerResponse::Error {
kind: "model_not_loaded".into(),
message: format!("model '{model_id}' not loaded on rank {}", self.config.rank),
};
};
match model.snapshot_kv_cache() {
Ok(snap) => {
let bytes = snap.size_bytes();
self.kv_snapshots
.insert((model_id.to_string(), snapshot_id), snap);
tracing::debug!(
rank = self.config.rank,
model = %model_id,
snapshot_id,
bytes,
stored = self.kv_snapshots.len(),
"kv snapshot captured"
);
WorkerResponse::KvSnapshotStored { bytes }
}
Err(e) => WorkerResponse::Error {
kind: "snapshot_failed".into(),
message: format!("snapshot_kv_cache: {e}"),
},
}
}
#[cfg(not(feature = "cuda"))]
fn handle_snapshot_kv_cache(&mut self, _model_id: &str, _snapshot_id: u64) -> WorkerResponse {
WorkerResponse::Error {
kind: "cuda_feature_not_enabled".into(),
message: "SnapshotKvCache requires --features cuda".into(),
}
}
#[cfg(feature = "cuda")]
fn handle_restore_kv_cache(&mut self, model_id: &str, snapshot_id: u64) -> WorkerResponse {
let key = (model_id.to_string(), snapshot_id);
let (Some(model), Some(snap)) =
(self.models.get_mut(model_id), self.kv_snapshots.get(&key))
else {
return WorkerResponse::Error {
kind: "snapshot_not_found".into(),
message: format!(
"model '{model_id}' / snapshot {snapshot_id} not present on rank {}",
self.config.rank
),
};
};
match model.restore_kv_cache(snap) {
Ok(()) => WorkerResponse::KvCacheRestored,
Err(e) => WorkerResponse::Error {
kind: "restore_failed".into(),
message: format!("restore_kv_cache: {e}"),
},
}
}
#[cfg(not(feature = "cuda"))]
fn handle_restore_kv_cache(&mut self, _model_id: &str, _snapshot_id: u64) -> WorkerResponse {
WorkerResponse::Error {
kind: "cuda_feature_not_enabled".into(),
message: "RestoreKvCache requires --features cuda".into(),
}
}
#[cfg(feature = "cuda")]
fn handle_drop_kv_snapshot(&mut self, model_id: &str, snapshot_id: u64) -> WorkerResponse {
let was_present = self
.kv_snapshots
.remove(&(model_id.to_string(), snapshot_id))
.is_some();
tracing::debug!(
rank = self.config.rank,
model = %model_id,
snapshot_id,
was_present,
stored = self.kv_snapshots.len(),
"kv snapshot dropped"
);
WorkerResponse::KvSnapshotDropped
}
#[cfg(not(feature = "cuda"))]
fn handle_drop_kv_snapshot(&mut self, _model_id: &str, _snapshot_id: u64) -> WorkerResponse {
// Dropping is bookkeeping-only; reply success so leader-side
// eviction never wedges on a no-cuda build.
WorkerResponse::KvSnapshotDropped
}
#[cfg(feature = "cuda")]
fn handle_unload_model(&mut self, model_id: &str) -> WorkerResponse {
if self.models.remove(model_id).is_none() {
@@ -599,6 +743,9 @@ impl WorkerState {
message: format!("model '{model_id}' not loaded on rank {}", self.config.rank),
};
}
// Snapshots are scoped to the shard — a reloaded model gets a
// fresh prefix cache, so stale ids must not resurrect tensors.
self.kv_snapshots.retain(|(m, _), _| m != model_id);
tracing::info!(rank = self.config.rank, model = %model_id, "unloaded TP shard");
WorkerResponse::Unloaded
}

View File

@@ -6,4 +6,5 @@ pub mod discovery;
pub mod harness;
pub mod health;
pub mod startup;
pub mod version;
pub mod wire;

View File

@@ -20,6 +20,7 @@ use tracing_subscriber::EnvFilter;
#[command(name = "neuron")]
#[command(about = "Per-node daemon for cortex inference clusters")]
#[command(version)]
#[command(long_version = neuron::version::long_version_static())]
struct Args {
/// Run in tensor-parallel worker mode. The leader process spawns
/// one of these per non-zero NCCL rank and drives it over
@@ -170,6 +171,14 @@ async fn daemon(args: Args) -> Result<()> {
devices = discovery_result.devices.len(),
"discovery complete"
);
// Driver/library mismatch preflight (#19): make the un-rebooted
// driver-update failure mode instantly legible at startup instead
// of a cryptic nccl_init_failed minutes later inside the first
// model load. One loud line; the reason also rides on /discovery
// so cortex can route around this node.
if let Some(reason) = &discovery_result.cuda_unavailable_reason {
tracing::error!(reason = %reason, "CUDA UNAVAILABLE on this host");
}
// Build harness registry from config. In-process harnesses (candle)
// need to know neuron's own bind URL so they can return it from
@@ -223,7 +232,15 @@ async fn daemon(args: Args) -> Result<()> {
// mutex, so concurrent on-demand loads and pre-warm loads
// do not race on the same model.
let registry = state_for_prewarm.registry.read().await;
startup::load_default_models(&registry, &default_models, &state_for_prewarm.activation)
startup::load_default_models(
&registry,
&default_models,
&state_for_prewarm.activation,
state_for_prewarm
.discovery
.cuda_unavailable_reason
.as_deref(),
)
.await;
});
}

View File

@@ -33,11 +33,31 @@ pub async fn load_default_models(
registry: &HarnessRegistry,
specs: &[ModelSpec],
activation: &ActivationTracker,
cuda_unavailable_reason: Option<&str>,
) {
if specs.is_empty() {
activation.mark_ready().await;
return;
}
// Driver/library mismatch preflight (#19): every CUDA load on this
// host is guaranteed to fail (cuInit → CUDA_ERROR_SYSTEM_DRIVER
// MISMATCH, surfacing as a cryptic NCCL/driver error). Don't
// attempt them — mark each default model failed with the
// operator-actionable reason so `/health` activation shows the
// real cause, and let the host run API-only until it's rebooted.
if let Some(reason) = cuda_unavailable_reason {
tracing::error!(
count = specs.len(),
reason = %reason,
"skipping default model loads: CUDA unavailable"
);
for spec in specs {
activation.start_loading(&spec.model_id).await;
activation.fail_loading(&spec.model_id, reason).await;
}
activation.mark_ready().await;
return;
}
tracing::info!(count = specs.len(), "loading default models");
for spec in specs {
let start = Instant::now();

View File

@@ -0,0 +1,83 @@
//! The daemon's own build identity, captured at compile time by
//! `build.rs` and served from `GET /version`.
//!
//! The `env!()` reads below resolve to the `cargo:rustc-env=` values
//! emitted by `build.rs::emit_build_metadata`. When neuron is built
//! from a source tarball with no git metadata and no injected
//! `HELEXA_BUILD_SHA`, `HELEXA_GIT_SHA` is the literal `"unknown"`.
use cortex_core::build_info::BuildInfo;
/// Assemble the compiled-in build metadata into a [`BuildInfo`].
pub fn build_info() -> BuildInfo {
BuildInfo {
package_version: env!("CARGO_PKG_VERSION").to_string(),
git_sha: env!("HELEXA_GIT_SHA").to_string(),
git_sha_long: non_empty(env!("HELEXA_GIT_SHA_LONG")),
git_dirty: env!("HELEXA_GIT_DIRTY") == "true",
build_timestamp: non_empty(env!("HELEXA_BUILD_TIMESTAMP")),
rustc_version: non_empty(env!("HELEXA_RUSTC_VERSION")),
profile: non_empty(env!("HELEXA_BUILD_PROFILE")),
target: non_empty(env!("HELEXA_TARGET")),
features: split_features(env!("HELEXA_FEATURES")),
candle_version: non_empty(env!("HELEXA_CANDLE_VERSION")),
}
}
/// A one-line version string for clap's `--version` long form, as a
/// `&'static str` (clap requires `'static`). Computed once.
pub fn long_version_static() -> &'static str {
static V: std::sync::OnceLock<String> = std::sync::OnceLock::new();
V.get_or_init(long_version).as_str()
}
/// A one-line version string for clap's `--version` long form.
pub fn long_version() -> String {
let info = build_info();
let dirty = if info.git_dirty { "-dirty" } else { "" };
let features = if info.features.is_empty() {
String::new()
} else {
format!(" [{}]", info.features.join(","))
};
format!(
"{} ({}{}){}",
info.package_version, info.git_sha, dirty, features
)
}
fn non_empty(s: &str) -> Option<String> {
if s.is_empty() {
None
} else {
Some(s.to_string())
}
}
fn split_features(s: &str) -> Vec<String> {
s.split(',')
.map(str::trim)
.filter(|f| !f.is_empty())
.map(str::to_string)
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn build_info_is_populated() {
let info = build_info();
// Always present regardless of git availability.
assert_eq!(info.package_version, env!("CARGO_PKG_VERSION"));
assert!(!info.git_sha.is_empty());
}
#[test]
fn long_version_includes_sha() {
let v = long_version();
assert!(v.contains(env!("CARGO_PKG_VERSION")));
assert!(v.contains(env!("HELEXA_GIT_SHA")));
}
}

View File

@@ -40,7 +40,7 @@ async fn test_load_default_models_skips_unknown_harness() {
];
let activation = ActivationTracker::new(&specs);
startup::load_default_models(&registry, &specs, &activation).await;
startup::load_default_models(&registry, &specs, &activation, None).await;
let listed = registry
.list_all_models()
@@ -71,7 +71,52 @@ async fn test_load_default_models_skips_unknown_harness() {
async fn test_load_default_models_empty_is_noop() {
let registry = HarnessRegistry::new();
let activation = ActivationTracker::new(&[]);
startup::load_default_models(&registry, &[], &activation).await;
startup::load_default_models(&registry, &[], &activation, None).await;
let snapshot = activation.snapshot().await;
assert_eq!(snapshot.state, ActivationState::Ready);
}
#[tokio::test]
async fn test_load_default_models_skipped_on_driver_mismatch() {
// #19: when the host has a driver/library mismatch, no load is
// attempted (it would die in cuInit/NCCL with a cryptic error);
// every default model lands in `failed` carrying the actionable
// reason, and the tracker still flips to ready so /health serves.
let registry = HarnessRegistry::from_configs(
&[HarnessConfig {
name: "candle".into(),
}],
"http://localhost:0",
&HarnessSettings::default(),
);
let specs = vec![ModelSpec {
model_id: "Qwen/Qwen3.6-27B".into(),
harness: "candle".into(),
quant: Some("q6k".into()),
tensor_parallel: Some(2),
devices: None,
}];
let activation = ActivationTracker::new(&specs);
let reason = "host NVIDIA driver/library mismatch (userspace NVML 580.159 vs loaded \
kernel module 580.159.03) — reboot the host to reload the kernel module; \
all CUDA inference is unavailable until then";
startup::load_default_models(&registry, &specs, &activation, Some(reason)).await;
let listed = registry
.list_all_models()
.await
.expect("list_all_models should succeed");
assert!(
listed.is_empty(),
"no load may be attempted on a mismatch host"
);
let snapshot = activation.snapshot().await;
assert_eq!(snapshot.state, ActivationState::Ready);
assert_eq!(snapshot.failed.len(), 1);
assert!(
snapshot.failed[0].error.contains("driver/library mismatch"),
"failure must carry the actionable reason, got: {}",
snapshot.failed[0].error
);
}

View File

@@ -50,6 +50,7 @@ fn fake_discovery() -> DiscoveryResponse {
},
],
harnesses: vec![],
cuda_unavailable_reason: None,
}
}
@@ -76,6 +77,27 @@ async fn test_discovery_endpoint() {
assert_eq!(devices[0]["vram_total_mb"], 32614);
}
#[tokio::test]
async fn test_version_endpoint() {
let url = spawn_neuron(fake_discovery()).await;
let client = reqwest::Client::new();
let resp = client
.get(format!("{url}/version"))
.send()
.await
.expect("request should succeed");
assert_eq!(resp.status(), 200);
// Deserialize into the shared type to lock the wire contract.
let body: cortex_core::build_info::BuildInfo = resp.json().await.unwrap();
assert_eq!(body.package_version, env!("CARGO_PKG_VERSION"));
// git_sha is always present — a real short SHA in a git checkout, or
// the literal "unknown" in a tarball build. Either way, non-empty.
assert!(!body.git_sha.is_empty());
}
#[tokio::test]
async fn test_health_endpoint() {
let url = spawn_neuron(fake_discovery()).await;
@@ -103,6 +125,7 @@ async fn test_discovery_no_gpus() {
driver_version: None,
devices: vec![],
harnesses: vec![],
cuda_unavailable_reason: None,
};
let url = spawn_neuron(disc).await;
@@ -487,3 +510,74 @@ async fn test_responses_streaming_model_not_loaded() {
.unwrap();
assert_eq!(resp.status(), 404);
}
#[tokio::test]
async fn test_driver_mismatch_rejects_load_and_rides_discovery() {
// #19: a host with the driver/library mismatch advertises the
// reason on /discovery (so cortex routes around it) and fast-
// rejects /models/load with 503 + the actionable message instead
// of dying minutes later inside cuInit/NCCL.
let reason = "host NVIDIA driver/library mismatch (userspace NVML 580.159 vs loaded \
kernel module 580.159.03) — reboot the host to reload the kernel module; \
all CUDA inference is unavailable until then";
let disc = DiscoveryResponse {
hostname: "mismatched".into(),
os: "Linux".into(),
kernel: "6.19.0".into(),
cuda_version: Some("13.0".into()),
driver_version: None,
devices: vec![],
harnesses: vec!["candle".into()],
cuda_unavailable_reason: Some(reason.into()),
};
let url = spawn_neuron(disc).await;
let client = reqwest::Client::new();
let body: serde_json::Value = client
.get(format!("{url}/discovery"))
.send()
.await
.expect("discovery request")
.json()
.await
.unwrap();
assert_eq!(body["cuda_unavailable_reason"], reason);
let resp = client
.post(format!("{url}/models/load"))
.json(&serde_json::json!({
"model_id": "Qwen/Qwen3.6-27B",
"harness": "candle",
"quant": "q6k",
"tensor_parallel": 2
}))
.send()
.await
.expect("load request");
assert_eq!(resp.status(), 503);
let body: serde_json::Value = resp.json().await.unwrap();
assert_eq!(body["code"], "cuda_unavailable");
assert!(
body["error"].as_str().unwrap().contains("reboot the host"),
"error must be operator-actionable: {body}"
);
}
#[tokio::test]
async fn test_healthy_discovery_omits_cuda_unavailable_reason() {
// No false positives: the field must be absent (not null) from the
// wire format on healthy hosts.
let url = spawn_neuron(fake_discovery()).await;
let body: serde_json::Value = reqwest::Client::new()
.get(format!("{url}/discovery"))
.send()
.await
.expect("discovery request")
.json()
.await
.unwrap();
assert!(
body.get("cuda_unavailable_reason").is_none(),
"healthy host must omit the field entirely: {body}"
);
}

View File

@@ -0,0 +1,67 @@
# Numerical-reference fixtures (#15)
Reference tensors captured from the HF `transformers` implementation
by [`script/dump_reference.py`](../../../../../script/dump_reference.py),
replayed and compared by
[`tests/numerical_reference.rs`](../../numerical_reference.rs). These
pin the README's "implemented in this repository, ported against the
HuggingFace reference" claim to checked-in numbers.
| fixture | model | case | dtype | compared by |
|---|---|---|---|---|
| `qwen3_5-0.8b-text` | Qwen/Qwen3.5-0.8B | text (>64-token prompt → chunked GDN prefill) | f32 | `text_logits_match_reference` |
| `qwen3_5-0.8b-vision` | Qwen/Qwen3.5-0.8B | 448×448 synthetic image + prompt | f32 | `vision_tower_and_logits_match_reference` |
| `qwen3_6-27b-text` | Qwen/Qwen3.6-27B | text | bf16 | manual (see below) |
## Running the comparison
On a host with the model snapshot (beast):
```sh
NEURON_REF_MODEL_PATH=/archive3/llm-cache/models--Qwen--Qwen3.5-0.8B/snapshots/<rev> \
cargo test -p neuron --test numerical_reference -- --nocapture
```
Without `NEURON_REF_MODEL_PATH` the tests compile and self-skip, so CI
stays green without weights.
## Why f32 fixtures
f32-vs-f32 isolates implementation differences: observed agreement is
text max_abs 0.000 / cosine 1.000000, vision tower cosine 0.999998.
Cross-dtype comparisons drown in bf16 rounding chaos through the
27-layer tower (global cosine ~0.997, worst patch ~0.92, worst index
unstable across runs) — that is production-dtype noise, not
implementation error. The mutation check: rerunning with
`NEURON_VISION_LEGACY_POS=1` (the deliberately-wrong sequential
pos-embed lookup) collapses tower cosine to 0.75 / worst patch 0.28
and fails the test loudly.
## The 27B fixture
`qwen3_6-27b-text` is captured in bf16 on CPU (an f32 27B forward
needs ~108 GB; beast has 91 GB free). The automated tests run against
the 0.8B because both models execute the *same* arch modules — the
27B differs only in hyperparameters — and an apples-to-apples 27B
replay needs either TP=2 bf16 (idle GPUs, no neuron running) or a
bigger-RAM host. Manual procedure when wanted: stop neuron on beast,
replay the manifest's token ids through a TP=2 bf16 load, compare
argmax + cosine against `logits.f32` with bf16-calibrated tolerances.
## Regenerating
Regenerate whenever the pinned snapshot or the transformers reference
changes; record both versions (in each `manifest.json`) in the commit
message:
```sh
# on beast; processor files may be missing from neuron's snapshot —
# point the processor at the repo id with a scratch cache
SNAP=$(ls -d /archive3/llm-cache/models--Qwen--Qwen3.5-0.8B/snapshots/*/ | head -1)
HF_HUB_CACHE=/tmp/hf-ref-cache python3 script/dump_reference.py \
--model-path "$SNAP" --processor-path Qwen/Qwen3.5-0.8B \
--case text --out crates/neuron/tests/fixtures/numerical/qwen3_5-0.8b-text
HF_HUB_CACHE=/tmp/hf-ref-cache python3 script/dump_reference.py \
--model-path "$SNAP" --processor-path Qwen/Qwen3.5-0.8B \
--case vision --out crates/neuron/tests/fixtures/numerical/qwen3_5-0.8b-vision
```

Binary file not shown.

View File

@@ -0,0 +1,143 @@
{
"model_path": "/archive3/llm-cache/models--Qwen--Qwen3.5-0.8B/snapshots/2fc06364715b967f1860aea9cf38778875588b17/",
"case": "text",
"transformers_version": "5.9.0",
"torch_version": "2.9.1+cu128",
"files": {
"logits": {
"file": "logits.f32",
"shape": [
248320
]
}
},
"dtype": "float32",
"prompt": "The helexa fleet serves near-frontier language models on consumer graphics cards. Each host runs a small daemon that discovers its hardware, loads the configured models, and answers OpenAI-compatible requests over the private mesh network. The gateway routes each request to the host that already holds the model, restores any cached prefix state, and streams the generated tokens back to the caller one chunk at a time. Operators care about three numbers: the time to the first token, the steady decode rate, and the time a cold model takes to become ready after a deploy. This paragraph exists only to be tokenized identically by two implementations.",
"token_ids": [
760,
551,
2486,
64,
24303,
16545,
3043,
61478,
1223,
3992,
3983,
383,
11171,
13775,
7176,
13,
8618,
3357,
8213,
264,
2526,
37993,
421,
49296,
1141,
11436,
11,
20269,
279,
19152,
3983,
11,
321,
10926,
5097,
15015,
77450,
7154,
888,
279,
843,
10967,
3790,
13,
561,
27853,
10964,
1754,
1622,
310,
279,
3357,
421,
2582,
9687,
279,
1558,
11,
84728,
866,
19954,
8978,
1528,
11,
321,
22327,
279,
7658,
10885,
1142,
310,
279,
19260,
799,
11540,
506,
264,
854,
13,
62244,
2373,
883,
2250,
4947,
25,
279,
854,
310,
279,
1118,
3817,
11,
279,
23271,
16401,
4238,
11,
321,
279,
854,
264,
8981,
1558,
4829,
310,
3512,
5354,
1238,
264,
10204,
13,
1061,
13901,
6513,
1132,
310,
381,
3817,
1452,
3408,
2586,
539,
1330,
37066,
13
]
}

Binary file not shown.

After

Width:  |  Height:  |  Size: 72 KiB

View File

@@ -0,0 +1,247 @@
{
"model_path": "/archive3/llm-cache/models--Qwen--Qwen3.5-0.8B/snapshots/2fc06364715b967f1860aea9cf38778875588b17/",
"case": "vision",
"transformers_version": "5.9.0",
"torch_version": "2.9.1+cu128",
"files": {
"visual_out": {
"file": "visual_out.f32",
"shape": [
196,
1024
]
},
"logits": {
"file": "logits.f32",
"shape": [
248320
]
}
},
"dtype": "float32",
"prompt": "Describe this image in one sentence.",
"token_ids": [
248045,
846,
198,
248053,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248056,
248054,
72240,
411,
2099,
303,
799,
11316,
13,
248046,
198,
248045,
74455,
198,
248068,
271,
248069,
271
],
"image_grid_thw": [
1,
28,
28
]
}

Binary file not shown.

View File

@@ -0,0 +1,143 @@
{
"model_path": "/archive3/llm-cache/models--Qwen--Qwen3.6-27B/snapshots/6a9e13bd6fc8f0983b9b99948120bc37f49c13e9/",
"case": "text",
"transformers_version": "5.9.0",
"torch_version": "2.9.1+cu128",
"files": {
"logits": {
"file": "logits.f32",
"shape": [
248320
]
}
},
"dtype": "bfloat16",
"prompt": "The helexa fleet serves near-frontier language models on consumer graphics cards. Each host runs a small daemon that discovers its hardware, loads the configured models, and answers OpenAI-compatible requests over the private mesh network. The gateway routes each request to the host that already holds the model, restores any cached prefix state, and streams the generated tokens back to the caller one chunk at a time. Operators care about three numbers: the time to the first token, the steady decode rate, and the time a cold model takes to become ready after a deploy. This paragraph exists only to be tokenized identically by two implementations.",
"token_ids": [
760,
551,
2486,
64,
24303,
16545,
3043,
61478,
1223,
3992,
3983,
383,
11171,
13775,
7176,
13,
8618,
3357,
8213,
264,
2526,
37993,
421,
49296,
1141,
11436,
11,
20269,
279,
19152,
3983,
11,
321,
10926,
5097,
15015,
77450,
7154,
888,
279,
843,
10967,
3790,
13,
561,
27853,
10964,
1754,
1622,
310,
279,
3357,
421,
2582,
9687,
279,
1558,
11,
84728,
866,
19954,
8978,
1528,
11,
321,
22327,
279,
7658,
10885,
1142,
310,
279,
19260,
799,
11540,
506,
264,
854,
13,
62244,
2373,
883,
2250,
4947,
25,
279,
854,
310,
279,
1118,
3817,
11,
279,
23271,
16401,
4238,
11,
321,
279,
854,
264,
8981,
1558,
4829,
310,
3512,
5354,
1238,
264,
10204,
13,
1061,
13901,
6513,
1132,
310,
381,
3817,
1452,
3408,
2586,
539,
1330,
37066,
13
]
}

View File

@@ -0,0 +1,339 @@
//! Numerical validation against the HF transformers reference (#15).
//!
//! Replays the fixtures captured by `script/dump_reference.py` (token
//! ids, image, reference tensors) through neuron's own qwen3_5
//! implementation and compares. This is what pins the README's
//! "implemented in this repository, ported against the HuggingFace
//! reference" claim to numbers.
//!
//! Needs the model weights on disk, so it self-skips unless
//! `NEURON_REF_MODEL_PATH` points at the HF snapshot directory the
//! fixtures were captured from (see each fixture's `manifest.json`).
//! Run on a host with the snapshot (CUDA used when available):
//!
//! ```sh
//! NEURON_REF_MODEL_PATH=/path/to/models--Qwen--Qwen3.5-0.8B/snapshots/<rev> \
//! cargo test -p neuron --test numerical_reference -- --nocapture
//! ```
//!
//! The text prompt is longer than 64 tokens on purpose: the replay
//! prefill goes through the chunked delta-rule path, so the
//! comparison validates the production prefill math, not just the
//! per-token recurrence.
//!
//! Fixtures are captured in **f32** (script default) so the
//! comparison pins the math itself: observed f32-vs-f32 agreement is
//! text max_abs 0.000 / cosine 1.000000 and vision-tower cosine
//! 0.999998 (worst patch 0.99994), so the thresholds below sit far
//! above noise and far below any real bug (a wrong RoPE base, a
//! missing projector bias, an off-by-one in position handling).
//! For context: comparing across dtypes is dominated by bf16
//! rounding chaos through the 27-layer tower (global cosine ~0.997,
//! worst patch ~0.92, worst index unstable across dtypes) — that is
//! production-dtype noise, not implementation error, and is why the
//! fixtures are not captured in bf16.
use candle_core::{DType, Device, Tensor};
use serde::Deserialize;
use std::path::{Path, PathBuf};
#[derive(Deserialize)]
struct Manifest {
case: String,
token_ids: Vec<u32>,
#[serde(default)]
image_grid_thw: Option<Vec<usize>>,
files: std::collections::HashMap<String, FileEntry>,
}
#[derive(Deserialize)]
struct FileEntry {
file: String,
shape: Vec<usize>,
}
fn fixture_root() -> PathBuf {
Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/numerical")
}
fn read_f32(path: &Path) -> Vec<f32> {
let bytes = std::fs::read(path).unwrap_or_else(|e| panic!("read {path:?}: {e}"));
assert!(bytes.len().is_multiple_of(4), "truncated f32 file {path:?}");
bytes
.chunks_exact(4)
.map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
.collect()
}
struct Comparison {
max_abs: f32,
cosine: f64,
argmax_ours: usize,
argmax_ref: usize,
}
fn compare(ours: &[f32], reference: &[f32]) -> Comparison {
assert_eq!(ours.len(), reference.len(), "length mismatch");
let mut max_abs = 0f32;
let (mut dot, mut na, mut nb) = (0f64, 0f64, 0f64);
for (&a, &b) in ours.iter().zip(reference) {
max_abs = max_abs.max((a - b).abs());
dot += a as f64 * b as f64;
na += a as f64 * a as f64;
nb += b as f64 * b as f64;
}
let argmax = |xs: &[f32]| {
xs.iter()
.enumerate()
.max_by(|x, y| x.1.total_cmp(y.1))
.map(|(i, _)| i)
.unwrap_or(0)
};
Comparison {
max_abs,
cosine: dot / (na.sqrt() * nb.sqrt()),
argmax_ours: argmax(ours),
argmax_ref: argmax(reference),
}
}
/// bf16 on CUDA (matching production and the reference capture);
/// f32 on CPU, where candle has no bf16 matmul — the comparison
/// tolerances absorb the reference's bf16 rounding either way.
fn load_dtype(device: &Device) -> DType {
if device.is_cuda() {
DType::BF16
} else {
DType::F32
}
}
fn load_model(
model_path: &str,
device: &Device,
) -> neuron::harness::arch::qwen3_5::Qwen3_5ForCausalLM {
let cfg_text =
std::fs::read_to_string(Path::new(model_path).join("config.json")).expect("config.json");
let cfg: neuron::harness::arch::qwen3_5::Config =
serde_json::from_str(&cfg_text).expect("parse config");
let index_text =
std::fs::read_to_string(Path::new(model_path).join("model.safetensors.index.json"));
let paths: Vec<PathBuf> = match index_text {
Ok(text) => {
let v: serde_json::Value = serde_json::from_str(&text).expect("parse index");
let mut names: Vec<String> = v["weight_map"]
.as_object()
.expect("weight_map")
.values()
.filter_map(|x| x.as_str().map(String::from))
.collect();
names.sort();
names.dedup();
names
.into_iter()
.map(|n| Path::new(model_path).join(n))
.collect()
}
Err(_) => vec![Path::new(model_path).join("model.safetensors")],
};
// SAFETY: mmap of read-only snapshot files, same justification as
// the production loader.
let vb = unsafe {
candle_nn::var_builder::ShardedSafeTensors::var_builder(&paths, load_dtype(device), device)
.expect("var_builder")
};
neuron::harness::arch::qwen3_5::Qwen3_5ForCausalLM::new(cfg, vb).expect("build model")
}
fn pick_device() -> Device {
Device::new_cuda(0).unwrap_or(Device::Cpu)
}
fn ref_model_path() -> Option<String> {
match std::env::var("NEURON_REF_MODEL_PATH") {
Ok(p) if !p.is_empty() => Some(p),
_ => {
eprintln!("NEURON_REF_MODEL_PATH unset — skipping numerical reference test");
None
}
}
}
#[test]
fn text_logits_match_reference() {
let Some(model_path) = ref_model_path() else {
return;
};
let fixture = fixture_root().join("qwen3_5-0.8b-text");
let manifest: Manifest =
serde_json::from_str(&std::fs::read_to_string(fixture.join("manifest.json")).unwrap())
.unwrap();
assert_eq!(manifest.case, "text");
let reference = read_f32(&fixture.join(&manifest.files["logits"].file));
let device = pick_device();
let mut model = load_model(&model_path, &device);
let input = Tensor::new(manifest.token_ids.as_slice(), &device)
.unwrap()
.unsqueeze(0)
.unwrap();
// Single full-prompt forward; the prompt is >64 tokens so the
// GDN layers take the chunked prefill path.
let logits = model.forward(&input, 0).unwrap();
let ours: Vec<f32> = logits
.to_dtype(DType::F32)
.unwrap()
.flatten_all()
.unwrap()
.to_vec1()
.unwrap();
let c = compare(&ours, &reference);
eprintln!(
"text: max_abs={:.4} cosine={:.6} argmax ours={} ref={}",
c.max_abs, c.cosine, c.argmax_ours, c.argmax_ref
);
assert_eq!(c.argmax_ours, c.argmax_ref, "argmax token diverged");
assert!(c.cosine > 0.9995, "cosine {:.6} too low", c.cosine);
assert!(c.max_abs < 0.1, "max abs diff {:.4} too high", c.max_abs);
}
#[test]
fn vision_tower_and_logits_match_reference() {
let Some(model_path) = ref_model_path() else {
return;
};
let fixture = fixture_root().join("qwen3_5-0.8b-vision");
let manifest: Manifest =
serde_json::from_str(&std::fs::read_to_string(fixture.join("manifest.json")).unwrap())
.unwrap();
assert_eq!(manifest.case, "vision");
let ref_visual = read_f32(&fixture.join(&manifest.files["visual_out"].file));
let ref_logits = read_f32(&fixture.join(&manifest.files["logits"].file));
let visual_shape = manifest.files["visual_out"].shape.clone();
let device = pick_device();
let model = load_model(&model_path, &device);
let tower = model.vision().expect("model has a vision tower");
let image_token_id = model.image_token_id().expect("image_token_id");
// Same preprocessing path production requests take. The fixture
// image is 448×448 (factor-aligned) so resize is the identity and
// any mismatch below is normalization/patchify/tower math.
let img = image::open(fixture.join("image.png")).expect("open fixture image");
let profile = neuron::harness::preprocess::PreprocessProfile::qwen3_6();
let (pixels, h, w) = neuron::harness::preprocess::preprocess(&img, &profile).unwrap();
let image = Tensor::from_vec(pixels, (3, h as usize, w as usize), &device).unwrap();
let embeds = tower.forward(&image).unwrap();
assert_eq!(
embeds.dims(),
visual_shape.as_slice(),
"tower output shape vs reference"
);
let ours_visual: Vec<f32> = embeds
.to_dtype(DType::F32)
.unwrap()
.flatten_all()
.unwrap()
.to_vec1()
.unwrap();
let cv = compare(&ours_visual, &ref_visual);
// Per-patch cosine: a positional bug (pos-embed interpolation,
// rotary grid, merger order) concentrates error in specific
// patches; dtype noise spreads uniformly.
let hidden = visual_shape[1];
let mut worst = (1.0f64, 0usize);
for (i, (a, b)) in ours_visual
.chunks(hidden)
.zip(ref_visual.chunks(hidden))
.enumerate()
{
let c = compare(a, b);
if c.cosine < worst.0 {
worst = (c.cosine, i);
}
}
eprintln!(
"vision tower: max_abs={:.4} cosine={:.6} worst_patch={} (cosine {:.6})",
cv.max_abs, cv.cosine, worst.1, worst.0
);
assert!(cv.cosine > 0.9995, "tower cosine {:.6} too low", cv.cosine);
assert!(
worst.0 > 0.995,
"patch {} cosine {:.6} — positional divergence",
worst.1,
worst.0
);
// Full LM forward with the splice — the fixture token ids are
// already pad-expanded by the HF processor. The LM grid is the
// post-merge grid: grid_thw / spatial_merge.
let grid = manifest.image_grid_thw.as_ref().expect("grid in manifest");
let lm_grid = (grid[1] / 2, grid[2] / 2);
let mut model = model;
let input = Tensor::new(manifest.token_ids.as_slice(), &device)
.unwrap()
.unsqueeze(0)
.unwrap();
let logits = model
.forward_with_vision(&input, 0, &embeds, image_token_id, &[lm_grid])
.unwrap();
let ours_logits: Vec<f32> = logits
.to_dtype(DType::F32)
.unwrap()
.flatten_all()
.unwrap()
.to_vec1()
.unwrap();
let cl = compare(&ours_logits, &ref_logits);
eprintln!(
"vision logits: max_abs={:.4} cosine={:.6} argmax ours={} ref={}",
cl.max_abs, cl.cosine, cl.argmax_ours, cl.argmax_ref
);
assert_eq!(cl.argmax_ours, cl.argmax_ref, "argmax token diverged");
assert!(cl.cosine > 0.9995, "logits cosine {:.6} too low", cl.cosine);
assert!(cl.max_abs < 0.1, "max abs diff {:.4} too high", cl.max_abs);
// #18: the chunked single-GPU vision prefill must produce the same
// logits as the single-shot path. chunk_size 64 over a 217-token
// prompt forces 4 chunks, and the ~196-token image-pad run spans
// them — exercising the per-chunk splice + img_off accounting and
// the GDN/KV cross-chunk state carry. Comparing to BOTH the
// single-shot result and the HF reference pins chunked == single-
// shot == reference. Re-encodes the image internally (same tower),
// so it takes pixel tensors, not the pre-encoded `embeds`.
model.clear_kv_cache();
let chunked = model
.prefill_with_images_chunked(
manifest.token_ids.as_slice(),
0,
std::slice::from_ref(&image),
image_token_id,
64,
)
.unwrap();
let chunked_logits: Vec<f32> = chunked
.to_dtype(DType::F32)
.unwrap()
.flatten_all()
.unwrap()
.to_vec1()
.unwrap();
let cc = compare(&chunked_logits, &ours_logits);
let cr = compare(&chunked_logits, &ref_logits);
eprintln!(
"vision chunked(64): vs single-shot max_abs={:.4} cosine={:.6}; vs ref argmax={}",
cc.max_abs, cc.cosine, cr.argmax_ours
);
assert_eq!(
cc.argmax_ours, cl.argmax_ours,
"chunked vision argmax diverged from single-shot"
);
assert!(
cc.cosine > 0.99995,
"chunked vs single-shot cosine {:.6} too low — chunking changed the math",
cc.cosine
);
assert_eq!(cr.argmax_ours, cr.argmax_ref, "chunked argmax vs reference");
}

View File

@@ -0,0 +1,3 @@
g helexa-bench - -
u helexa-bench - "helexa-bench harness" /var/lib/helexa-bench /sbin/nologin
m helexa-bench helexa-bench

21
data/helexa-bench.service Normal file
View File

@@ -0,0 +1,21 @@
[Unit]
Description=helexa-bench — continuous version-aware benchmark harness for the neuron fleet
After=network-online.target
Wants=network-online.target
[Service]
Type=simple
ExecStart=/usr/bin/helexa-bench run --config /etc/helexa-bench/helexa-bench.toml
# Long-running sweep loop; restart unconditionally if it ever exits.
Restart=always
RestartSec=10
User=helexa-bench
Group=helexa-bench
# /var/lib/helexa-bench holds the SQLite system-of-record
# (bench.sqlite) and is the service user's $HOME. StateDirectory makes
# systemd create it with the right ownership on start.
StateDirectory=helexa-bench
StateDirectoryMode=0755
[Install]
WantedBy=multi-user.target

154
doc/benchmarks.md Normal file
View File

@@ -0,0 +1,154 @@
# Benchmarks
Batch-1 numbers for the helexa fleet — what one operator at a keyboard
feels. Produced by [`script/bench.py`](../script/bench.py), which works
against any OpenAI-compatible `/v1` endpoint so the same table can be
extended with llama.cpp / Ollama / vLLM columns by pointing it at their
servers (issue #22 tracks adding those baselines).
## Method
- **Workload**: streamed `chat/completions`, one request at a time
(helexa's regime is operators and their agents, not QPS — see
README "What helexa is not").
- **TTFT**: request send → first SSE content chunk. For thinking
models this includes any visible-token delay; the bench prompt
appends Qwen's `/no_think` soft switch so the budget isn't burned
invisibly.
- **decode tok/s**: visible completion tokens over the first→last
chunk window. neuron emits exactly one SSE chunk per generated
token, so the chunk count is engine-truth (streaming
`stream_options.include_usage` is not implemented yet). Reported
only when the window exceeds 200 ms — short coalesced replies don't
produce an honest rate.
- **Prompts**: synthetic filler at ~128 and ~4096 tokens plus a
~300-word generation task (`--max-tokens 256`, temperature 0).
- **Runs**: median of 3 after 1 unmeasured warmup, per cell.
- Requests go through the cortex gateway (`hanzalova:31313`), so
numbers include the proxy hop — the path real clients use. The
gateway also exports the same quantities per-request as Prometheus
histograms (`cortex_time_to_first_token_seconds`,
`cortex_tokens_per_second`, see #21).
## Fleet
| host | GPU(s) | model under test | quant / placement |
|---|---|---|---|
| beast | 2× RTX 5090 (32 GB, cc 12.0) | Qwen/Qwen3.6-27B | Q6K, TP=2 |
| benjy | RTX 4090 (24 GB, cc 8.9) | Qwen/Qwen3-8B | BF16, single GPU |
| quadbrat | RTX 3060 (12 GB, cc 8.6) | Qwen/Qwen3-1.7B | BF16, single GPU |
Driver 580.159, CUDA 13.0, Fedora 43. Models as configured in each
host's `default_models`.
## Results — 2026-06-12 (`8f6f1d3`)
| engine | model | prompt tok | TTFT (s) | decode tok/s | total (s) |
|---|---|---:|---:|---:|---:|
| helexa | Qwen/Qwen3-1.7B | ~128 | 0.685 | 81.3 | 3.741 |
| helexa | Qwen/Qwen3-1.7B | ~4096 | 2.743 | 35.4 | 9.884 |
| helexa | Qwen/Qwen3-8B | ~128 | 0.884 | 62.4 | 4.938 |
| helexa | Qwen/Qwen3-8B | ~4096 | 1.818 | 46.5 | 7.27 |
| helexa | Qwen/Qwen3.6-27B | ~128 | 1.658 | 35.0 | 8.981 |
| helexa | Qwen/Qwen3.6-27B | ~4096 | 7.067 | 33.7 | 14.63 |
Reading the table:
- Long-context decode degradation (81→35 tok/s on the 1.7B) is the
attention cost of a fuller KV cache — expected, and the kind of
number the prefix-KV-cache work (#11) and chunked prefill (#23)
exist to improve at the TTFT end.
- The 27B rows are the headline case: a near-frontier hybrid
linear-attention model decoding at a steady ~35 tok/s on two
consumer cards, with essentially no decode degradation from 128 to
4k context (the Gated DeltaNet recurrent state is O(1) in sequence
length — this is the architecture doing what it promises). The
4k-prompt TTFT (7.1 s) is dominated by the recurrent, non-chunked
delta-rule prefill — issue #23 tracks the fix, and this row is its
before number.
## Results — 2026-06-12, post prefix-KV-cache (#11, `a1952a4`)
| engine | model | prompt tok | TTFT (s) | decode tok/s | total (s) |
|---|---|---:|---:|---:|---:|
| helexa | Qwen/Qwen3-1.7B | ~128 | 0.702 | 104.8 | 1.895 |
| helexa | Qwen/Qwen3-1.7B | ~4096 | 2.749 | 44.9 | 5.534 |
| helexa | Qwen/Qwen3-8B | ~128 | 0.886 | 78.6 | 2.478 |
| helexa | Qwen/Qwen3-8B | ~4096 | 1.824 | 58.3 | 3.969 |
| helexa | Qwen/Qwen3.6-27B | ~128 | 1.355 | 45.8 | 4.147 |
| helexa | Qwen/Qwen3.6-27B | ~4096 | 1.431 | 43.3 | 4.387 |
Reading the table:
- **Methodology note since #11**: neuron now caches cache-state
snapshots per prompt prefix (qwen3_5-arch models only). The bench
repeats one prompt per cell after a warmup, and the snapshot
boundary sits just before the prompt's volatile tail, so the
measured runs hit the cache — qwen3_5 TTFT rows are **warm** TTFT.
The cold number is the warmup run (unchanged from the baseline
table above). For repeated-prefix workloads — agents, chat — warm
is what the operator feels.
- The 27B @4k warm TTFT collapsed 7.07 s → 1.43 s. The controlled
multi-turn measurement (turn N+1 = turn N + a new question, ~5k
context, journal-verified) shows the prefill itself at 8.07 s cold
→ 0.22 s warm with ~98% of the prompt reused — see the closing
numbers on #11.
- The Qwen3 rows (1.7B, 8B) are candle-transformers archs with no
snapshot support — their unchanged TTFT vs the baseline is the
no-regression control. Their decode tok/s moved with the 27B's, so
the decode drift is environmental, not a #11 effect.
## Reproducing
```sh
# the whole fleet (all loaded models), defaults shown
./script/bench.py --base-url http://hanzalova.internal:31313/v1 \
--runs 3 --prompt-tokens 128,4096 --max-tokens 256 \
--json bench-results.jsonl
# a competitor engine for comparison columns
./script/bench.py --base-url http://localhost:8080/v1 \
--label llama.cpp --model <model-id>
```
Append-only JSON rows (`--json`) keep a longitudinal record across
commits; the `engine` label column makes cross-engine tables a
concatenation, not a merge.
## Automated harness (`helexa-bench`)
`script/bench.py` above is the manual, ad-hoc probe (any `/v1`
endpoint, run by hand). The `helexa-bench` crate is its continuous,
version-aware successor: a daemon (one systemd unit, typically on the
metrics host) that hits each neuron **directly** on `:13131`, exercises
every **warm** model, and records each run into a SQLite
system-of-record stamped with the neuron's full build identity — git
SHA, enabled cargo features, rustc/candle versions — read from the new
neuron `GET /version` endpoint.
It is keyed by build SHA: a given neuron build is benchmarked only until
it has `samples_per_version` results per (model, scenario), then skipped
until a new SHA ships. So the table below can be regenerated
automatically per neuron update instead of edited by hand:
```sh
helexa-bench once --config helexa-bench.toml # single sweep
helexa-bench report --config helexa-bench.toml # markdown table by SHA
```
The scenario method (synthetic 128/4096-token prompts, `/no_think`,
streamed TTFT + decode-window tok/s) is ported verbatim from bench.py,
so its columns stay comparable. The OpenAI-target seam for cross-engine
comparison rows is scaffolded but not yet wired (see gaps).
## Known gaps
- **No competitor baselines yet** — requires llama.cpp / Ollama
serving the same checkpoints on the same hosts; the harness is
ready for them.
- **Cold-load time** is not yet measured here; it is visible per
deploy in the `loaded default model … elapsed_ms=…` journal line
and the deploy workflow's validation step, and is tracked as #1.
- **Streaming usage**: neuron does not emit a final usage frame on
SSE streams yet, so token counts rely on the chunk-per-token
invariant.

129
doc/learnings/p1.md Normal file
View File

@@ -0,0 +1,129 @@
# P1 learnings — briefing for the P2 session
Written 2026-06-12 at the close of the P1 sprint (#20, #19, #21, #24
closed; #22 landed minus competitor columns — see the pinned tracking
issue #27). Everything below was learned the hard way during P1 and is
directly useful for P2 (#11 prefix KV caching, #23 chunked prefill,
#1 cold-load, #15 numerical validation).
> Relocated from `doc/plan/p1-learnings.md` (gitignored) into
> `doc/learnings/` so it lives in source control alongside
> [`p2.md`](p2.md). Content is the original P1 briefing, verbatim.
## Working agreements that are now live
- **merge-on-green is standing policy** for roadmap PRs: open the PR,
background-poll `commits/<sha>/status`, merge with branch-delete the
moment it's green. On red: investigate, never merge.
- **Branch → PR per roadmap issue**, `Closes #N` in the commit message
(Gitea auto-closes on merge). Tick the checkbox on #27 afterwards.
- **Every merge to main self-validates on the fleet**: deploy.yml waits
for `/health` activation `ready` (per-host timeouts: beast 900 s,
others 300 s), fails on any `activation.failed` entry with the
per-model error, then asks the loaded model to say a specific word
("LLM probe"). A P2 engine change that breaks load or inference
turns the deploy run red — trust it, watch it after engine merges.
- **Deploy gating is manifest-equality**: a host restarts only when
`rpm -q` differs from the newest packages.json entry. Never
reintroduce unprivileged `dnf check-update` — it both hung (quadbrat)
and silently lied (benjy/beast) in the same run.
## CI behaviour you must plan around
- **`crates/cortex-core/` touches trigger a FULL fleet rebuild and
neuron restarts** (it's in both change-detection regexes in
build-prerelease.yml). Each neuron restart costs a cold-load
(~3m40s on beast, measured by the deploy validation). For P2 engine
work this is unavoidable (neuron changes anyway); for incidental
core type tweaks, batch them into the engine PR rather than landing
them separately.
- Gateway-only and docs-only pushes skip neuron builds/restarts
entirely. The skip logic diffs against the *last published RPM's*
sha, so nothing is ever silently missed.
- **`GITHUB_ENV` does not propagate between steps** on these runners
(observed: probe step set RUSTC_WRAPPER, build ran unwrapped).
Job-level `env:` and same-step `export` work; cross-JOB outputs
(`needs.X.outputs.Y`) work. Set env where you use it.
- **sccache**: present in the CUDA runner image, wired with escalation
(retry → server restart → final attempt uncached) in lint/test/build
jobs. A sick cache costs minutes, never the run. Same-ref pushes
coalesce (`cancel-in-progress: true` on build-prerelease) — pushing
twice in quick succession wastes the first build, so batch commits.
- **`deploy-dev` workflow** (Actions → deploy-dev → pick host): builds
ONE flavour and scp's the raw binary onto the host — the fast
iteration path for engine work (~build + restart, no RPM/publish).
The sudoers rule requires the exact command form in the workflow;
change both together. The next regular deploy reconciles the host.
- Reading Gitea job logs via API: runner "expression evaluated"
noise and the rendered-script dump both match naive greps — filter
with `grep -v 'expression\|rewritten'` and remember the masker
replaces any string equal to a secret value with `***` (including,
comically, the word "sccache").
## Engine facts P2 will need (learned from live traffic)
- **neuron streams exactly one SSE chunk per generated visible
token.** The gateway token metrics (#21) and bench harness (#22)
rely on this invariant when no usage frame is present. If P2 work
changes chunking, fix those consumers or land #31 first.
- **No usage frame on streams** (`stream_options.include_usage`
ignored — #31, p2-labelled). Non-streaming responses DO carry usage.
Gateway streaming token counters stay silent until #31 lands.
- **Reasoning deltas are off-wire by default** (dropped unless the
request opts into thinking). Consequences: thinking models burn
`max_tokens` invisibly, and short `max_tokens` can yield ZERO
visible chunks (`finish_reason: length`, empty stream). Benchmarks
and probes append Qwen's `/no_think` soft switch — it renders an
empty think block and works on all fleet models. The deploy probe
uses `max_tokens: 512` headroom instead.
- **27B baseline (doc/benchmarks.md, 2026-06-12)**: decode ~35 tok/s
FLAT from 128→4k context (Gated DeltaNet's O(1) recurrent state);
TTFT 1.66 s @128 / 7.07 s @4k. The 4k TTFT is the #23 before-number.
Re-run `script/bench.py` (same flags as the doc) after every engine
change — it takes ~5 min for the fleet and the table is
append-friendly via `--json`.
## Design notes for #11 (prefix KV caching) specifically
- The per-request cache clear lives in `harness/candle.rs` (~line
1393, `clear_kv_cache()` before every inference) — that call site is
the entry point: "stop deleting it" is the issue's framing.
- **Hybrid-architecture caveat**: for the qwen3_5/3.6 family, a
"prefix cache" is NOT just attention KV. Three out of four decoder
layers are GatedDeltaNet carrying `conv_state` + `recurrent_state`
(see `harness/arch/qwen3_5/linear_attn.rs` doc-comment), and the
vision/M-RoPE position counters also persist (`tp_qwen3_5.rs`
clear_kv_cache resets rope counters). A reusable prefix snapshot at
token N = attention KV **plus** GDN states **plus** position
counters, all consistent at exactly N. Recurrent state makes partial
prefix reuse impossible mid-stride — match must be exact-prefix at
the snapshot boundary.
- **All cache state must live inside the device worker** (the slab,
per CLAUDE.md "Per-device worker thread"): tensors never escape the
worker thread; new state types mean new `Job` variants in
`device_worker/jobs.rs` + handlers in `dispatch.rs`. The recovery
path (#17/#20) unloads and reloads models — cached prefixes must be
invalidated on unload (the snapshot pattern used by the `recovering`
map in candle.rs is a reference for state that must outlive registry
slots — but for caches, dropping on unload is correct and simpler).
- The eviction story (#11 body: LRU bounded by a per-device VRAM
budget) can reuse the worker's existing VRAM-query plumbing
(`device_vram_mb` routes through the worker since Phase 1).
- Measurement: `cortex_time_to_first_token_seconds{model,node}` is
live in Prometheus, and `script/bench.py` gives the before/after
table. For agent-shaped workloads (the win case), benchmark a
multi-turn conversation: send turn N, then turn N+1 = same prefix +
new question, and compare TTFT-with-cache vs today's full re-prefill.
## Misc operational facts
- Fleet: beast 2×5090 (27B Q6K TP=2), benjy 4090 (8B), quadbrat 3060
(1.7B); gateway on hanzalova:31313; all reachable by ssh as the
operator (passwordless sudo) — used for sudoers sync, live probes,
and killing wedged processes during P1.
- The tea CLI token (`~/.config/tea/config.yml`) authenticates against
git.lair.cafe for everything the gitea-mcp tools don't cover
(milestones, labels, pin, merge).
- Anthropic streaming is now production-real: Claude Code can be
pointed at the gateway. Real-traffic feedback from that is useful
input for P2 priorities (especially tool-use-heavy streams).

251
doc/learnings/p2.md Normal file
View File

@@ -0,0 +1,251 @@
# P2 learnings — briefing for the P3 session
Written 2026-06-13 at the close of the P2 sprint. The whole P2 queue
landed in one session: #11 (prefix KV caching), #23 (chunked
delta-rule prefill), #1 (TP cold-load), #15 (numerical validation) —
plus #42 filed (CI flake). Numbers and PRs are on the issues and the
tracker (#27); this doc is the hard-won *how*, for #18, #25, #7, #26,
#4 and anything after.
Read [`p1.md`](p1.md) first — its working agreements (merge-on-green,
branch→PR per issue, fleet self-validation, the `deploy-dev` fast
path) all held through P2 and are not repeated here except where P2
changed or sharpened them.
## The single most important lesson: CI green ≠ correct
P2 shipped four correctness bugs that **passed every CI gate** and
only surfaced on the live fleet. Each now has a regression test that
forbids the wrong thing. The pattern is the lesson: the engine math
is exercised by CPU unit tests on tiny random tensors, and random
data is *benign* — it doesn't reproduce the conditions real model
weights and real prompts create.
1. **Prefix snapshot taken post-generation never re-matched** (#11).
The cached token sequence included the model's `<think>` reasoning
tokens, which the client strips when it echoes the assistant turn
back — so the next prompt was never a token-prefix of the cached
sequence. Caught only by watching `reused=0` in the journal on a
real multi-turn conversation against the 27B. Fix: snapshot at the
*prefill boundary* (prompt-only), not after generation (PR #36).
2. **Full-prompt snapshots broke on BPE retokenization** (#11). A
prompt ends `…<|im_start|>assistant\n`; when the next turn
re-tokenizes that text followed by the reply, the trailing `\n`
merges with the reply's first characters into *different* token
ids — so the last 12 tokens of the cached sequence diverge and the
exact-prefix match (forced by GDN recurrent state — no partial
rewind) never fires. Fix: snapshot one past the **last special
token** (`<|im_start|>`); special tokens are hard tokenizer
segmentation points, so that prefix is provably stable across
renders (PR #37).
3. **The 0.8B "passed" both of the above by luck** — its replies
started with the atomic `<think>` special token, which both keeps
the reasoning marker in the echoed text *and* blocks the BPE merge.
**Validate any model-text-dependent behaviour on more than one
model.** A single model can mask a class of bug entirely.
4. **The chunked delta-rule's tempting optimization was numerically
wrong in f32** (#23). `(I T)⁻¹` for the strictly-lower-triangular
`T` *can* be written as the nilpotent squaring product
`Π(I + T^(2^j))` (6 matmuls vs 63 sequential row updates), and it
**passed parity on random data**. On real prompts with repetitive
text → correlated keys, raw powers of `T` grow combinatorially
(path counts ≈ C(62,31) ≈ 4.6e17) before nilpotency collapses
them, destroying f32 precision → NaN logits → `"!!!"` replies.
The HF reference uses forward substitution for a reason; port it
faithfully. The regression test
(`chunked_ut_transform_survives_correlated_keys`) builds
near-identical keys with β≈1 and diverges to ~8e30 under the
squaring form.
**Takeaway for P3:** when you write an engine-math unit test, ask
"what would real weights/prompts do that random tensors don't?" —
correlated keys, special-token boundaries, reasoning markers,
repetitive text. Build at least one adversarial fixture per new path.
The live two-turn / fixed-prompt probe on the fleet caught all four;
budget time for it after every engine merge, not just at the end.
## Validating on the fleet (what worked)
- **Live multi-turn / fixed-prompt probe over ssh** is the fastest
truth: `curl <host>:13131/v1/chat/completions` and read
`journalctl -u neuron` (fleet runs `RUST_LOG=debug`) for the
request-path lines — `prefix cache: hit reused_tokens=…`,
`prefill complete … reused=… elapsed_ms=…`, the new per-phase
load timings. Greedy (`temperature: 0`) so outputs are comparable.
- **A/B a change behind an env kill-switch** without redeploying:
drop `Environment=NEURON_GDN_CHUNKED=0` (or whatever the switch is)
into `/etc/systemd/system/neuron.service.d/<name>.conf`,
`daemon-reload`, `restart`, measure; remove the drop-in, restart,
measure again. This gave the clean 8156 ms → 3694 ms chunked-vs-
recurrent prefill comparison on the *same binary, same prompt*. Add
such a switch to any perf-sensitive path you land — it's the only
way to get an honest before/after without a second deploy.
- **`deploy-dev` is the iteration loop for engine work**: it builds
one flavour and scp's the raw binary, ~build + restart, no
RPM/publish (~10 min on beast vs ~25 for a full pipeline). Dispatch
it from the feature branch via the API
(`POST actions/workflows/deploy-dev.yml/dispatches`,
`{"ref":"<branch>","inputs":{"target":"beast"}}`). The next regular
deploy reconciles the host back to the packaged binary. Used it to
catch the `"!!!"` NaN bug before merging, then again to verify the
fix.
- **`nvcc` is on beast** (`/usr/local/cuda-13.0/bin/nvcc`, not on
PATH) — so beast can build the `cuda` feature locally for an
apples-to-apples bf16 numerical comparison, and has the full Rust
toolchain (`~/.cargo/bin`). Useful for the #15 harness and any
cuda-gated test you can't run on the dev box.
## CI behaviour P2 added to the P1 list
- **The "CUDA type-check" job is your only pre-merge check of
`#[cfg(feature = "cuda")]` code.** Local clippy/test on the dev box
compile no-cuda only (no nvcc), so cuda-gated mistakes compile
clean locally and fail in CI. P2 hit two:
- a `std::sync::MutexGuard` (the prefix-cache registry) held across
an `.await` inside a `let`-chain made the TP request futures
non-`Send``tokio::spawn` rejected them. **Bind the guard's
result before any await; never hold a registry lock across a
suspension point.**
- a `match` arm returning `candle_core::Error` vs `anyhow::Error`
— fine until the cuda arm widened the type set.
Push the branch early and watch that one job after touching
cuda-gated code; don't wait for the full green.
- **Change-detection skip is per-package, and a flavour build failure
splits the fleet.** build-prerelease diffs each package against the
last *published* RPM's sha and skips unchanged ones — good. But when
one flavour's *build* fails (see #42), only that package is skipped
from publish; the other two publish and deploy, leaving the fleet on
mixed versions. **Watch all three `Build neuron-<flavour>` +
`Package …` jobs, not just the combined commit status** (which can
read `success`/`pending` while a flavour quietly failed).
- **Retrigger a skipped/failed flavour with an empty commit.** Gitea
on this version has no run-rerun API endpoint
(`actions_run_write rerun` returns 404). `git commit --allow-empty`
+ push re-runs the pipeline; change-detection sees the still-stale
flavour (its last *published* sha is older) and rebuilds it. P2
needed this twice for #42.
- **The RPM `git<sha>` stamp is the per-package input-change sha, NOT
main HEAD.** Don't confirm a rollout by grepping `rpm -qa` for the
merge commit's short sha — it won't match. Confirm via an **uptime
reset** on `/health` plus a behaviour only the new build has (a new
log line, a new metric). The fleet-validation memory has this.
- **`cancel-in-progress` can split a host mid-deploy.** A same-ref
push that lands between a host's `dnf upgrade` and its `restart`
cancels the restart, leaving the host with the *new* RPM installed
but the *old* binary running. The next successful deploy heals it
(the rpm-vs-packages.json compare sees the next version). Filed as
part of #42; avoid pushing again while a deploy is mid-flight.
## Engine facts P2 established (for #18, #25, and beyond)
- **The device-worker discipline scales cleanly to new state.** Adding
prefix snapshots was: new `Job` variants
(`SnapshotKv`/`RestoreKv`/`DropKvSnapshot`) + handlers in
`dispatch.rs` + a `HashMap` beside the model slab, dropped with the
model in `DropArch`. The async side holds only an opaque id + the
token sequence + a byte size — tensors never escape the worker.
Reuse this shape for any new per-model device state (#25's drafter
KV, for instance).
- **TP mirrors single-GPU through one pool-minted id.** For TP the
leader's snapshot lives in its device worker (`Job::Tp*`) and each
subprocess rank stores its shard's snapshot in-process via new
`WorkerRequest` variants in `tp/rpc.rs`, all keyed by the *same*
id the pool mints and broadcasts. Step fan-out is synchronous, so
all ranks sit at the same token boundary — that's what makes a
cross-rank snapshot consistent. Partial-failure rule: any rank fails
restore → clear all ranks + full prefill; any rank fails snapshot →
drop the id everywhere. The TP shard state has the *same shape* as
single-GPU, so the `arch/qwen3_5/snapshot.rs` types are shared
verbatim.
- **GDN state copy semantics are not uniform.** The CUDA delta-rule
kernels mutate the recurrent-state buffer **in place**
(`&mut state_bh`; `flatten`/`contiguous` on a contiguous tensor is a
view), so GDN `conv_state`/`recurrent_state` snapshots must
**deep-copy** (`Tensor::copy`) in both directions. Attention
`ConcatKvCache` k/v can share storage (its `append` cats into fresh
allocations and never mutates stored tensors). Get this wrong and
the snapshot silently tracks the live cache.
- **The reusable-prefix snapshot boundary is "one past the last
`<|im_start|>`"** — not the full prompt (BPE retokenization), not
post-generation (reasoning/tool-call tokens stripped on echo). The
request path prefills to that boundary, snapshots, then finishes the
~2-token `assistant\n` tail. This generalises: any future
cross-request state reuse must key on a tokenizer-stable boundary,
and special tokens are the only provably-stable ones.
- **Chunked prefill composes with the prefix cache.** A restored
conversation's divergent suffix still prefills chunked (the chunked
path takes a non-zero `initial_state`, validated by a parity test).
Decode and short (<64-token) prompts keep the recurrent path —
decode was deliberately untouched. #18 (single-GPU vision prefill
chunking) and #25 (speculative decoding) both build on these paths.
- **Cold-load is now disk-read-bound.** Parallel ISQ (#1) cut the 27B
TP cold-load 221 s → 86 s by fanning candle's per-block k-quant math
(`k_quants::GgmlType::from_float`, public API — no candle fork)
across the rayon pool, **byte-identical** to `QTensor::quantize`.
The residual 79 s layer-loop is now ~1.2 s/layer dominated by
reading the 54 GB bf16 safetensors (~700 MB/s with both ranks
reading), not quantization. The next lever (issue #1 item 3) is a
post-ISQ disk cache; raise a fresh issue only if 86 s still hurts.
Phase timing is now instrumented (`layer loop complete elapsed_ms`,
`lm_head loaded elapsed_ms`) — read it before optimizing further.
## Numerical validation (#15) — reusable rig for every arch
- **Capture reference fixtures in f32, not bf16.** f32-vs-f32 pins the
*math*; the implementation matched HF exactly (text logits
`max_abs 0.000`, vision end-to-end `cosine 1.000000`). Cross-dtype
comparison is dominated by bf16 rounding chaos through a deep tower
(global cosine ~0.997, *worst patch* ~0.92, worst-patch index
unstable across runs) — that's production-dtype noise, not bug, and
it drowns the signal. The script defaults to f32 for this reason.
- **The 0.8B validates the same arch modules as the 27B** — they
differ only in hyperparameters — so the automated test runs against
the 0.8B (an f32 27B forward needs ~108 GB; beast has ~91 GB free).
A bf16 27B text fixture is checked in for the manual procedure.
- **Mutation sensitivity is the test that the test works.** Re-running
with `NEURON_VISION_LEGACY_POS=1` (the deliberately-wrong sequential
pos-embed lookup) collapses tower cosine 0.999998 → 0.753 and fails
loudly. Every numerical harness needs a known-bad toggle to prove it
isn't asserting on noise.
- **The harness found a real production bug.** The pos-embed bilinear
blend was rounding interpolation weights to bf16 *before* blending;
the reference keeps them f32 and casts once at the end. Fixed in the
#15 PR. Numerical validation pays for itself.
- **Reproducing HF on beast needs two shims** (documented in
`crates/neuron/tests/fixtures/numerical/README.md` and the dump
script): (a) transformers 5.9 ↔ kernels 0.15 import breakage —
monkeypatch `LayerRepository.__init__` / `FuncRepository.__init__`
to inject a `revision`, with `USE_HUB_KERNELS=NO`; (b) neuron's
local HF snapshot lacks `preprocessor_config.json` (hf-hub only
fetched what neuron needed) — load the *processor* from the repo id
with `HF_HUB_CACHE` pointed at a scratch dir. The test self-skips
without `NEURON_REF_MODEL_PATH`, so CI compiles it without weights.
- **Extend this for any new arch.** `script/dump_reference.py` +
`tests/numerical_reference.rs` are the template; capture
text + vision fixtures, calibrate thresholds against the observed
f32 agreement, add a mutation toggle.
## State of the board for P3
- **P1 + P2 are fully closed** with published before/after numbers on
every issue and on #27. The milestone-7→8 perf story is in hand:
prefix cache (warm 27B turn ~10 s → ~2.3 s), chunked prefill (cold
~5k prefill 8.2 s → 3.7 s), cold-load (221 s → 86 s), and a
fidelity harness proving f32-exact parity with HF.
- **#26 (tagged release + public writeup) is now unblocked** — it was
gated on the #22/#23 numbers, which all exist. It's the natural next
pick: the engine work it would showcase is done.
- **#25 (speculative decoding) is sequenced after #23** (done) and
reuses the device-worker + TP-rank state patterns above for the
drafter. **#18 (single-GPU vision prefill chunking)** is the
single-card-tier parity counterpart to #23 and the chunked-prefill
machinery already exists.
- **#42 (ampere `ptxas` SIGSEGV in candle-flash-attn sm_86 kernels)**
is open runner-infra: intermittent, sticky within a run, clears
across runs; cost two retriggers in the P2 session. Mitigation ideas
(serialize the flash-attn kernel compile, cap concurrent flavour
builds, check runner memory headroom) are on the issue. Until it's
fixed, watch the three flavour jobs and retrigger with an empty
commit when one segfaults.
- The fleet, gateway, tea-CLI auth, and ssh/sudo access are unchanged
from P1 — see [`p1.md`](p1.md)'s "Misc operational facts".

View File

@@ -5,7 +5,7 @@ Sourced from beast's local cache on 2026-06-01:
Single source of truth for Stages AD of the vision plan in
`~/.claude/plans/foamy-twirling-catmull.md`. Umbrella issue:
[#3](https://git.lair.cafe/helexa/cortex/issues/3).
[#3](https://git.lair.cafe/helexa/helexa/issues/3).
---
@@ -92,7 +92,7 @@ Reading:
- `image_mean = image_std = 0.5` → normalisation is simply `(x/255 - 0.5) / 0.5 = 2*x/255 - 1`, mapping `[0,255]``[-1, 1]`. No imagenet-style mean/std.
- `size.{shortest_edge, longest_edge}` are **pixel counts**, not edge lengths. The `Qwen2VLImageProcessorFast` recipe picks a resolution within `[65,536 = 256², 16,777,216 = 4096²]` total pixels, snapping `h` and `w` to multiples of `patch_size × spatial_merge_size = 32` pixels.
- Stage A ships **fixed resolution**: pick a target pixel count (e.g. 448×448 = 200,704 px → 28×28 patches → 14×14 LM tokens after merger). Variable resolution deferred to issue [#14](https://git.lair.cafe/helexa/cortex/issues/14).
- Stage A ships **fixed resolution**: pick a target pixel count (e.g. 448×448 = 200,704 px → 28×28 patches → 14×14 LM tokens after merger). Variable resolution deferred to issue [#14](https://git.lair.cafe/helexa/helexa/issues/14).
## Chat template (`chat_template.jinja`)
@@ -140,7 +140,7 @@ rope_parameters: {
}
```
MRoPE encodes spatial position alongside text position so the LM attention layers can reason about image-token spatial structure. The LM's existing forward path *may or may not* already implement this — the qwen3_5 module's doc-comment notes "numerical correctness vs the reference Python is not yet validated." Verifying MRoPE behaviour in the language model is out of Stage A scope (vision tower only) but will be required in Stage B (LM splice) and is tracked under the numerical-validation issue [#15](https://git.lair.cafe/helexa/cortex/issues/15).
MRoPE encodes spatial position alongside text position so the LM attention layers can reason about image-token spatial structure. The LM's existing forward path *may or may not* already implement this — the qwen3_5 module's doc-comment notes "numerical correctness vs the reference Python is not yet validated." Verifying MRoPE behaviour in the language model is out of Stage A scope (vision tower only) but will be required in Stage B (LM splice) and is tracked under the numerical-validation issue [#15](https://git.lair.cafe/helexa/helexa/issues/15).
`max_position_embeddings = 262144` (256 K context), so context-length limits are not a constraint for vision.
@@ -152,7 +152,7 @@ The vision tower has its own self-contained weight tree and is small (~333 tenso
- Run unit tests with random tensor weights matching the exact shapes → assert forward produces correct output shape with finite values.
- Optionally: a CUDA-integration test that loads just the 2 vision shards from beast's cache (or on a smaller GPU like quadbrat's Ampere) and runs encode on a real image. Doesn't require loading the 27B LM at all.
This sidesteps the "develop against a smaller VL model" question for Stage A. Stage B (LM splice → end-to-end chat with vision) is where iteration speed becomes pressing; revisit there. The default scope pick 2a (smaller iteration model) is therefore deferred to Stage B planning — issue [#13](https://git.lair.cafe/helexa/cortex/issues/13) covers deployment validation regardless.
This sidesteps the "develop against a smaller VL model" question for Stage A. Stage B (LM splice → end-to-end chat with vision) is where iteration speed becomes pressing; revisit there. The default scope pick 2a (smaller iteration model) is therefore deferred to Stage B planning — issue [#13](https://git.lair.cafe/helexa/helexa/issues/13) covers deployment validation regardless.
## Concrete Stage A1+ inputs
@@ -167,10 +167,10 @@ This sidesteps the "develop against a smaller VL model" question for Stage A. St
## What this doc does NOT settle (deferred to issues)
- Numerical correctness of `VisionTower` output vs Python transformers
→ issue [#15](https://git.lair.cafe/helexa/cortex/issues/15).
→ issue [#15](https://git.lair.cafe/helexa/helexa/issues/15).
- Variable image resolution
→ issue [#14](https://git.lair.cafe/helexa/cortex/issues/14).
→ issue [#14](https://git.lair.cafe/helexa/helexa/issues/14).
- TP-vision (multi-rank vision tower)
→ issue [#12](https://git.lair.cafe/helexa/cortex/issues/12).
→ issue [#12](https://git.lair.cafe/helexa/helexa/issues/12).
- 27B production deployment
→ issue [#13](https://git.lair.cafe/helexa/cortex/issues/13).
→ issue [#13](https://git.lair.cafe/helexa/helexa/issues/13).

50
helexa-bench.example.toml Normal file
View File

@@ -0,0 +1,50 @@
# helexa-bench — continuous, version-aware fleet benchmark harness.
#
# Hits each neuron directly, exercises warm models, and records every run
# with full build/version provenance into SQLite. Once a neuron build has
# `samples_per_version` results for a (model, scenario), later sweeps skip
# it until a new build SHA ships — so a steady fleet costs only cheap
# version polls.
#
# Env overrides: BENCH_-prefixed, `__` for nesting
# (e.g. BENCH_BENCH__SAMPLES_PER_VERSION=10).
[bench]
# Pause between full sweeps of all targets (seconds).
sweep_interval_secs = 1800
# Target measured samples per (target, build SHA, model, scenario).
samples_per_version = 5
# Pause between successive measured iterations against one model.
iteration_pause_secs = 2
# Per-request timeout (seconds); generous for cold lazy-loads.
request_timeout_secs = 600
# SQLite system-of-record.
db_path = "/var/lib/helexa-bench/bench.sqlite"
[scenarios]
# One chat-latency scenario is generated per size (chat:128, chat:4096).
prompt_sizes = [128, 4096]
max_tokens = 256
# One [[targets]] block per neuron on the fleet. `kind = "neuron"` (the
# default) gets build metadata via GET /version and warm-model discovery
# via GET /models.
[[targets]]
name = "beast"
endpoint = "http://beast.hanzalova.internal:13131"
[[targets]]
name = "benjy"
endpoint = "http://benjy.hanzalova.internal:13131"
[[targets]]
name = "quadbrat"
endpoint = "http://quadbrat.hanzalova.internal:13131"
# Future: compare against a non-neuron OpenAI-compatible engine. `kind =
# "openai"` skips neuron-only metadata; point `endpoint` at the /v1 base.
# [[targets]]
# name = "llamacpp-ref"
# kind = "openai"
# endpoint = "http://benjy.hanzalova.internal:8080/v1"
# label = "llama.cpp"

View File

@@ -7,7 +7,7 @@ Summary: Per-node GPU discovery and harness management daemon for cortex
# unit, and system user are still called "neuron" for brevity.
License: GPL-3.0-or-later
URL: https://git.lair.cafe/helexa/cortex
URL: https://git.lair.cafe/helexa/helexa
Source0: %{name}-%{version}.tar.gz
Source1: %{name}-%{version}-vendor.tar.gz
@@ -54,6 +54,11 @@ directory = "vendor"
EOF
%build
# Source tarballs carry no .git, so build.rs can't recover the commit on
# its own — it would report "unknown" from GET /version. Pass the commit
# in with `rpmbuild --define "helexa_commit <sha>"`; absent that, it
# degrades to "unknown" rather than failing the build.
export HELEXA_BUILD_SHA="%{?helexa_commit}"
cargo build --release -p neuron
%install

View File

@@ -64,6 +64,19 @@ name = "candle"
# auth_env = "HELEXA_TOKEN"
# cache_dir = "/archive3/llm-cache/helexa"
# -- Prefix KV cache ----------------------------------------------------------
# Reuse cache state across requests when a new prompt starts with the
# exact token sequence of a previous one (chat/agent workloads), so
# prefill only runs on the new suffix. Applies per loaded model, on
# architectures that expose their cache state (qwen3_5). Snapshots
# live in device memory: budget_mb is per loaded model and comes out
# of the same VRAM that serves inference.
#
# [harness.candle.prefix_cache]
# enabled = true
# budget_mb = 1024
# max_entries = 8
# -- Default models ----------------------------------------------------------
# Models listed here are loaded automatically when the neuron service
# activates. Loading is sequential — a slow or failing entry doesn't

View File

@@ -29,7 +29,7 @@ Release: %{cortex_release}%{?dist}
Summary: Inference gateway for multi-node GPU clusters (prebuilt)
License: GPL-3.0-or-later
URL: https://git.lair.cafe/helexa/cortex
URL: https://git.lair.cafe/helexa/helexa
Source0: cortex
Source1: cortex.service

View File

@@ -0,0 +1,98 @@
# Prebuilt-binary spec for helexa-bench.
#
# Wraps a pre-built `helexa-bench` binary produced by an upstream CI job
# and packages it for rpm.lair.cafe. The %build phase is a no-op.
# helexa-bench is a pure-Rust, non-CUDA, outbound-only daemon (no
# listener), so there is no firewalld service to install.
#
# Required defines at rpmbuild time:
# bench_version e.g. "0.1.16"
# bench_prerelease e.g. "0.1.20260518140530.gitabcdef0"
# ^^^^^^^^^^^^^^^^^^ ^^^^^^^^
# commit time (sec) commit sha
# (used as Release; the timestamp prefix
# keeps same-day builds strictly ordered.)
%global _build_id_links none
%global debug_package %{nil}
%global __strip /usr/bin/true
%{!?bench_version: %global bench_version 0.0.0}
%if 0%{?bench_prerelease:1}
%global bench_release %{bench_prerelease}
%else
%global bench_release 1
%endif
Name: helexa-bench
Version: %{bench_version}
Release: %{bench_release}%{?dist}
Summary: Continuous version-aware benchmark harness for the neuron fleet (prebuilt)
License: GPL-3.0-or-later
URL: https://git.lair.cafe/helexa/helexa
Source0: helexa-bench
Source1: helexa-bench.service
Source2: helexa-bench-sysusers.conf
Source3: helexa-bench.example.toml
Source4: LICENSE
ExclusiveArch: x86_64
Requires(pre): shadow-utils
Requires: systemd
Provides: user(helexa-bench)
%description
helexa-bench hits each neuron on the fleet directly, exercises an
extensible benchmark suite against every warm model, and records each
run with full build/version provenance into a SQLite store. It runs
continuously under systemd and is version-aware: a given neuron build is
benchmarked only until it has the configured number of samples, then
skipped until a new build ships. Replaces manual bench.py runs.
%prep
cp %{SOURCE0} ./helexa-bench
cp %{SOURCE1} .
cp %{SOURCE2} .
cp %{SOURCE3} .
cp %{SOURCE4} .
%build
# Already built in the upstream CI build job.
%install
install -Dm755 helexa-bench %{buildroot}%{_bindir}/helexa-bench
install -Dm644 helexa-bench.service %{buildroot}%{_unitdir}/helexa-bench.service
install -Dm644 helexa-bench-sysusers.conf %{buildroot}%{_sysusersdir}/helexa-bench.conf
install -dm755 %{buildroot}%{_sysconfdir}/helexa-bench
install -Dm644 helexa-bench.example.toml %{buildroot}%{_sysconfdir}/helexa-bench/helexa-bench.toml
%pre
getent group helexa-bench >/dev/null || groupadd -r helexa-bench
getent passwd helexa-bench >/dev/null || \
useradd -r -g helexa-bench -d /var/lib/helexa-bench -s /sbin/nologin \
-c "helexa-bench harness" helexa-bench
%post
%systemd_post helexa-bench.service
%preun
%systemd_preun helexa-bench.service
%postun
%systemd_postun_with_restart helexa-bench.service
%files
%license LICENSE
%{_bindir}/helexa-bench
%{_unitdir}/helexa-bench.service
%{_sysusersdir}/helexa-bench.conf
%dir %{_sysconfdir}/helexa-bench
%config(noreplace) %{_sysconfdir}/helexa-bench/helexa-bench.toml
%changelog
* Sat Jun 13 2026 Gitea Actions <actions@git.lair.cafe> - %{bench_version}-%{bench_release}
- Prerelease build from upstream CI binary.

View File

@@ -36,7 +36,7 @@ Release: %{neuron_release}%{?dist}
Summary: Per-node GPU inference daemon (candle, %{neuron_flavour} flavour)
License: GPL-3.0-or-later
URL: https://git.lair.cafe/helexa/cortex
URL: https://git.lair.cafe/helexa/helexa
Source0: neuron-%{neuron_flavour}
Source1: neuron.service

205
script/bench.py Executable file
View File

@@ -0,0 +1,205 @@
#!/usr/bin/env python3
"""Reproducible batch-1 benchmark harness for helexa (#22).
Measures what one operator at a keyboard feels, per model:
- TTFT time from request send to the first SSE content chunk
- decode completion tokens per second over the first->last chunk
window (token count from the final `usage` object when
the server sends one, else the content-chunk count)
- total wall-clock for the whole request
Works against ANY OpenAI-compatible /v1 endpoint (helexa's cortex,
llama.cpp's llama-server, Ollama's compat endpoint, vLLM, ...), so the
same invocation produces comparable columns across engines:
./script/bench.py --base-url http://hanzalova.internal:31313/v1
./script/bench.py --base-url http://localhost:8080/v1 --model qwen3:8b
stdlib-only on purpose: no venv, no pip, runs from any Fedora host.
Results print as a markdown table; --json appends machine-readable
rows for longitudinal tracking (doc/benchmarks.md records the method).
"""
import argparse
import json
import statistics
import sys
import time
import urllib.error
import urllib.request
# A paragraph of filler re-used to synthesise prompts of a target
# approximate token count (~4 chars/token heuristic — close enough for
# bucketing; real token counts are read back from the usage object).
FILLER = (
"The quick brown fox jumps over the lazy dog while the band plays "
"a slow waltz in the background and somebody counts the beats. "
)
# /no_think: Qwen3-family soft switch rendered by the chat template;
# keeps thinking models from burning the token budget invisibly
# (reasoning deltas are not on the wire by default). Harmless for
# non-thinking models.
QUESTION = (
"\n\nRetell the scene above as a vivid story of about 300 words. /no_think"
)
def build_prompt(approx_tokens: int) -> str:
target_chars = max(approx_tokens, 16) * 4
body = (FILLER * (target_chars // len(FILLER) + 1))[:target_chars]
return body + QUESTION
def one_run(base_url: str, model: str, prompt: str, max_tokens: int, timeout: float):
"""Single streamed request. Returns dict with ttft, decode_tps,
total_s, completion_tokens, prompt_tokens (None where unknown)."""
payload = {
"model": model,
"messages": [{"role": "user", "content": prompt}],
"max_tokens": max_tokens,
"temperature": 0,
"stream": True,
"stream_options": {"include_usage": True},
}
req = urllib.request.Request(
f"{base_url}/chat/completions",
data=json.dumps(payload).encode(),
headers={"content-type": "application/json"},
)
start = time.monotonic()
first = last = None
chunk_count = 0
prompt_tokens = completion_tokens = None
tail = ""
with urllib.request.urlopen(req, timeout=timeout) as resp:
buf = b""
while True:
block = resp.read(8192)
if not block:
break
now = time.monotonic()
buf += block
while b"\n\n" in buf:
event, buf = buf.split(b"\n\n", 1)
line = event.decode("utf-8", "replace").strip()
if not line.startswith("data:"):
continue
data = line[len("data:") :].strip()
if data == "[DONE]":
continue
tail = data # last data frame wins (usage rides there)
try:
obj = json.loads(data)
except json.JSONDecodeError:
continue
choices = obj.get("choices") or []
delta = (choices[0].get("delta") or {}) if choices else {}
if delta.get("content"):
if first is None:
first = now
last = now
chunk_count += 1
usage = obj.get("usage")
if usage:
prompt_tokens = usage.get("prompt_tokens")
completion_tokens = usage.get("completion_tokens")
end = time.monotonic()
if first is None:
raise RuntimeError(f"no content chunks received (last frame: {tail[:200]})")
# neuron emits exactly one SSE chunk per generated visible token,
# so chunk count is an engine-truth count when no usage frame is
# sent (streaming include_usage is not implemented yet).
tokens = completion_tokens if completion_tokens else chunk_count
# decode rate is only meaningful over a real inter-chunk window;
# short replies can arrive coalesced into one TCP read (window=0).
window = (last - first) if (last and last > first) else 0.0
return {
"ttft_s": first - start,
"decode_tps": tokens / window if window > 0.2 else None,
"total_s": end - start,
"prompt_tokens": prompt_tokens,
"completion_tokens": tokens,
}
def discover_models(base_url: str, timeout: float) -> list[str]:
with urllib.request.urlopen(f"{base_url}/models", timeout=timeout) as resp:
data = json.load(resp).get("data", [])
# helexa extension: prefer loaded models; plain OpenAI lists lack
# the field, in which case take everything.
loaded = [m["id"] for m in data if m.get("loaded")]
return loaded or [m["id"] for m in data]
def median(values):
vals = [v for v in values if v is not None]
return statistics.median(vals) if vals else None
def main():
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("--base-url", default="http://hanzalova.internal:31313/v1")
ap.add_argument("--model", action="append", help="repeatable; default: all loaded models")
ap.add_argument("--runs", type=int, default=3, help="measured runs per cell (after 1 warmup)")
ap.add_argument(
"--prompt-tokens",
default="128,4096",
help="comma-separated approximate prompt sizes",
)
ap.add_argument("--max-tokens", type=int, default=128)
ap.add_argument("--timeout", type=float, default=600.0)
ap.add_argument("--json", help="append JSON rows to this file")
ap.add_argument("--label", default="helexa", help="engine label for the output rows")
args = ap.parse_args()
models = args.model or discover_models(args.base_url, args.timeout)
sizes = [int(s) for s in args.prompt_tokens.split(",")]
rows = []
for model in models:
for size in sizes:
prompt = build_prompt(size)
try:
one_run(args.base_url, model, prompt, args.max_tokens, args.timeout) # warmup
runs = [
one_run(args.base_url, model, prompt, args.max_tokens, args.timeout)
for _ in range(args.runs)
]
except (RuntimeError, urllib.error.URLError, TimeoutError) as e:
print(f"!! {model} @~{size} tok: {e}", file=sys.stderr)
continue
row = {
"engine": args.label,
"model": model,
"approx_prompt_tokens": size,
"actual_prompt_tokens": runs[0]["prompt_tokens"],
"runs": args.runs,
"ttft_s_median": round(median(r["ttft_s"] for r in runs), 3),
"decode_tps_median": round(median(r["decode_tps"] for r in runs), 1),
"total_s_median": round(median(r["total_s"] for r in runs), 3),
"completion_tokens": runs[0]["completion_tokens"],
"timestamp": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
}
rows.append(row)
print(f".. {model} @~{size} tok done", file=sys.stderr)
print(f"\n| engine | model | prompt tok | TTFT (s) | decode tok/s | total (s) |")
print("|---|---|---:|---:|---:|---:|")
for r in rows:
ptok = r["actual_prompt_tokens"] or f"~{r['approx_prompt_tokens']}"
print(
f"| {r['engine']} | {r['model']} | {ptok} | {r['ttft_s_median']} "
f"| {r['decode_tps_median']} | {r['total_s_median']} |"
)
if args.json:
with open(args.json, "a") as f:
for r in rows:
f.write(json.dumps(r) + "\n")
if __name__ == "__main__":
main()

221
script/dump_reference.py Normal file
View File

@@ -0,0 +1,221 @@
#!/usr/bin/env python3
"""Capture numerical-reference fixtures from HF transformers (#15).
Runs the reference Python implementation of an architecture neuron
serves (today: qwen3_5) on a fixed input and dumps the tensors a
companion Rust test (crates/neuron/tests/numerical_reference.rs)
replays and compares against. The fixtures pin the README's
"implemented in this repository, ported against the HuggingFace
reference" claim to checked-in numbers.
Cases:
text — a fixed >64-token prompt (long enough that neuron's
chunked delta-rule prefill path is exercised), dumping
the token ids and the final-position logits.
vision — a deterministic synthetic 448x448 image (factor-aligned,
so resize is the identity and pixel-level preprocessing
parity is part of what the comparison validates) plus a
short prompt, dumping the expanded token ids, the image
PNG, the LM grid, the vision tower's post-merger output,
and the final-position logits.
Fixture layout (one directory per model+case):
manifest.json — model id, case, token ids, shapes, versions
<name>.f32 — raw little-endian f32 tensor data
image.png — (vision only) the input image
Usage (on a host with torch + transformers and the model snapshot):
python3 script/dump_reference.py \
--model-path /path/to/hf/snapshot --case text \
--out crates/neuron/tests/fixtures/numerical/qwen3_5-0.8b-text
Regenerate fixtures whenever the pinned model snapshot or the
transformers reference implementation changes; record both versions
from the manifest in the commit message.
"""
import argparse
import json
import os
import struct
import sys
# ---------------------------------------------------------------------------
# Compat shim: transformers 5.9 constructs kernels-hub repository
# objects at import time without the revision/version that kernels
# 0.15 requires. The hub kernels are never used here
# (USE_HUB_KERNELS=NO below); the constructors just must not throw.
os.environ.setdefault("USE_HUB_KERNELS", "NO")
try:
import kernels.layer.layer as _kl
import kernels.layer.func as _kf
def _patch(cls):
orig = cls.__init__
def patched(self, *a, **kw):
if "revision" not in kw and "version" not in kw:
kw["revision"] = "main"
orig(self, *a, **kw)
cls.__init__ = patched
_patch(_kl.LayerRepository)
_patch(_kf.FuncRepository)
except Exception: # noqa: BLE001 — older/newer kernels may not need it
pass
# ---------------------------------------------------------------------------
import torch # noqa: E402
# Long enough (>64 tokens) that neuron's replay takes the chunked
# delta-rule prefill path; plain prose so the tokenization is stable.
TEXT_PROMPT = (
"The helexa fleet serves near-frontier language models on consumer "
"graphics cards. Each host runs a small daemon that discovers its "
"hardware, loads the configured models, and answers OpenAI-compatible "
"requests over the private mesh network. The gateway routes each "
"request to the host that already holds the model, restores any "
"cached prefix state, and streams the generated tokens back to the "
"caller one chunk at a time. Operators care about three numbers: the "
"time to the first token, the steady decode rate, and the time a "
"cold model takes to become ready after a deploy. This paragraph "
"exists only to be tokenized identically by two implementations."
)
VISION_PROMPT = "Describe this image in one sentence."
def write_f32(path, tensor):
data = tensor.detach().to(torch.float32).cpu().contiguous().reshape(-1)
with open(path, "wb") as f:
f.write(struct.pack(f"<{data.numel()}f", *data.tolist()))
def synthetic_image(size=448):
"""Deterministic, NON-periodic RGB pattern. Every patch must be
unique: periodic patterns (checkerboards) make many patches exact
duplicates, and attention over near-identical keys is
ill-conditioned — tiny dtype rounding then amplifies chaotically
and the fixture comparison drowns in noise. The x*y term breaks
all translational symmetry while staying byte-deterministic."""
from PIL import Image
img = Image.new("RGB", (size, size))
px = img.load()
for y in range(size):
for x in range(size):
r = (x * 255) // size
g = (y * 255) // size
b = (x * y) % 251
px[x, y] = (r, g, b)
return img
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model-path", required=True, help="HF snapshot dir or repo id")
ap.add_argument("--case", choices=["text", "vision"], required=True)
ap.add_argument("--out", required=True, help="fixture directory to write")
ap.add_argument("--device", default="cuda", choices=["cuda", "cpu"])
ap.add_argument(
"--dtype",
default="float32",
choices=["float32", "bfloat16"],
help="reference compute dtype. float32 (default) pins the math "
"itself — the Rust replay compares f32-to-f32 and implementation "
"bugs are not masked by (or blamed on) bf16 rounding chaos.",
)
ap.add_argument(
"--processor-path",
default=None,
help="where to load the tokenizer/processor from (defaults to "
"--model-path; pass the repo id with HF_HUB_CACHE pointed at a "
"writable scratch dir when the local snapshot is missing "
"preprocessor_config.json)",
)
args = ap.parse_args()
import transformers
from transformers import AutoProcessor, AutoTokenizer
from transformers.models.qwen3_5.modeling_qwen3_5 import (
Qwen3_5ForConditionalGeneration,
)
os.makedirs(args.out, exist_ok=True)
manifest = {
"model_path": args.model_path,
"case": args.case,
"transformers_version": transformers.__version__,
"torch_version": torch.__version__,
"files": {},
}
dtype = torch.float32 if args.dtype == "float32" else torch.bfloat16
manifest["dtype"] = args.dtype
model = Qwen3_5ForConditionalGeneration.from_pretrained(
args.model_path, dtype=dtype, device_map=args.device
)
model.eval()
if args.case == "text":
tok = AutoTokenizer.from_pretrained(args.processor_path or args.model_path)
ids = tok(TEXT_PROMPT, return_tensors="pt").input_ids
manifest["prompt"] = TEXT_PROMPT
manifest["token_ids"] = ids[0].tolist()
with torch.no_grad():
logits = model(input_ids=ids.to(model.device)).logits[0, -1]
write_f32(os.path.join(args.out, "logits.f32"), logits)
manifest["files"]["logits"] = {"file": "logits.f32", "shape": [logits.shape[-1]]}
else:
processor = AutoProcessor.from_pretrained(args.processor_path or args.model_path)
img = synthetic_image()
img.save(os.path.join(args.out, "image.png"))
messages = [
{
"role": "user",
"content": [
{"type": "image", "image": img},
{"type": "text", "text": VISION_PROMPT},
],
}
]
inputs = processor.apply_chat_template(
messages,
add_generation_prompt=True,
tokenize=True,
return_dict=True,
return_tensors="pt",
)
manifest["prompt"] = VISION_PROMPT
manifest["token_ids"] = inputs["input_ids"][0].tolist()
manifest["image_grid_thw"] = inputs["image_grid_thw"][0].tolist()
with torch.no_grad():
visual_out = model.model.visual(
inputs["pixel_values"].to(model.device, dtype),
grid_thw=inputs["image_grid_thw"].to(model.device),
)
# transformers 5.x returns BaseModelOutputWithPooling:
# pooler_output is the post-merger embedding the LM
# splices (= neuron's VisionTower::forward output);
# last_hidden_state is the pre-merger grid.
if hasattr(visual_out, "pooler_output"):
visual_out = visual_out.pooler_output
logits = model(
**{k: v.to(model.device) for k, v in inputs.items()}
).logits[0, -1]
write_f32(os.path.join(args.out, "visual_out.f32"), visual_out)
manifest["files"]["visual_out"] = {
"file": "visual_out.f32",
"shape": list(visual_out.shape),
}
write_f32(os.path.join(args.out, "logits.f32"), logits)
manifest["files"]["logits"] = {"file": "logits.f32", "shape": [logits.shape[-1]]}
with open(os.path.join(args.out, "manifest.json"), "w") as f:
json.dump(manifest, f, indent=2)
print(f"wrote fixture: {args.out}", file=sys.stderr)
if __name__ == "__main__":
main()