Compare commits
59 Commits
d311c8ca7a
...
feat/versi
| Author | SHA1 | Date | |
|---|---|---|---|
|
42da25a37c
|
|||
| 30d50d6215 | |||
|
9a312098dd
|
|||
| 98e9749f22 | |||
|
ec764a2cac
|
|||
| 4c1bdba31d | |||
|
988ef5afc2
|
|||
| a1450789d2 | |||
|
2eaa776d85
|
|||
|
7918995e5a
|
|||
| 538cc87572 | |||
|
1c4b53cbf1
|
|||
| 49a8dbcd28 | |||
|
90e971dcf5
|
|||
|
92273eb936
|
|||
| 128b3818cb | |||
|
812d191e50
|
|||
|
2a9def6d2d
|
|||
| ddb331e1a3 | |||
|
df0bf4c518
|
|||
| a1952a4522 | |||
|
4f266dbd82
|
|||
| 43a6d96d5f | |||
|
3fd1989b2b
|
|||
| f7952547e7 | |||
|
7e66f77851
|
|||
|
e629e1872c
|
|||
| bb558451db | |||
|
c5378d532d
|
|||
| 9f383e7bc7 | |||
|
569c528c4b
|
|||
| 06e4ffc25c | |||
|
a2e73a8907
|
|||
|
8f6f1d3205
|
|||
| b0d0b939af | |||
|
6a36d15ef1
|
|||
| b463439416 | |||
|
716558c8ff
|
|||
|
112e4e124a
|
|||
|
dc6feec6dc
|
|||
| 02f20bc9e1 | |||
|
2a231e49de
|
|||
|
2dadea5d8d
|
|||
|
9bb9678f93
|
|||
|
df9c490614
|
|||
|
f5fa840dfb
|
|||
|
7557c5e877
|
|||
|
91e95ca979
|
|||
|
1a74cb0c56
|
|||
|
60f5598542
|
|||
|
7945240646
|
|||
|
0c74d89d15
|
|||
|
c94a2ae755
|
|||
|
99920dd322
|
|||
|
c4f239ceb9
|
|||
|
ac445c1569
|
|||
|
abc6e605b8
|
|||
|
4f2957af9e
|
|||
|
75cd088b61
|
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
134
.gitea/workflows/deploy-dev.yml
Normal file
134
.gitea/workflows/deploy-dev.yml
Normal 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'
|
||||
@@ -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 '
|
||||
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
|
||||
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 '
|
||||
sudo /usr/bin/systemctl daemon-reload
|
||||
sudo /usr/bin/systemctl start cortex.service'
|
||||
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
|
||||
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
|
||||
sudo /usr/bin/systemctl daemon-reload
|
||||
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 }} '
|
||||
if systemctl is-active --quiet neuron.service; then
|
||||
sudo /usr/bin/systemctl stop neuron.service
|
||||
fi'
|
||||
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
|
||||
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 "${pkg}"
|
||||
fi
|
||||
sudo /usr/bin/systemctl daemon-reload
|
||||
sudo /usr/bin/systemctl start neuron.service
|
||||
|
||||
- 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 }}
|
||||
else
|
||||
sudo /usr/bin/dnf install --refresh --allowerasing -y helexa-neuron-${{ matrix.flavour }}
|
||||
fi"
|
||||
# ── 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.
|
||||
|
||||
60
CLAUDE.md
60
CLAUDE.md
@@ -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
84
Cargo.lock
generated
@@ -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"
|
||||
|
||||
12
Cargo.toml
12
Cargo.toml
@@ -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
190
README.md
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
119
crates/cortex-core/src/build_info.rs
Normal file
119
crates/cortex-core/src/build_info.rs
Normal 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());
|
||||
}
|
||||
}
|
||||
@@ -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.
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
pub mod anthropic;
|
||||
pub mod build_info;
|
||||
pub mod catalogue;
|
||||
pub mod config;
|
||||
pub mod discovery;
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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>,
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
|
||||
178
crates/cortex-gateway/src/anthropic_sse.rs
Normal file
178
crates/cortex-gateway/src/anthropic_sse.rs
Normal 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")
|
||||
}
|
||||
@@ -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) => {
|
||||
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()
|
||||
}
|
||||
if !resp.status().is_success() {
|
||||
metrics::counter!("cortex_request_errors_total", &labels).increment(1);
|
||||
}
|
||||
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 {
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
pub mod anthropic_sse;
|
||||
pub mod evictor;
|
||||
pub mod handlers;
|
||||
pub mod metrics;
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -197,6 +197,7 @@ fn parse_status(s: &str) -> ModelStatus {
|
||||
"unloaded" => ModelStatus::Unloaded,
|
||||
"reloading" => ModelStatus::Reloading,
|
||||
"loading" => ModelStatus::Loading,
|
||||
"recovering" => ModelStatus::Recovering,
|
||||
_ => ModelStatus::Loaded,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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?;
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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}"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 3a–3e and 6a/6b that the
|
||||
|
||||
38
crates/helexa-bench/Cargo.toml
Normal file
38
crates/helexa-bench/Cargo.toml
Normal 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"] }
|
||||
159
crates/helexa-bench/src/client.rs
Normal file
159
crates/helexa-bench/src/client.rs
Normal 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())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
210
crates/helexa-bench/src/config.rs
Normal file
210
crates/helexa-bench/src/config.rs
Normal 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(())
|
||||
});
|
||||
}
|
||||
}
|
||||
12
crates/helexa-bench/src/lib.rs
Normal file
12
crates/helexa-bench/src/lib.rs
Normal 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;
|
||||
126
crates/helexa-bench/src/main.rs
Normal file
126
crates/helexa-bench/src/main.rs
Normal 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(())
|
||||
}
|
||||
106
crates/helexa-bench/src/report.rs
Normal file
106
crates/helexa-bench/src/report.rs
Normal 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("—"));
|
||||
}
|
||||
}
|
||||
238
crates/helexa-bench/src/scenario.rs
Normal file
238
crates/helexa-bench/src/scenario.rs
Normal 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);
|
||||
}
|
||||
}
|
||||
400
crates/helexa-bench/src/store.rs
Normal file
400
crates/helexa-bench/src/store.rs
Normal 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);
|
||||
}
|
||||
}
|
||||
250
crates/helexa-bench/src/sweep.rs
Normal file
250
crates/helexa-bench/src/sweep.rs
Normal 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",
|
||||
}
|
||||
}
|
||||
132
crates/helexa-bench/tests/sweep_integration.rs
Normal file
132
crates/helexa-bench/tests/sweep_integration.rs
Normal 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);
|
||||
}
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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`
|
||||
|
||||
@@ -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",
|
||||
&[
|
||||
&format!("--query-gpu={NVIDIA_SMI_DISCOVERY_QUERY}"),
|
||||
"--format=csv,noheader,nounits",
|
||||
],
|
||||
)
|
||||
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"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
@@ -115,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"
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -165,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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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;
|
||||
@@ -395,6 +396,42 @@ impl Qwen3_5Model {
|
||||
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> {
|
||||
let minf = f32::NEG_INFINITY;
|
||||
let mask: Vec<_> = (0..tgt)
|
||||
@@ -404,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.
|
||||
@@ -440,9 +504,16 @@ impl Qwen3_5Model {
|
||||
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,
|
||||
@@ -450,19 +521,15 @@ impl Qwen3_5Model {
|
||||
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)?;
|
||||
|
||||
// Vision path: splice image embeddings at `image_token_id`
|
||||
// positions and build interleaved M-RoPE cos/sin so image tokens
|
||||
// carry their 2D (lm_gh × lm_gw) grid coordinates. Text / decode skip the
|
||||
// device→host id copy entirely and take the plain-RoPE fast path
|
||||
// — bit-for-bit the pre-M-RoPE behaviour when `rope_delta == 0`.
|
||||
let (cos, sin) = if let (Some(img), Some(tok_id)) = (image_embeds, image_token_id) {
|
||||
// Token ids on CPU — reused for the splice + position ids.
|
||||
// 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) {
|
||||
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() {
|
||||
if *id == tok_id {
|
||||
@@ -472,9 +539,9 @@ 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,
|
||||
);
|
||||
@@ -485,7 +552,20 @@ impl Qwen3_5Model {
|
||||
let img = img.to_dtype(self.dtype)?;
|
||||
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;
|
||||
@@ -615,9 +695,173 @@ impl Qwen3_5ForCausalLM {
|
||||
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)]
|
||||
|
||||
299
crates/neuron/src/harness/arch/qwen3_5/snapshot.rs
Normal file
299
crates/neuron/src/harness/arch/qwen3_5/snapshot.rs
Normal 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}");
|
||||
}
|
||||
}
|
||||
@@ -457,19 +457,27 @@ impl VisionTower {
|
||||
}
|
||||
}
|
||||
|
||||
// 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)?; // (n, hidden), pos_embed dtype
|
||||
let wt = Tensor::from_vec(std::mem::take(&mut wts[corner]), (n, 1), &self.device)?
|
||||
.to_dtype(self.dtype)?;
|
||||
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,
|
||||
});
|
||||
}
|
||||
Ok(acc.expect("4 corners accumulated"))
|
||||
acc.expect("4 corners accumulated")
|
||||
.to_dtype(self.dtype)
|
||||
.map_err(Into::into)
|
||||
}
|
||||
|
||||
/// Encode one image.
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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,
|
||||
@@ -353,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,
|
||||
}
|
||||
}
|
||||
@@ -366,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(),
|
||||
}
|
||||
}
|
||||
@@ -871,21 +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)
|
||||
})?;
|
||||
|
||||
// pixel→LM-grid divisor (patch×merge) for this tower; each image's
|
||||
// LM grid is (h/factor, w/factor) (#14 dynamic resolution).
|
||||
let factor = arch.vision_grid_factor().ok_or_else(|| {
|
||||
anyhow::anyhow!("ForwardLogitsWithImages: loaded model has no vision tower")
|
||||
})?;
|
||||
|
||||
// Encode every image on the worker's device, collecting per-image
|
||||
// post-merger embeddings as device-resident tensors plus their LM
|
||||
// grids (for the interleaved-M-RoPE position ids).
|
||||
let mut per_image: Vec<Tensor> = Vec::with_capacity(images.len());
|
||||
let mut grids: Vec<(usize, usize)> = 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,
|
||||
@@ -895,20 +1072,26 @@ fn forward_logits_with_images(
|
||||
img.h,
|
||||
img.w,
|
||||
);
|
||||
grids.push((img.h / factor, img.w / factor));
|
||||
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, &grids)?;
|
||||
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()?
|
||||
@@ -989,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()));
|
||||
}
|
||||
@@ -1004,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(),
|
||||
@@ -1023,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()));
|
||||
}
|
||||
|
||||
@@ -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
|
||||
@@ -105,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
|
||||
@@ -192,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
|
||||
@@ -224,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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
266
crates/neuron/src/harness/prefix_cache.rs
Normal file
266
crates/neuron/src/harness/prefix_cache.rs
Normal 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());
|
||||
}
|
||||
}
|
||||
@@ -55,12 +55,23 @@ pub struct PreprocessProfile {
|
||||
pub image_std: [f32; 3],
|
||||
}
|
||||
|
||||
/// Default pixel budget for Qwen3.6 (`256² … 1024²` → 64 … 1024 LM
|
||||
/// tokens/image). Generous for documents/OCR, bounded for serving on
|
||||
/// 2×RTX5090. Operators tune with `NEURON_VISION_MIN_PIXELS` /
|
||||
/// `NEURON_VISION_MAX_PIXELS` (matching the other `NEURON_VISION_*` knobs).
|
||||
/// 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;
|
||||
const QWEN3_6_MAX_PIXELS: u32 = 1_048_576;
|
||||
|
||||
fn env_pixels(name: &str, default: u32) -> u32 {
|
||||
std::env::var(name)
|
||||
@@ -72,15 +83,19 @@ fn env_pixels(name: &str, default: u32) -> u32 {
|
||||
impl PreprocessProfile {
|
||||
/// 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`], overridable via the
|
||||
/// `NEURON_VISION_MIN_PIXELS` / `NEURON_VISION_MAX_PIXELS` env vars.
|
||||
/// The budget is clamped sane: `min ≥ factor²` (at least one LM token)
|
||||
/// and `max ≥ min`.
|
||||
/// [`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);
|
||||
let max_pixels = env_pixels("NEURON_VISION_MAX_PIXELS", QWEN3_6_MAX_PIXELS).max(min_pixels);
|
||||
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 {
|
||||
factor,
|
||||
min_pixels,
|
||||
@@ -388,6 +403,28 @@ mod tests {
|
||||
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
|
||||
|
||||
234
crates/neuron/src/harness/speculative.rs
Normal file
234
crates/neuron/src/harness/speculative.rs
Normal 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()
|
||||
);
|
||||
}
|
||||
}
|
||||
171
crates/neuron/src/harness/tp/isq.rs
Normal file
171
crates/neuron/src/harness/tp/isq.rs
Normal 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 32–256 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);
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
@@ -93,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(),
|
||||
@@ -245,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.
|
||||
@@ -324,6 +414,8 @@ impl WorkerPool {
|
||||
workers,
|
||||
exe,
|
||||
leader_worker,
|
||||
#[cfg(feature = "cuda")]
|
||||
leader_comm: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -404,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(())
|
||||
}
|
||||
|
||||
@@ -628,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
|
||||
@@ -767,17 +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(
|
||||
leader_handle,
|
||||
tokens,
|
||||
offset,
|
||||
image_token_id,
|
||||
image_data_uris,
|
||||
chunk_size,
|
||||
)
|
||||
.await;
|
||||
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,
|
||||
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 {
|
||||
@@ -887,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).
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -120,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`.
|
||||
@@ -168,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,
|
||||
|
||||
@@ -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()
|
||||
)
|
||||
})?;
|
||||
|
||||
@@ -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> {
|
||||
@@ -585,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 ────────────────────────────────────────────────────────────
|
||||
@@ -828,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 ─────────────────────────────────────────────────────
|
||||
@@ -885,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(),
|
||||
@@ -902,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)?;
|
||||
|
||||
@@ -950,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,
|
||||
@@ -962,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)?;
|
||||
|
||||
@@ -987,6 +1097,38 @@ impl TpQwen3_5Model {
|
||||
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.
|
||||
@@ -1365,6 +1507,16 @@ impl TpQwen3_5ForCausalLM {
|
||||
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
|
||||
}
|
||||
@@ -1387,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)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -82,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(),
|
||||
@@ -164,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 {
|
||||
@@ -172,6 +209,7 @@ impl WorkerState {
|
||||
config,
|
||||
nccl: NcclState::new(),
|
||||
models: HashMap::new(),
|
||||
kv_snapshots: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -211,6 +249,18 @@ impl WorkerState {
|
||||
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,
|
||||
}
|
||||
@@ -592,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() {
|
||||
@@ -600,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
|
||||
}
|
||||
|
||||
@@ -6,4 +6,5 @@ pub mod discovery;
|
||||
pub mod harness;
|
||||
pub mod health;
|
||||
pub mod startup;
|
||||
pub mod version;
|
||||
pub mod wire;
|
||||
|
||||
@@ -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,8 +232,16 @@ 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(®istry, &default_models, &state_for_prewarm.activation)
|
||||
.await;
|
||||
startup::load_default_models(
|
||||
®istry,
|
||||
&default_models,
|
||||
&state_for_prewarm.activation,
|
||||
state_for_prewarm
|
||||
.discovery
|
||||
.cuda_unavailable_reason
|
||||
.as_deref(),
|
||||
)
|
||||
.await;
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
@@ -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();
|
||||
|
||||
83
crates/neuron/src/version.rs
Normal file
83
crates/neuron/src/version.rs
Normal 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")));
|
||||
}
|
||||
}
|
||||
@@ -40,7 +40,7 @@ async fn test_load_default_models_skips_unknown_harness() {
|
||||
];
|
||||
|
||||
let activation = ActivationTracker::new(&specs);
|
||||
startup::load_default_models(®istry, &specs, &activation).await;
|
||||
startup::load_default_models(®istry, &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(®istry, &[], &activation).await;
|
||||
startup::load_default_models(®istry, &[], &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(®istry, &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
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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}"
|
||||
);
|
||||
}
|
||||
|
||||
67
crates/neuron/tests/fixtures/numerical/README.md
vendored
Normal file
67
crates/neuron/tests/fixtures/numerical/README.md
vendored
Normal 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
|
||||
```
|
||||
BIN
crates/neuron/tests/fixtures/numerical/qwen3_5-0.8b-text/logits.f32
vendored
Normal file
BIN
crates/neuron/tests/fixtures/numerical/qwen3_5-0.8b-text/logits.f32
vendored
Normal file
Binary file not shown.
143
crates/neuron/tests/fixtures/numerical/qwen3_5-0.8b-text/manifest.json
vendored
Normal file
143
crates/neuron/tests/fixtures/numerical/qwen3_5-0.8b-text/manifest.json
vendored
Normal 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
|
||||
]
|
||||
}
|
||||
BIN
crates/neuron/tests/fixtures/numerical/qwen3_5-0.8b-vision/image.png
vendored
Normal file
BIN
crates/neuron/tests/fixtures/numerical/qwen3_5-0.8b-vision/image.png
vendored
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 72 KiB |
BIN
crates/neuron/tests/fixtures/numerical/qwen3_5-0.8b-vision/logits.f32
vendored
Normal file
BIN
crates/neuron/tests/fixtures/numerical/qwen3_5-0.8b-vision/logits.f32
vendored
Normal file
Binary file not shown.
247
crates/neuron/tests/fixtures/numerical/qwen3_5-0.8b-vision/manifest.json
vendored
Normal file
247
crates/neuron/tests/fixtures/numerical/qwen3_5-0.8b-vision/manifest.json
vendored
Normal 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
|
||||
]
|
||||
}
|
||||
BIN
crates/neuron/tests/fixtures/numerical/qwen3_5-0.8b-vision/visual_out.f32
vendored
Normal file
BIN
crates/neuron/tests/fixtures/numerical/qwen3_5-0.8b-vision/visual_out.f32
vendored
Normal file
Binary file not shown.
BIN
crates/neuron/tests/fixtures/numerical/qwen3_6-27b-text/logits.f32
vendored
Normal file
BIN
crates/neuron/tests/fixtures/numerical/qwen3_6-27b-text/logits.f32
vendored
Normal file
Binary file not shown.
143
crates/neuron/tests/fixtures/numerical/qwen3_6-27b-text/manifest.json
vendored
Normal file
143
crates/neuron/tests/fixtures/numerical/qwen3_6-27b-text/manifest.json
vendored
Normal 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
|
||||
]
|
||||
}
|
||||
339
crates/neuron/tests/numerical_reference.rs
Normal file
339
crates/neuron/tests/numerical_reference.rs
Normal 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");
|
||||
}
|
||||
3
data/helexa-bench-sysusers.conf
Normal file
3
data/helexa-bench-sysusers.conf
Normal 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
21
data/helexa-bench.service
Normal 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
154
doc/benchmarks.md
Normal 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
129
doc/learnings/p1.md
Normal 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
251
doc/learnings/p2.md
Normal 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 1–2 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".
|
||||
@@ -5,7 +5,7 @@ Sourced from beast's local cache on 2026-06-01:
|
||||
|
||||
Single source of truth for Stages A–D 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
50
helexa-bench.example.toml
Normal 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"
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
98
rpm/helexa-bench-prerelease.spec
Normal file
98
rpm/helexa-bench-prerelease.spec
Normal 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.
|
||||
@@ -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
205
script/bench.py
Executable 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
221
script/dump_reference.py
Normal 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()
|
||||
Reference in New Issue
Block a user