Compare commits
73 Commits
feat/F1-th
...
feat/98-ba
| Author | SHA1 | Date | |
|---|---|---|---|
|
d3ef9513f8
|
|||
|
ec7a5b750a
|
|||
|
7f2fcc3527
|
|||
|
9ac1a1586c
|
|||
|
05b6c70a3a
|
|||
| 9c037c5a0b | |||
|
cb6ddbbd98
|
|||
| e824ea3274 | |||
|
9139b07d71
|
|||
|
4762537225
|
|||
| 9b635bfe5c | |||
|
28aa06b768
|
|||
| fc95faaed6 | |||
|
2df829ed73
|
|||
| f611e3479b | |||
|
41ec05e40c
|
|||
|
3242cce39b
|
|||
| 35041f11de | |||
|
40fecb2cf7
|
|||
| 8d50918a1a | |||
|
29a7054c23
|
|||
|
7489a49d24
|
|||
| 341f97f705 | |||
|
fc56eafbf7
|
|||
|
a157c8fa27
|
|||
|
201cc54a7b
|
|||
| 781d92d9a3 | |||
|
9bf13f09dd
|
|||
|
63fb8c6e76
|
|||
|
a1426f177c
|
|||
| a9f5ac7ab4 | |||
|
0126144a93
|
|||
|
869033d08e
|
|||
|
7c5011e5d3
|
|||
|
d1366614f9
|
|||
| 1884535542 | |||
|
140f9124ee
|
|||
| 3b77dbfa32 | |||
|
842caf5d4d
|
|||
| a488ade675 | |||
|
d472b6428a
|
|||
| 0488768afc | |||
|
209bf58207
|
|||
| 1ea7425585 | |||
|
a415454962
|
|||
| 06a36566d1 | |||
| 4cb52e3144 | |||
|
afc1f7a706
|
|||
|
6f956dfda3
|
|||
|
6e0f15c888
|
|||
| 66eb9f558f | |||
|
f96a2e7ed3
|
|||
| b17b555a3d | |||
|
13daf95514
|
|||
| 319b01e0b2 | |||
|
6731adca51
|
|||
|
7e11a7688c
|
|||
|
5600575ba2
|
|||
|
bc7476bf1b
|
|||
|
5a8f6bc7b3
|
|||
|
452d7d9b3d
|
|||
|
21eb211d6a
|
|||
|
508b326bf7
|
|||
|
0de99a8cc7
|
|||
|
f4117224fc
|
|||
|
ce29e0c171
|
|||
|
1bf3348c8c
|
|||
|
7c12b9ea98
|
|||
|
c596519dbd
|
|||
|
a6b1fdc33d
|
|||
|
8dd82776f1
|
|||
|
8600d4fbf2
|
|||
|
2348cc2234
|
@@ -66,6 +66,7 @@ jobs:
|
||||
build_cortex: ${{ steps.changes.outputs.build_cortex }}
|
||||
build_neuron: ${{ steps.changes.outputs.build_neuron }}
|
||||
build_bench: ${{ steps.changes.outputs.build_bench }}
|
||||
build_upstream: ${{ steps.changes.outputs.build_upstream }}
|
||||
check_rust: ${{ steps.changes.outputs.check_rust }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
@@ -86,7 +87,14 @@ jobs:
|
||||
# rpmvercmp ranks digit-prefixed segments above alpha ones.
|
||||
# The SHA stays only as a debug identifier; sort order is
|
||||
# decided entirely by the timestamp.
|
||||
COMMIT_TIMESTAMP=$(git log -1 --format=%cd --date=format:%Y%m%d%H%M%S HEAD)
|
||||
# format-local + TZ=UTC0 renders the stamp in UTC regardless
|
||||
# of the committer's recorded timezone. Plain `format:` uses
|
||||
# each commit's own TZ offset — local commits (UTC+3) and
|
||||
# Gitea server-side merge commits (UTC) interleaved
|
||||
# non-monotonically, letting an older build out-rank newer
|
||||
# ones in RPM EVR comparison (the 2026-07-01 "235959" stamp
|
||||
# that froze the fleet on a stale build).
|
||||
COMMIT_TIMESTAMP=$(TZ=UTC0 git log -1 --format=%cd --date=format-local:%Y%m%d%H%M%S HEAD)
|
||||
RELEASE="0.1.${COMMIT_TIMESTAMP}.git${SHORT_SHA}"
|
||||
echo "version=${VERSION}" >> "$GITHUB_OUTPUT"
|
||||
echo "release=${RELEASE}" >> "$GITHUB_OUTPUT"
|
||||
@@ -104,6 +112,7 @@ jobs:
|
||||
BUILD_CORTEX=true
|
||||
BUILD_NEURON=true
|
||||
BUILD_BENCH=true
|
||||
BUILD_UPSTREAM=true
|
||||
CHECK_RUST=true
|
||||
|
||||
if [ "${GITHUB_EVENT_NAME}" = "push" ]; then
|
||||
@@ -149,6 +158,7 @@ jobs:
|
||||
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$'
|
||||
UPSTREAM_RE='^crates/helexa-upstream/|^crates/cortex-core/|^Cargo\.toml$|^Cargo\.lock$|^rpm/helexa-upstream-prerelease\.spec$|^data/helexa-upstream|^helexa-upstream\.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$'
|
||||
@@ -156,10 +166,12 @@ jobs:
|
||||
CORTEX_BASE=$(base_for cortex)
|
||||
NEURON_BASE=$(base_for helexa-neuron-blackwell)
|
||||
BENCH_BASE=$(base_for helexa-bench)
|
||||
UPSTREAM_BASE=$(base_for helexa-upstream)
|
||||
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
|
||||
BUILD_UPSTREAM=$(decide "$UPSTREAM_BASE" "$UPSTREAM_RE")
|
||||
if [ "$BUILD_CORTEX" = "true" ] || [ "$BUILD_NEURON" = "true" ] || [ "$BUILD_BENCH" = "true" ] || [ "$BUILD_UPSTREAM" = "true" ]; then
|
||||
CHECK_RUST=true
|
||||
else
|
||||
CHECK_RUST=$(decide "$CORTEX_BASE" "$RUST_RE")
|
||||
@@ -170,8 +182,9 @@ jobs:
|
||||
echo "build_cortex=${BUILD_CORTEX}" >> "$GITHUB_OUTPUT"
|
||||
echo "build_neuron=${BUILD_NEURON}" >> "$GITHUB_OUTPUT"
|
||||
echo "build_bench=${BUILD_BENCH}" >> "$GITHUB_OUTPUT"
|
||||
echo "build_upstream=${BUILD_UPSTREAM}" >> "$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}"
|
||||
echo "### change detection: build_cortex=${BUILD_CORTEX} build_neuron=${BUILD_NEURON} build_bench=${BUILD_BENCH} build_upstream=${BUILD_UPSTREAM} 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
|
||||
@@ -303,6 +316,45 @@ jobs:
|
||||
path: artifacts/helexa-bench
|
||||
retention-days: 1
|
||||
|
||||
build-upstream:
|
||||
name: Build helexa-upstream binary
|
||||
timeout-minutes: 25
|
||||
needs: prepare
|
||||
if: needs.prepare.outputs.build_upstream == 'true'
|
||||
# Pure-Rust, non-CUDA binary — same runner as cortex/bench.
|
||||
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 }}
|
||||
# helexa-upstream uses the sqlx runtime query API (no compile-time
|
||||
# query macros), so it builds without a database or a .sqlx cache.
|
||||
# Set OFFLINE defensively so a stray macro can never reach for a DB.
|
||||
SQLX_OFFLINE: "true"
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
ref: ${{ inputs.ref }}
|
||||
|
||||
- name: Build helexa-upstream (release, sccache escalation)
|
||||
run: script/ci-cargo-escalate.sh cargo build --release -p helexa-upstream
|
||||
|
||||
- name: Stage binary
|
||||
run: |
|
||||
mkdir --parents artifacts
|
||||
cp target/release/helexa-upstream artifacts/helexa-upstream
|
||||
./artifacts/helexa-upstream --version || true
|
||||
|
||||
- uses: actions/upload-artifact@v3
|
||||
with:
|
||||
name: upstream-fc43
|
||||
path: artifacts/helexa-upstream
|
||||
retention-days: 1
|
||||
|
||||
build-neuron:
|
||||
name: Build neuron-${{ matrix.flavour }}
|
||||
timeout-minutes: 35
|
||||
@@ -332,7 +384,10 @@ jobs:
|
||||
cuda_home: /usr/local/cuda-13.0
|
||||
build_jobs: 8
|
||||
nvcc_threads: 4
|
||||
cargo_features: "cuda cudnn"
|
||||
# flash-attn on blackwell first (#95): beast carries the
|
||||
# agentic prefill pain; ada/ampere follow once the win is
|
||||
# measured. NEURON_FLASH_ATTN=0 is the runtime rollback.
|
||||
cargo_features: "cuda cudnn flash-attn"
|
||||
runs-on: ${{ matrix.runner }}
|
||||
env:
|
||||
SCCACHE_BUCKET: sccache
|
||||
@@ -459,6 +514,44 @@ jobs:
|
||||
path: ~/rpmbuild/RPMS/x86_64/*.rpm
|
||||
retention-days: 7
|
||||
|
||||
package-upstream:
|
||||
name: Package helexa-upstream RPM
|
||||
timeout-minutes: 20
|
||||
needs: [prepare, build-upstream]
|
||||
runs-on: rpm
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
ref: ${{ inputs.ref }}
|
||||
|
||||
- uses: actions/download-artifact@v3
|
||||
with:
|
||||
name: upstream-fc43
|
||||
path: artifacts/
|
||||
|
||||
- name: Build RPM
|
||||
run: |
|
||||
set -eux
|
||||
rm -f ~/.rpmmacros
|
||||
rpmdev-setuptree
|
||||
cp artifacts/helexa-upstream ~/rpmbuild/SOURCES/
|
||||
cp data/helexa-upstream.service ~/rpmbuild/SOURCES/
|
||||
cp data/helexa-upstream-sysusers.conf ~/rpmbuild/SOURCES/
|
||||
cp data/helexa-upstream-firewalld.xml ~/rpmbuild/SOURCES/
|
||||
cp helexa-upstream.example.toml ~/rpmbuild/SOURCES/
|
||||
cp LICENSE ~/rpmbuild/SOURCES/
|
||||
rpmbuild -bb rpm/helexa-upstream-prerelease.spec \
|
||||
--define "upstream_version ${{ needs.prepare.outputs.version }}" \
|
||||
--define "upstream_prerelease ${{ needs.prepare.outputs.release }}" \
|
||||
--undefine dist \
|
||||
--define "dist .fc43"
|
||||
|
||||
- uses: actions/upload-artifact@v3
|
||||
with:
|
||||
name: rpm-upstream-fc43
|
||||
path: ~/rpmbuild/RPMS/x86_64/*.rpm
|
||||
retention-days: 7
|
||||
|
||||
package-neuron:
|
||||
name: Package helexa-neuron-${{ matrix.flavour }} RPM
|
||||
timeout-minutes: 20
|
||||
@@ -508,7 +601,7 @@ jobs:
|
||||
publish:
|
||||
name: Publish to rpm.lair.cafe (unstable)
|
||||
timeout-minutes: 25
|
||||
needs: [lint, test, package-cortex, package-neuron, package-bench]
|
||||
needs: [lint, test, package-cortex, package-neuron, package-bench, package-upstream]
|
||||
# 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
|
||||
@@ -518,10 +611,11 @@ jobs:
|
||||
!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 == 'success' || needs.package-neuron.result == 'success' || needs.package-bench.result == 'success' || needs.package-upstream.result == 'success')
|
||||
&& needs.package-cortex.result != 'failure'
|
||||
&& needs.package-neuron.result != 'failure'
|
||||
&& needs.package-bench.result != 'failure'
|
||||
&& needs.package-upstream.result != 'failure'
|
||||
}}
|
||||
runs-on: rpm
|
||||
concurrency:
|
||||
|
||||
@@ -82,7 +82,9 @@ jobs:
|
||||
# see commit history).
|
||||
cuda-check:
|
||||
name: CUDA type-check
|
||||
timeout-minutes: 35
|
||||
# flash-attn kernel compilation dominates the first uncached run;
|
||||
# sccache + the cargo target cache absorb it afterwards.
|
||||
timeout-minutes: 70
|
||||
runs-on: cuda-13.0
|
||||
# The workflow-level env sets `RUSTC_WRAPPER: sccache`
|
||||
# unconditionally, which hard-fails cargo if the CUDA image
|
||||
@@ -115,7 +117,7 @@ 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:-}"
|
||||
script/ci-cargo-escalate.sh cargo check -p neuron --features cuda --all-targets
|
||||
script/ci-cargo-escalate.sh cargo check -p neuron --features cuda,flash-attn --all-targets
|
||||
|
||||
srpm-cortex:
|
||||
name: Build cortex SRPM
|
||||
|
||||
@@ -274,19 +274,52 @@ jobs:
|
||||
|
||||
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
|
||||
# The probe races real traffic: a deploy publishes a new build
|
||||
# SHA, which is exactly what triggers helexa-bench to re-sweep
|
||||
# every scenario against this neuron — long capability-probe
|
||||
# generations can hold the batch-1 model for minutes. Admission
|
||||
# control answers concurrent requests with 429/503 +
|
||||
# Retry-After (#53/#63); treat those as "busy, try again", not
|
||||
# deploy failure. Overall budget ~6 min.
|
||||
probe_deadline=$(( $(date +%s) + 360 ))
|
||||
attempt=0
|
||||
while :; do
|
||||
attempt=$(( attempt + 1 ))
|
||||
hdrs=$(mktemp)
|
||||
resp=$(curl -sS --max-time 180 -D "${hdrs}" \
|
||||
-H "content-type: application/json" \
|
||||
-d "${probe_body}" http://localhost:13131/v1/chat/completions) || resp=""
|
||||
status=$(awk 'toupper($1) ~ /^HTTP/ {code=$2} END {print code}' "${hdrs}")
|
||||
retry_after=$(awk 'tolower($1) == "retry-after:" {print $2+0; exit}' "${hdrs}")
|
||||
rm -f "${hdrs}"
|
||||
if [ "${status}" = "200" ]; then
|
||||
if printf %s "${resp}" | grep -qi pineapple; then
|
||||
echo "LLM probe passed (attempt ${attempt})"
|
||||
break
|
||||
fi
|
||||
echo "FAIL: probe response missing expected token"
|
||||
printf %s "${resp}" | head -c 2000
|
||||
echo
|
||||
exit 1
|
||||
fi
|
||||
case "${status}" in
|
||||
429|503) ;; # busy/queue-full — retryable
|
||||
*)
|
||||
echo "FAIL: probe request errored (HTTP ${status:-none})"
|
||||
printf %s "${resp}" | head -c 2000
|
||||
echo
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
if [ "$(date +%s)" -ge "${probe_deadline}" ]; then
|
||||
echo "FAIL: model still busy (HTTP ${status}) after probe budget"
|
||||
exit 1
|
||||
fi
|
||||
wait_s=${retry_after:-15}
|
||||
[ "${wait_s}" -ge 5 ] 2>/dev/null || wait_s=15
|
||||
echo "model busy (HTTP ${status}), retrying in ${wait_s}s (attempt ${attempt})"
|
||||
sleep "${wait_s}"
|
||||
done
|
||||
DEPLOY
|
||||
|
||||
- name: Ensure firewalld allows helexa-neuron
|
||||
|
||||
1
Cargo.lock
generated
1
Cargo.lock
generated
@@ -2972,6 +2972,7 @@ dependencies = [
|
||||
"axum",
|
||||
"base64 0.22.1",
|
||||
"candle-core",
|
||||
"candle-flash-attn",
|
||||
"candle-nn",
|
||||
"candle-transformers",
|
||||
"clap",
|
||||
|
||||
@@ -18,6 +18,35 @@ db_path = "/var/lib/helexa-bench/bench.sqlite"
|
||||
[scenarios]
|
||||
prompt_sizes = [128, 4096]
|
||||
max_tokens = 256
|
||||
# Concurrency / agentic-load scenarios (#89), enabled so the 27B baseline
|
||||
# carries p95-under-concurrency data before the F3 A/B gate (#94) needs
|
||||
# it for comparison. Levels mirror the real a0/hermes/opencode fan-out.
|
||||
concurrency_levels = [2, 4, 8]
|
||||
concurrency_prompt_tokens = 512
|
||||
|
||||
# Capability probes (#91) — the reasoning/planning axis the speed
|
||||
# scenarios miss; scored manually via `helexa-bench score` (O7). Enabled
|
||||
# for the same reason: the F3 gate compares 80B-A3B variants against the
|
||||
# 27B on planning quality, so the 27B needs scored artifacts first.
|
||||
[[scenarios.capability_probes]]
|
||||
name = "rust-plan"
|
||||
max_tokens = 4096
|
||||
prompt = """
|
||||
Write an implementation plan for adding rate limiting to an Axum service.
|
||||
Honor existing conventions, call out trade-offs, and sequence the work.
|
||||
"""
|
||||
|
||||
[[scenarios.capability_probes]]
|
||||
name = "debug-reason"
|
||||
max_tokens = 4096
|
||||
prompt = """
|
||||
A Rust axum server streams SSE responses through a reverse proxy. Clients
|
||||
report that streams stall for exactly 60 seconds and then resume, but only
|
||||
when response chunks are small and infrequent. Curling the backend directly
|
||||
never stalls. List the most likely causes in order of probability, explain
|
||||
the mechanism behind each, and describe the smallest experiment that would
|
||||
confirm or eliminate each cause.
|
||||
"""
|
||||
|
||||
# Read-only JSON API consumed by the bench UI (hosted separately) and for
|
||||
# programmatic access. Served alongside the sweep loop.
|
||||
|
||||
@@ -105,3 +105,5 @@ enabled = false
|
||||
# upstream). Override via CORTEX_UPSTREAM__BEARER in prod.
|
||||
# bearer = "replace-with-operator-client-secret"
|
||||
# timeout_secs = 5
|
||||
# How often to flush served-usage counters to upstream for reconciliation (#58).
|
||||
# served_usage_report_interval_secs = 60
|
||||
|
||||
@@ -48,11 +48,18 @@ pub struct UpstreamClientConfig {
|
||||
/// Per-call timeout (seconds) to upstream.
|
||||
#[serde(default = "default_upstream_timeout")]
|
||||
pub timeout_secs: u64,
|
||||
/// How often (seconds) to flush served-usage counters to upstream for
|
||||
/// reconciliation (#58).
|
||||
#[serde(default = "default_served_usage_interval")]
|
||||
pub served_usage_report_interval_secs: u64,
|
||||
}
|
||||
|
||||
fn default_upstream_timeout() -> u64 {
|
||||
5
|
||||
}
|
||||
fn default_served_usage_interval() -> u64 {
|
||||
60
|
||||
}
|
||||
|
||||
/// `[entitlements]` — the local/static [`crate::entitlements::EntitlementProvider`]
|
||||
/// source of truth (#50). Accounts, keys, and hard caps live here; the
|
||||
|
||||
@@ -147,6 +147,39 @@ pub struct CortexModelEntry {
|
||||
/// `true` when any neuron reports this model supports reasoning tokens.
|
||||
#[serde(default)]
|
||||
pub reasoning: bool,
|
||||
// ── Flat ecosystem context-window fields (issue #78) ──────
|
||||
// Duplicates of `limit` under the flat, vLLM-convention key names
|
||||
// (`max_model_len` et al.) that OpenAI-ecosystem clients (Hermes
|
||||
// Agent, vLLM tooling) probe for — they cannot see `limit.context`.
|
||||
// Additive: `limit` stays the opencode-oriented source of truth.
|
||||
// Derived, never set directly — call [`sync_flat_limit`] after the
|
||||
// final `limit` value is known. Omitted (not `0`) when the window
|
||||
// is unknown; absent-vs-zero is load-bearing, as with `cost`.
|
||||
//
|
||||
// [`sync_flat_limit`]: CortexModelEntry::sync_flat_limit
|
||||
/// Served max-seq-len in tokens — mirrors `limit.context`.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub max_model_len: Option<usize>,
|
||||
/// Usable input budget — mirrors `limit.input` when present.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub max_input_tokens: Option<usize>,
|
||||
/// Maximum generation tokens — mirrors `limit.output`.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub max_output_tokens: Option<usize>,
|
||||
}
|
||||
|
||||
impl CortexModelEntry {
|
||||
/// Re-derive the flat ecosystem fields (#78) from `limit`.
|
||||
///
|
||||
/// Must run after the final `limit` is known (post merge/tightening),
|
||||
/// immediately before serialization. Fully overwrites: a `None` limit
|
||||
/// clears the flat fields, so stale values can't survive a merge that
|
||||
/// dropped the limit.
|
||||
pub fn sync_flat_limit(&mut self) {
|
||||
self.max_model_len = self.limit.as_ref().map(|l| l.context);
|
||||
self.max_input_tokens = self.limit.as_ref().and_then(|l| l.input);
|
||||
self.max_output_tokens = self.limit.as_ref().map(|l| l.output);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
|
||||
@@ -116,6 +116,23 @@ pub struct Usage {
|
||||
/// prompt caching lands (#11); `None` until then.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub prompt_tokens_details: Option<PromptTokensDetails>,
|
||||
/// helexa extension (non-OpenAI): server-measured prefill/decode
|
||||
/// timing, so the bench harness can compute true prefill vs decode
|
||||
/// tok/s instead of inferring both from client-side SSE arrival
|
||||
/// (#85). Additive and optional — standard OpenAI clients ignore
|
||||
/// it; cortex forwards usage verbatim so it survives proxying.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub helexa_timing: Option<HelexaTiming>,
|
||||
}
|
||||
|
||||
/// helexa extension carried on [`Usage::helexa_timing`]. Mirrors
|
||||
/// neuron's internal `FinishTiming`. All fields are server-measured;
|
||||
/// `prefill_tokens` is the prefill-rate denominator.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct HelexaTiming {
|
||||
pub prefill_ms: u64,
|
||||
pub decode_ms: u64,
|
||||
pub prefill_tokens: u64,
|
||||
}
|
||||
|
||||
/// Sub-counts of `Usage::completion_tokens`.
|
||||
|
||||
@@ -66,14 +66,48 @@ pub struct ResponsesRequest {
|
||||
pub extra: Value,
|
||||
}
|
||||
|
||||
/// `input` is either a single string or an array of typed items.
|
||||
/// `input` is either a single string or an array of items.
|
||||
/// `#[serde(untagged)]` so the wire shape `"input": "hi"` and
|
||||
/// `"input": [{...}]` both deserialize.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum ResponsesInput {
|
||||
Text(String),
|
||||
Items(Vec<ResponsesInputItem>),
|
||||
Items(Vec<ResponsesInputElement>),
|
||||
}
|
||||
|
||||
/// One element of an `input` array.
|
||||
///
|
||||
/// OpenAI's Responses API accepts three shapes here, and real clients
|
||||
/// use all of them — most notably agent-zero (via litellm), which
|
||||
/// sends the bare "easy message" form. We must tolerate every shape,
|
||||
/// because `input` is an `#[serde(untagged)]` array: a single element
|
||||
/// that matches no variant fails the *entire* request with a 422
|
||||
/// (`did not match any variant of untagged enum ResponsesInput`).
|
||||
///
|
||||
/// 1. [`Self::Typed`] — an item carrying an explicit `"type"`
|
||||
/// discriminant (`message`, `function_call`, `function_call_output`,
|
||||
/// `reasoning`).
|
||||
/// 2. [`Self::EasyMessage`] — a bare `{role, content}` with **no**
|
||||
/// `type` field. This is OpenAI's `EasyInputMessage` and what
|
||||
/// litellm emits for every turn. `content` is optional so an
|
||||
/// assistant turn carrying only tool calls (`content: null`) still
|
||||
/// parses.
|
||||
/// 3. [`Self::Other`] — anything else, captured as raw JSON and
|
||||
/// dropped during translation. This is the forward-compat escape
|
||||
/// hatch that mirrors [`ResponsesRequest::extra`] at the item
|
||||
/// level: an unmodeled item type can never again reject the whole
|
||||
/// request.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum ResponsesInputElement {
|
||||
Typed(ResponsesInputItem),
|
||||
EasyMessage {
|
||||
role: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
content: Option<ResponsesMessageContent>,
|
||||
},
|
||||
Other(Value),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
@@ -91,8 +125,11 @@ pub enum ResponsesInputItem {
|
||||
name: String,
|
||||
arguments: String,
|
||||
},
|
||||
/// User is feeding a tool result back into the model.
|
||||
FunctionCallOutput { call_id: String, output: String },
|
||||
/// User is feeding a tool result back into the model. `output`
|
||||
/// is a `Value` because OpenAI allows it to be either a plain
|
||||
/// string or an array of content parts; the translator renders
|
||||
/// either form to text rather than losing the tool result.
|
||||
FunctionCallOutput { call_id: String, output: Value },
|
||||
/// Reasoning items emitted by o-series models. Accepted but
|
||||
/// not forwarded to the model — neuron's candle path doesn't
|
||||
/// surface reasoning separately yet.
|
||||
@@ -132,6 +169,11 @@ pub enum ResponsesContentPart {
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
annotations: Vec<Value>,
|
||||
},
|
||||
/// Any content-part type we don't model (e.g. `refusal`, audio).
|
||||
/// Captured as a unit so an unknown part can't reject the whole
|
||||
/// request; dropped during translation.
|
||||
#[serde(other)]
|
||||
Unknown,
|
||||
}
|
||||
|
||||
// ── Response (non-streaming) ─────────────────────────────────────────
|
||||
@@ -277,20 +319,116 @@ mod tests {
|
||||
ResponsesInput::Items(items) => {
|
||||
assert_eq!(items.len(), 1);
|
||||
match &items[0] {
|
||||
ResponsesInputItem::Message { role, content } => {
|
||||
ResponsesInputElement::Typed(ResponsesInputItem::Message { role, content }) => {
|
||||
assert_eq!(role, "user");
|
||||
match content {
|
||||
ResponsesMessageContent::Text(t) => assert_eq!(t, "hi"),
|
||||
other => panic!("expected Text content, got {other:?}"),
|
||||
}
|
||||
}
|
||||
other => panic!("expected Message item, got {other:?}"),
|
||||
other => panic!("expected typed Message item, got {other:?}"),
|
||||
}
|
||||
}
|
||||
other => panic!("expected Items, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deserialises_bare_easy_message_without_type() {
|
||||
// The shape agent-zero (via litellm) actually sends: `input`
|
||||
// items are bare `{role, content}` with NO `type` field. This
|
||||
// is the exact payload that was returning 422.
|
||||
let raw = r#"{
|
||||
"model": "Qwen/Qwen3.6-27B",
|
||||
"store": true,
|
||||
"tools": [{"type": "function", "name": "x", "description": "d", "parameters": {}}],
|
||||
"input": [
|
||||
{"role": "system", "content": "you are helpful"},
|
||||
{"role": "assistant", "content": "{\"tool_name\":\"response\"}"},
|
||||
{"role": "user", "content": "hi"}
|
||||
]
|
||||
}"#;
|
||||
let req: ResponsesRequest = serde_json::from_str(raw).unwrap();
|
||||
let items = match req.input {
|
||||
ResponsesInput::Items(i) => i,
|
||||
other => panic!("expected Items, got {other:?}"),
|
||||
};
|
||||
assert_eq!(items.len(), 3);
|
||||
for el in &items {
|
||||
assert!(
|
||||
matches!(el, ResponsesInputElement::EasyMessage { .. }),
|
||||
"expected EasyMessage, got {el:?}"
|
||||
);
|
||||
}
|
||||
// `tools` / `store` ride through `extra`, not `input`.
|
||||
assert!(req.extra.get("tools").is_some());
|
||||
assert_eq!(req.extra.get("store"), Some(&Value::Bool(true)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tolerates_null_content_and_unknown_item_types() {
|
||||
// An assistant turn carrying only tool calls has `content: null`;
|
||||
// and a future/unmodeled item type must not 422 the request.
|
||||
let raw = r#"{
|
||||
"model": "m",
|
||||
"input": [
|
||||
{"role": "assistant", "content": null},
|
||||
{"type": "item_reference", "id": "abc"},
|
||||
{"type": "function_call_output", "call_id": "c1",
|
||||
"output": [{"type": "output_text", "text": "result"}]},
|
||||
{"role": "user", "content": "go"}
|
||||
]
|
||||
}"#;
|
||||
let req: ResponsesRequest = serde_json::from_str(raw).unwrap();
|
||||
let items = match req.input {
|
||||
ResponsesInput::Items(i) => i,
|
||||
other => panic!("expected Items, got {other:?}"),
|
||||
};
|
||||
assert_eq!(items.len(), 4);
|
||||
assert!(matches!(
|
||||
&items[0],
|
||||
ResponsesInputElement::EasyMessage { content: None, .. }
|
||||
));
|
||||
assert!(matches!(&items[1], ResponsesInputElement::Other(_)));
|
||||
assert!(matches!(
|
||||
&items[2],
|
||||
ResponsesInputElement::Typed(ResponsesInputItem::FunctionCallOutput { .. })
|
||||
));
|
||||
assert!(matches!(
|
||||
&items[3],
|
||||
ResponsesInputElement::EasyMessage { .. }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tolerates_unknown_content_part_type() {
|
||||
// A `refusal` (or any unmodeled) content part must parse, not 422.
|
||||
let raw = r#"{
|
||||
"model": "m",
|
||||
"input": [
|
||||
{"role": "assistant", "content": [
|
||||
{"type": "refusal", "refusal": "no"},
|
||||
{"type": "output_text", "text": "ok"}
|
||||
]}
|
||||
]
|
||||
}"#;
|
||||
let req: ResponsesRequest = serde_json::from_str(raw).unwrap();
|
||||
let items = match req.input {
|
||||
ResponsesInput::Items(i) => i,
|
||||
other => panic!("expected Items, got {other:?}"),
|
||||
};
|
||||
let parts = match &items[0] {
|
||||
ResponsesInputElement::EasyMessage {
|
||||
content: Some(ResponsesMessageContent::Parts(p)),
|
||||
..
|
||||
} => p,
|
||||
other => panic!("expected EasyMessage with Parts, got {other:?}"),
|
||||
};
|
||||
assert_eq!(parts.len(), 2);
|
||||
assert!(matches!(&parts[0], ResponsesContentPart::Unknown));
|
||||
assert!(matches!(&parts[1], ResponsesContentPart::OutputText { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deserialises_input_with_image() {
|
||||
let raw = r#"{
|
||||
@@ -308,10 +446,10 @@ mod tests {
|
||||
other => panic!("expected Items, got {other:?}"),
|
||||
};
|
||||
let parts = match &items[0] {
|
||||
ResponsesInputItem::Message {
|
||||
ResponsesInputElement::Typed(ResponsesInputItem::Message {
|
||||
content: ResponsesMessageContent::Parts(p),
|
||||
..
|
||||
} => p,
|
||||
}) => p,
|
||||
other => panic!("expected Parts, got {other:?}"),
|
||||
};
|
||||
assert_eq!(parts.len(), 2);
|
||||
|
||||
@@ -400,6 +400,7 @@ pub fn openai_to_anthropic(resp: ChatCompletionResponse) -> MessagesResponse {
|
||||
total_tokens: 0,
|
||||
completion_tokens_details: None,
|
||||
prompt_tokens_details: None,
|
||||
helexa_timing: None,
|
||||
});
|
||||
|
||||
MessagesResponse {
|
||||
@@ -772,6 +773,7 @@ mod stream_tests {
|
||||
total_tokens: 267,
|
||||
completion_tokens_details: None,
|
||||
prompt_tokens_details: None,
|
||||
helexa_timing: None,
|
||||
});
|
||||
t.on_chunk(&usage_chunk);
|
||||
let fin = t.finish();
|
||||
|
||||
@@ -322,7 +322,11 @@ async fn anthropic_messages(
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(guard) => Some(crate::metering::usage_sink(principal, guard)),
|
||||
Ok(guard) => Some(crate::metering::usage_sink(
|
||||
principal,
|
||||
guard,
|
||||
std::sync::Arc::clone(&fleet.served_usage),
|
||||
)),
|
||||
Err(env) => return crate::error::envelope_response(env),
|
||||
}
|
||||
}
|
||||
@@ -577,6 +581,11 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
||||
// Runtime-detected — will be OR-ed in Pass 2 from neuron data.
|
||||
tool_call: false,
|
||||
reasoning: false,
|
||||
// Flat #78 fields are derived from `limit` in the final
|
||||
// sync pass, once merging is done.
|
||||
max_model_len: None,
|
||||
max_input_tokens: None,
|
||||
max_output_tokens: None,
|
||||
},
|
||||
);
|
||||
}
|
||||
@@ -634,6 +643,9 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
||||
cost: None,
|
||||
tool_call: entry.tool_call,
|
||||
reasoning: entry.reasoning,
|
||||
max_model_len: None,
|
||||
max_input_tokens: None,
|
||||
max_output_tokens: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -690,6 +702,9 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
||||
cost: None,
|
||||
tool_call: false,
|
||||
reasoning: false,
|
||||
max_model_len: None,
|
||||
max_input_tokens: None,
|
||||
max_output_tokens: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -724,11 +739,24 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
||||
cost: target_entry.cost.clone(),
|
||||
tool_call: target_entry.tool_call,
|
||||
reasoning: target_entry.reasoning,
|
||||
max_model_len: None,
|
||||
max_input_tokens: None,
|
||||
max_output_tokens: None,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
let data: Vec<Value> = entries.values().map(|e| json!(e)).collect();
|
||||
// Final pass: derive the flat ecosystem context-window fields (#78)
|
||||
// from each entry's now-settled `limit`, so vLLM-convention clients
|
||||
// (Hermes Agent et al.) can read the window without knowing helexa's
|
||||
// `limit` schema.
|
||||
let data: Vec<Value> = entries
|
||||
.values_mut()
|
||||
.map(|e| {
|
||||
e.sync_flat_limit();
|
||||
json!(e)
|
||||
})
|
||||
.collect();
|
||||
Json(json!({
|
||||
"object": "list",
|
||||
"data": data,
|
||||
@@ -802,7 +830,11 @@ async fn proxy_with_metrics(
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(guard) => Some(crate::metering::usage_sink(principal, guard)),
|
||||
Ok(guard) => Some(crate::metering::usage_sink(
|
||||
principal,
|
||||
guard,
|
||||
std::sync::Arc::clone(&fleet.served_usage),
|
||||
)),
|
||||
Err(env) => return crate::error::envelope_response(env),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -11,6 +11,7 @@ pub mod metrics;
|
||||
pub mod poller;
|
||||
pub mod proxy;
|
||||
pub mod router;
|
||||
pub mod served_usage;
|
||||
pub mod state;
|
||||
|
||||
use anyhow::Result;
|
||||
@@ -57,6 +58,28 @@ pub async fn run(config: GatewayConfig) -> Result<()> {
|
||||
evictor::eviction_loop(evictor_fleet).await;
|
||||
});
|
||||
|
||||
// Served-usage reporter (#58): when this operator is part of the mesh,
|
||||
// periodically flush absolute per-principal served-token counters to
|
||||
// upstream for reconciliation.
|
||||
if config.upstream.enabled {
|
||||
let su_fleet = Arc::clone(&fleet);
|
||||
let url = config.upstream.url.clone();
|
||||
let bearer = config.upstream.bearer.clone();
|
||||
let interval =
|
||||
std::time::Duration::from_secs(config.upstream.served_usage_report_interval_secs);
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::time::sleep(interval).await;
|
||||
let rows = su_fleet.served_usage.snapshot();
|
||||
if let Err(e) =
|
||||
served_usage::report(&su_fleet.http_client, &url, &bearer, &rows).await
|
||||
{
|
||||
tracing::warn!(error = %e, "served-usage report failed (will retry)");
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
let app = build_app(Arc::clone(&fleet));
|
||||
|
||||
let listen_addr = config.gateway.listen.parse::<std::net::SocketAddr>()?;
|
||||
|
||||
@@ -117,9 +117,21 @@ impl Drop for ReservationGuard {
|
||||
/// Build the completion sink for an authenticated request: record spend and
|
||||
/// settle the reservation with the observed total. Dropping it unused (no
|
||||
/// usage observed) releases the reservation via the guard.
|
||||
pub fn usage_sink(principal: Principal, guard: ReservationGuard) -> UsageSink {
|
||||
pub fn usage_sink(
|
||||
principal: Principal,
|
||||
guard: ReservationGuard,
|
||||
served_usage: std::sync::Arc<crate::served_usage::ServedUsage>,
|
||||
) -> UsageSink {
|
||||
Box::new(move |prompt, completion| {
|
||||
record_spend(&principal, prompt, completion);
|
||||
// Per-principal served-usage tally for #58 reconciliation. Recorded
|
||||
// for every metered (authenticated) request; the flush task reports
|
||||
// it to upstream when the operator is part of the mesh.
|
||||
served_usage.add(
|
||||
&principal.account_id,
|
||||
&principal.key_id,
|
||||
prompt + completion,
|
||||
);
|
||||
guard.settle(prompt + completion);
|
||||
})
|
||||
}
|
||||
|
||||
105
crates/cortex-gateway/src/served_usage.rs
Normal file
105
crates/cortex-gateway/src/served_usage.rs
Normal file
@@ -0,0 +1,105 @@
|
||||
//! Served-usage ledger (#58): cortex meters, per principal and per UTC day,
|
||||
//! the tokens it has served on behalf of mesh accounts, and periodically
|
||||
//! reports **absolute** cumulative counters to helexa-upstream for
|
||||
//! reconciliation (operators are compensated for served tokens).
|
||||
//!
|
||||
//! Counters are cumulative-since-process-start for the current period;
|
||||
//! upstream upserts them monotonically (GREATEST), so re-sending the same
|
||||
//! value is idempotent and a flush that races another is harmless. (A
|
||||
//! process restart resets the in-memory counter; the monotonic upsert keeps
|
||||
//! upstream from regressing — at most it under-counts the restarted window,
|
||||
//! acceptable for beta. One cortex per operator token is assumed.)
|
||||
|
||||
use serde::Serialize;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Mutex;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
|
||||
pub struct ServedRow {
|
||||
pub account_id: String,
|
||||
pub key_id: String,
|
||||
pub period: String, // YYYY-MM-DD (UTC)
|
||||
pub served_tokens: u64,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct ServedUsage {
|
||||
inner: Mutex<HashMap<(String, String, String), u64>>,
|
||||
}
|
||||
|
||||
impl ServedUsage {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Add served tokens for a principal in today's (UTC) period.
|
||||
pub fn add(&self, account_id: &str, key_id: &str, tokens: u64) {
|
||||
if tokens == 0 {
|
||||
return;
|
||||
}
|
||||
let period = chrono::Utc::now().format("%Y-%m-%d").to_string();
|
||||
let mut m = self.inner.lock().expect("served-usage lock");
|
||||
*m.entry((account_id.to_string(), key_id.to_string(), period))
|
||||
.or_insert(0) += tokens;
|
||||
}
|
||||
|
||||
/// Absolute cumulative counters, for a flush to upstream.
|
||||
pub fn snapshot(&self) -> Vec<ServedRow> {
|
||||
let m = self.inner.lock().expect("served-usage lock");
|
||||
m.iter()
|
||||
.map(|((account_id, key_id, period), &served_tokens)| ServedRow {
|
||||
account_id: account_id.clone(),
|
||||
key_id: key_id.clone(),
|
||||
period: period.clone(),
|
||||
served_tokens,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// POST the absolute counters to upstream's `/authz/v1/served-usage`.
|
||||
pub async fn report(
|
||||
client: &reqwest::Client,
|
||||
base_url: &str,
|
||||
bearer: &str,
|
||||
rows: &[ServedRow],
|
||||
) -> Result<(), reqwest::Error> {
|
||||
if rows.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let url = format!("{}/authz/v1/served-usage", base_url.trim_end_matches('/'));
|
||||
client
|
||||
.post(url)
|
||||
.bearer_auth(bearer)
|
||||
.json(&serde_json::json!({ "rows": rows }))
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn accumulates_per_principal_and_period() {
|
||||
let su = ServedUsage::new();
|
||||
su.add("acct", "key", 10);
|
||||
su.add("acct", "key", 5);
|
||||
su.add("acct", "other", 7);
|
||||
su.add("acct", "key", 0); // no-op
|
||||
let mut rows = su.snapshot();
|
||||
rows.sort_by(|a, b| a.key_id.cmp(&b.key_id));
|
||||
assert_eq!(rows.len(), 2);
|
||||
let key_row = rows.iter().find(|r| r.key_id == "key").unwrap();
|
||||
assert_eq!(key_row.served_tokens, 15);
|
||||
assert_eq!(
|
||||
rows.iter()
|
||||
.find(|r| r.key_id == "other")
|
||||
.unwrap()
|
||||
.served_tokens,
|
||||
7
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -22,6 +22,9 @@ pub struct CortexState {
|
||||
/// Whether to reject unauthenticated requests (#49). Read by the auth
|
||||
/// middleware once it lands.
|
||||
pub require_auth: bool,
|
||||
/// Per-principal served-token tally (#58), reported to upstream for
|
||||
/// operator reconciliation by the flush task when upstream is enabled.
|
||||
pub served_usage: Arc<crate::served_usage::ServedUsage>,
|
||||
}
|
||||
|
||||
impl CortexState {
|
||||
@@ -73,6 +76,7 @@ impl CortexState {
|
||||
.expect("failed to build HTTP client"),
|
||||
entitlements,
|
||||
require_auth: config.entitlements.require_auth,
|
||||
served_usage: Arc::new(crate::served_usage::ServedUsage::new()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,8 +8,11 @@
|
||||
//! - `cost` from the catalogue profile (operator-set pricing).
|
||||
//! - `tool_call` / `reasoning` from the neuron's runtime detection (OR-ed in)
|
||||
//!
|
||||
//! Also a regression guard for the removal of `max_model_len` — the misnamed,
|
||||
//! unconsumed vLLM-ism that this contract replaces.
|
||||
//! Also asserts the flat, vLLM-convention duplicates (`max_model_len`,
|
||||
//! `max_input_tokens`, `max_output_tokens`) mirror `limit` (#78): the
|
||||
//! earlier removal of `max_model_len` as "unconsumed" was wrong — Hermes
|
||||
//! Agent (and the wider OpenAI client ecosystem) probes those flat keys
|
||||
//! and cannot see `limit.context`.
|
||||
|
||||
use cortex_core::config::{
|
||||
EvictionSettings, EvictionStrategy, GatewayConfig, GatewaySettings, NeuronEndpoint,
|
||||
@@ -85,6 +88,21 @@ capabilities = ["text"]
|
||||
}),
|
||||
},
|
||||
);
|
||||
// A model with no derivable limit: the flat #78 fields must be
|
||||
// OMITTED (absent-vs-zero is load-bearing), never 0 or a guess.
|
||||
node.models.insert(
|
||||
"no-limit-model".into(),
|
||||
ModelEntry {
|
||||
id: "no-limit-model".into(),
|
||||
status: ModelStatus::Loaded,
|
||||
last_accessed: None,
|
||||
vram_estimate_mb: None,
|
||||
capabilities: vec!["text".into()],
|
||||
tool_call: false,
|
||||
reasoning: false,
|
||||
limit: None,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
let app = cortex_gateway::build_app(Arc::clone(&fleet));
|
||||
@@ -123,11 +141,26 @@ capabilities = ["text"]
|
||||
assert_eq!(entry["tool_call"], true);
|
||||
assert_eq!(entry["reasoning"], true);
|
||||
|
||||
// Regression guard: the removed, unconsumed vLLM-ism must not reappear.
|
||||
assert!(
|
||||
entry.get("max_model_len").is_none(),
|
||||
"max_model_len was removed; /v1/models must not advertise it"
|
||||
);
|
||||
// Flat ecosystem duplicates (#78) mirror the advertised `limit` so
|
||||
// vLLM-convention probes (Hermes Agent) auto-detect the window.
|
||||
assert_eq!(entry["max_model_len"], 49152);
|
||||
assert_eq!(entry["max_input_tokens"], 40960);
|
||||
assert_eq!(entry["max_output_tokens"], 8192);
|
||||
|
||||
// No limit → flat fields omitted entirely, never 0 or a guess.
|
||||
let unknown = body["data"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.find(|m| m["id"] == "no-limit-model")
|
||||
.expect("no-limit-model present in /v1/models");
|
||||
assert!(unknown.get("limit").is_none());
|
||||
for key in ["max_model_len", "max_input_tokens", "max_output_tokens"] {
|
||||
assert!(
|
||||
unknown.get(key).is_none(),
|
||||
"{key} must be omitted when the window is unknown"
|
||||
);
|
||||
}
|
||||
|
||||
let _ = std::fs::remove_file(&cat_path);
|
||||
}
|
||||
|
||||
@@ -60,6 +60,7 @@ fn chain(local: LocalEntitlementProvider, url: &str) -> ChainedEntitlementProvid
|
||||
url: url.to_string(),
|
||||
bearer: "client-secret".into(),
|
||||
timeout_secs: 5,
|
||||
served_usage_report_interval_secs: 60,
|
||||
});
|
||||
ChainedEntitlementProvider::new(local, upstream)
|
||||
}
|
||||
|
||||
@@ -35,6 +35,9 @@ pub fn api_routes(state: ApiState) -> Router {
|
||||
.route("/api/health", get(health))
|
||||
.route("/api/dimensions", get(dimensions))
|
||||
.route("/api/summary", get(summary))
|
||||
.route("/api/scaling", get(scaling))
|
||||
.route("/api/swap", get(swap))
|
||||
.route("/api/capability", get(capability))
|
||||
.route("/api/series", get(series))
|
||||
.route("/api/runs", get(runs))
|
||||
.layer(CorsLayer::permissive())
|
||||
@@ -73,6 +76,30 @@ async fn summary(
|
||||
store.summary().map(Json).map_err(err500)
|
||||
}
|
||||
|
||||
/// Context-length scaling curves per (target, model) — prefill & decode
|
||||
/// tok/s vs context, with decode-flatness (#88).
|
||||
async fn scaling(
|
||||
State(s): State<ApiState>,
|
||||
) -> Result<Json<Vec<crate::store::ScalingCurve>>, ApiError> {
|
||||
let store = s.lock().await;
|
||||
store.scaling().map(Json).map_err(err500)
|
||||
}
|
||||
|
||||
/// Cold-load / model-swap costs per (target, model) — reload latency + cold
|
||||
/// first-request (#90).
|
||||
async fn swap(State(s): State<ApiState>) -> Result<Json<Vec<crate::store::SwapCost>>, ApiError> {
|
||||
let store = s.lock().await;
|
||||
store.swap_costs().map(Json).map_err(err500)
|
||||
}
|
||||
|
||||
/// Capability-probe runs — stored artifacts + quality scores (#91).
|
||||
async fn capability(
|
||||
State(s): State<ApiState>,
|
||||
) -> Result<Json<Vec<crate::store::CapabilityRun>>, ApiError> {
|
||||
let store = s.lock().await;
|
||||
store.capability_runs(false).map(Json).map_err(err500)
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct SeriesQuery {
|
||||
/// Optional — when omitted the store resolves the host serving this model.
|
||||
|
||||
@@ -3,10 +3,10 @@
|
||||
//! `openai` targets use the OpenAI-compatible surface (preliminary).
|
||||
|
||||
use crate::config::{TargetConfig, TargetKind};
|
||||
use anyhow::{Context, Result};
|
||||
use anyhow::{Context, Result, anyhow};
|
||||
use cortex_core::build_info::BuildInfo;
|
||||
use cortex_core::discovery::DiscoveryResponse;
|
||||
use cortex_core::harness::ModelInfo;
|
||||
use cortex_core::discovery::{DiscoveryResponse, HealthResponse};
|
||||
use cortex_core::harness::{ModelInfo, ModelSpec};
|
||||
use cortex_core::openai::ModelsResponse;
|
||||
use std::time::Duration;
|
||||
|
||||
@@ -94,6 +94,84 @@ impl TargetClient {
|
||||
Ok(Some(disco))
|
||||
}
|
||||
|
||||
/// Runtime device health (neuron only): per-GPU VRAM used/free,
|
||||
/// utilization, and temperature from `GET /health`. Bench samples this
|
||||
/// around each measured run to record VRAM high-water + GPU telemetry
|
||||
/// (#87). Returns `Ok(None)` for non-neuron targets; a soft `Ok(None)`
|
||||
/// (not an error) on transport failure so a flaky `/health` never fails
|
||||
/// a measurement.
|
||||
pub async fn fetch_health(&self, target: &TargetConfig) -> Result<Option<HealthResponse>> {
|
||||
if target.kind != TargetKind::Neuron {
|
||||
return Ok(None);
|
||||
}
|
||||
let base = target.endpoint.trim_end_matches('/');
|
||||
let health = self
|
||||
.http
|
||||
.get(format!("{base}/health"))
|
||||
.timeout(META_TIMEOUT)
|
||||
.send()
|
||||
.await
|
||||
.context("GET /health")?
|
||||
.error_for_status()
|
||||
.context("GET /health status")?
|
||||
.json::<HealthResponse>()
|
||||
.await
|
||||
.context("decoding /health")?;
|
||||
Ok(Some(health))
|
||||
}
|
||||
|
||||
/// Unload a model (neuron only): `POST /models/unload {model_id}`.
|
||||
/// Used by the deliberate swap-cost measurement (#90), never the sweep.
|
||||
pub async fn unload_model(&self, target: &TargetConfig, model_id: &str) -> Result<()> {
|
||||
let base = target.endpoint.trim_end_matches('/');
|
||||
self.http
|
||||
.post(format!("{base}/models/unload"))
|
||||
.json(&serde_json::json!({ "model_id": model_id }))
|
||||
.send()
|
||||
.await
|
||||
.context("POST /models/unload")?
|
||||
.error_for_status()
|
||||
.context("POST /models/unload status")?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Load a model from a spec (neuron only): `POST /models/load`. neuron
|
||||
/// returns synchronously once loaded, so the call duration is the reload
|
||||
/// cost the swap-cost measurement records (#90).
|
||||
pub async fn load_model(&self, target: &TargetConfig, spec: &ModelSpec) -> Result<()> {
|
||||
let base = target.endpoint.trim_end_matches('/');
|
||||
self.http
|
||||
.post(format!("{base}/models/load"))
|
||||
.json(spec)
|
||||
// A cold load can take tens of seconds; use the full request
|
||||
// timeout rather than the short metadata one.
|
||||
.send()
|
||||
.await
|
||||
.context("POST /models/load")?
|
||||
.error_for_status()
|
||||
.context("POST /models/load status")?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Reconstruct a reload [`ModelSpec`] from a model's `/models` entry.
|
||||
/// Tensor-parallel is inferred from the device count; `quant` is left
|
||||
/// `None` for neuron to resolve from the catalogue / its prior load.
|
||||
pub fn spec_from_info(info: &ModelInfo) -> Result<ModelSpec> {
|
||||
if info.devices.is_empty() {
|
||||
return Err(anyhow!(
|
||||
"model '{}' reports no devices; cannot reconstruct a load spec",
|
||||
info.id
|
||||
));
|
||||
}
|
||||
Ok(ModelSpec {
|
||||
model_id: info.id.clone(),
|
||||
harness: info.harness.clone(),
|
||||
quant: None,
|
||||
tensor_parallel: (info.devices.len() > 1).then_some(info.devices.len() as u32),
|
||||
devices: Some(info.devices.clone()),
|
||||
})
|
||||
}
|
||||
|
||||
/// Warm models — those ready to serve without a cold load.
|
||||
///
|
||||
/// Neuron: `GET /models` filtered to `status == "loaded"` (skips
|
||||
|
||||
@@ -104,6 +104,37 @@ pub struct ScenarioConfig {
|
||||
/// Max generated tokens per request.
|
||||
#[serde(default = "default_max_tokens")]
|
||||
pub max_tokens: u64,
|
||||
/// Concurrency levels (#89): one `concurrency:<n>` scenario per entry,
|
||||
/// each firing N simultaneous streams. Defaults to empty (opt-in) — a
|
||||
/// burst puts real load on a serving fleet, so operators enable it
|
||||
/// deliberately, e.g. `concurrency_levels = [2, 4, 8]`.
|
||||
#[serde(default)]
|
||||
pub concurrency_levels: Vec<u32>,
|
||||
/// Approximate prompt size (tokens) used by the concurrency scenarios.
|
||||
#[serde(default = "default_concurrency_prompt_tokens")]
|
||||
pub concurrency_prompt_tokens: u32,
|
||||
/// Capability probes (#91): one `capability:<name>` scenario per entry,
|
||||
/// each running a fixed prompt and storing the full output artifact for
|
||||
/// quality scoring (manual now, LLM-judge later). Defaults to empty
|
||||
/// (opt-in) — these generate long outputs and exist to compare reasoning
|
||||
/// quality across models, not to run on every sweep by default.
|
||||
#[serde(default)]
|
||||
pub capability_probes: Vec<CapabilityProbe>,
|
||||
}
|
||||
|
||||
/// One capability probe: a named prompt whose output is stored and scored
|
||||
/// for quality (#91). The probe is deterministic (temperature 0) so the
|
||||
/// same model+build produces a stable artifact to score.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct CapabilityProbe {
|
||||
/// Stable id fragment — the scenario becomes `capability:<name>`.
|
||||
pub name: String,
|
||||
/// The prompt sent verbatim as the user message.
|
||||
pub prompt: String,
|
||||
/// Generation budget for the probe (planning answers are long; the
|
||||
/// default 256 is too small). Defaults to 2048.
|
||||
#[serde(default = "default_capability_max_tokens")]
|
||||
pub max_tokens: u64,
|
||||
}
|
||||
|
||||
impl Default for ScenarioConfig {
|
||||
@@ -111,6 +142,9 @@ impl Default for ScenarioConfig {
|
||||
ScenarioConfig {
|
||||
prompt_sizes: default_prompt_sizes(),
|
||||
max_tokens: default_max_tokens(),
|
||||
concurrency_levels: Vec::new(),
|
||||
concurrency_prompt_tokens: default_concurrency_prompt_tokens(),
|
||||
capability_probes: Vec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -187,6 +221,12 @@ fn default_prompt_sizes() -> Vec<u32> {
|
||||
fn default_max_tokens() -> u64 {
|
||||
256
|
||||
}
|
||||
fn default_concurrency_prompt_tokens() -> u32 {
|
||||
512
|
||||
}
|
||||
fn default_capability_max_tokens() -> u64 {
|
||||
2048
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
// Jail's closure must return figment::Result; the large-Err type is
|
||||
|
||||
@@ -43,6 +43,33 @@ enum Command {
|
||||
#[arg(short, long, default_value = "helexa-bench.toml")]
|
||||
config: String,
|
||||
},
|
||||
/// Measure cold-load / model-swap cost (#90): for each neuron target's
|
||||
/// warm models, unload → time reload → time a cold first request, recorded
|
||||
/// under scenario "swap". DELIBERATE — takes each model offline for its
|
||||
/// reload, so run it in a maintenance window, not against live traffic.
|
||||
SwapCost {
|
||||
#[arg(short, long, default_value = "helexa-bench.toml")]
|
||||
config: String,
|
||||
},
|
||||
/// Attach a quality score to a capability-probe run (#91). Find run ids
|
||||
/// with `report --capability`. `--scorer` records who scored it
|
||||
/// (defaults to "manual"); a future LLM-judge would set e.g. "llm:…".
|
||||
Score {
|
||||
#[arg(short, long, default_value = "helexa-bench.toml")]
|
||||
config: String,
|
||||
/// Override the SQLite path (skips reading the config file).
|
||||
#[arg(long)]
|
||||
db: Option<String>,
|
||||
/// The run id to score.
|
||||
#[arg(long)]
|
||||
id: i64,
|
||||
/// The quality score to attach (scale is the operator's rubric).
|
||||
#[arg(long)]
|
||||
score: f64,
|
||||
/// Who/what produced the score.
|
||||
#[arg(long, default_value = "manual")]
|
||||
scorer: String,
|
||||
},
|
||||
/// Render recorded results. Uses `--db` if given, else the db_path
|
||||
/// from `--config`.
|
||||
Report {
|
||||
@@ -54,6 +81,19 @@ enum Command {
|
||||
/// Output format.
|
||||
#[arg(long, default_value = "md")]
|
||||
format: Format,
|
||||
/// Render the context-length scaling view (prefill & decode tok/s
|
||||
/// vs context per model, with decode-flatness) instead of the flat
|
||||
/// results table (#88).
|
||||
#[arg(long)]
|
||||
scaling: bool,
|
||||
/// Render the cold-load / model-swap cost view (#90) instead of the
|
||||
/// flat results table.
|
||||
#[arg(long)]
|
||||
swap: bool,
|
||||
/// Render the capability-probe view (#91): stored artifacts + quality
|
||||
/// scores, with per-model median.
|
||||
#[arg(long)]
|
||||
capability: bool,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -122,16 +162,79 @@ async fn run(cli: Cli) -> Result<()> {
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Command::Report { config, db, format } => {
|
||||
Command::SwapCost { config } => {
|
||||
let cfg = load_config(&config)?;
|
||||
require_targets(&cfg)?;
|
||||
let sweeper = Sweeper::new(cfg)?;
|
||||
tracing::warn!(
|
||||
"swap-cost: cycling each warm model (unload → reload → cold request); models go offline during reload"
|
||||
);
|
||||
let summary = sweeper.swap_cost_once().await?;
|
||||
tracing::info!(
|
||||
measured = summary.measured,
|
||||
failed = summary.failed,
|
||||
unreachable = summary.targets_unreachable,
|
||||
"swap-cost measurement complete"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Command::Score {
|
||||
config,
|
||||
db,
|
||||
id,
|
||||
score,
|
||||
scorer,
|
||||
} => {
|
||||
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)?,
|
||||
match store.set_score(id, score, &scorer)? {
|
||||
0 => anyhow::bail!("no run with id {id}"),
|
||||
_ => {
|
||||
println!("scored run {id}: {score} ({scorer})");
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
Command::Report {
|
||||
config,
|
||||
db,
|
||||
format,
|
||||
scaling,
|
||||
swap,
|
||||
capability,
|
||||
} => {
|
||||
let db_path = match db {
|
||||
Some(p) => p,
|
||||
None => load_config(&config)?.bench.db_path,
|
||||
};
|
||||
let store = Store::open(&db_path)?;
|
||||
let rendered = if capability {
|
||||
let runs = store.capability_runs(false)?;
|
||||
match format {
|
||||
Format::Md => report::render_capability_markdown(&runs),
|
||||
Format::Json => report::render_capability_json(&runs)?,
|
||||
}
|
||||
} else if swap {
|
||||
let costs = store.swap_costs()?;
|
||||
match format {
|
||||
Format::Md => report::render_swap_markdown(&costs),
|
||||
Format::Json => report::render_swap_json(&costs)?,
|
||||
}
|
||||
} else if scaling {
|
||||
let curves = store.scaling()?;
|
||||
match format {
|
||||
Format::Md => report::render_scaling_markdown(&curves),
|
||||
Format::Json => report::render_scaling_json(&curves)?,
|
||||
}
|
||||
} else {
|
||||
let rows = store.report_rows()?;
|
||||
match format {
|
||||
Format::Md => report::render_markdown(&rows),
|
||||
Format::Json => report::render_json(&rows)?,
|
||||
}
|
||||
};
|
||||
println!("{rendered}");
|
||||
Ok(())
|
||||
|
||||
@@ -3,28 +3,36 @@
|
||||
//! 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 crate::store::{CapabilityRun, ReportRow, ScalingCurve, SwapCost};
|
||||
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",
|
||||
"| engine | model | prompt tok | prefill tok/s | TTFT (s) | TTFT p95 | \
|
||||
decode tok/s | total (s) | total p95 | VRAM (GB) | conc | queue ms | rej | build | n |\n",
|
||||
);
|
||||
out.push_str("|---|---|---:|---:|---:|---:|---|---:|\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",
|
||||
"| {} | {} | {} | {} | {} | {} | {} | {} | {} | {} | {} | {} | {} | `{}` | {} |\n",
|
||||
r.target_name,
|
||||
r.model_id,
|
||||
ptok,
|
||||
fmt_opt(r.prefill_tps_median, 1),
|
||||
fmt_opt(r.ttft_s_median, 3),
|
||||
fmt_opt(r.ttft_s_p95, 3),
|
||||
fmt_opt(r.decode_tps_median, 1),
|
||||
fmt_opt(r.total_s_median, 3),
|
||||
fmt_opt(r.total_s_p95, 3),
|
||||
fmt_vram(r.vram_used_mb_median, r.vram_total_mb),
|
||||
fmt_u64(r.concurrency),
|
||||
fmt_opt(r.queue_wait_ms_median, 0),
|
||||
fmt_opt(r.rejected_median, 0),
|
||||
r.git_sha,
|
||||
r.samples,
|
||||
));
|
||||
@@ -43,8 +51,23 @@ pub fn render_json(rows: &[ReportRow]) -> Result<String> {
|
||||
"prompt_size_approx": r.prompt_size_approx,
|
||||
"prompt_tokens": r.prompt_tokens,
|
||||
"ttft_s_median": r.ttft_s_median,
|
||||
"ttft_s_p95": r.ttft_s_p95,
|
||||
"ttft_s_p99": r.ttft_s_p99,
|
||||
"decode_tps_median": r.decode_tps_median,
|
||||
"total_s_median": r.total_s_median,
|
||||
"total_s_p95": r.total_s_p95,
|
||||
"total_s_p99": r.total_s_p99,
|
||||
"prefill_ms_median": r.prefill_ms_median,
|
||||
"decode_ms_median": r.decode_ms_median,
|
||||
"prefill_tps_median": r.prefill_tps_median,
|
||||
"vram_used_mb_median": r.vram_used_mb_median,
|
||||
"vram_total_mb": r.vram_total_mb,
|
||||
"gpu_util_pct_median": r.gpu_util_pct_median,
|
||||
"gpu_temp_c_median": r.gpu_temp_c_median,
|
||||
"concurrency": r.concurrency,
|
||||
"ttft_p95_load_s": r.ttft_p95_load_s,
|
||||
"queue_wait_ms_median": r.queue_wait_ms_median,
|
||||
"rejected_median": r.rejected_median,
|
||||
"git_sha": r.git_sha,
|
||||
"samples": r.samples,
|
||||
"gpu": r.gpu,
|
||||
@@ -54,6 +77,179 @@ pub fn render_json(rows: &[ReportRow]) -> Result<String> {
|
||||
Ok(serde_json::to_string_pretty(&arr)?)
|
||||
}
|
||||
|
||||
/// Context-length scaling view (#88): one block per (target, model) with
|
||||
/// prefill & decode tok/s vs context, then the decode-flatness verdict.
|
||||
pub fn render_scaling_markdown(curves: &[ScalingCurve]) -> String {
|
||||
let mut out = String::new();
|
||||
for c in curves {
|
||||
let gpu = c.gpu.as_deref().unwrap_or("");
|
||||
out.push_str(&format!(
|
||||
"### {} · {} (`{}`{})\n\n",
|
||||
c.target_name,
|
||||
c.model_id,
|
||||
c.git_sha,
|
||||
if gpu.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
format!(", {gpu}")
|
||||
},
|
||||
));
|
||||
out.push_str("| ctx tok | prefill tok/s | decode tok/s | n |\n");
|
||||
out.push_str("|---:|---:|---:|---:|\n");
|
||||
for p in &c.points {
|
||||
let ctx = p
|
||||
.prompt_tokens
|
||||
.map(|t| t.to_string())
|
||||
.unwrap_or_else(|| format!("~{}", p.prompt_size));
|
||||
out.push_str(&format!(
|
||||
"| {} | {} | {} | {} |\n",
|
||||
ctx,
|
||||
fmt_opt(p.prefill_tps, 1),
|
||||
fmt_opt(p.decode_tps, 1),
|
||||
p.samples,
|
||||
));
|
||||
}
|
||||
match c.decode_flatness {
|
||||
Some(f) => out.push_str(&format!(
|
||||
"\ndecode flatness: {f:.2} — decode tok/s {} across the context range \
|
||||
({})\n\n",
|
||||
if f >= 0.9 {
|
||||
"holds"
|
||||
} else if f >= 0.7 {
|
||||
"softens"
|
||||
} else {
|
||||
"drops sharply"
|
||||
},
|
||||
if f >= 0.9 {
|
||||
"Gated-DeltaNet O(1) decode confirmed"
|
||||
} else {
|
||||
"investigate where it breaks"
|
||||
},
|
||||
)),
|
||||
None => out.push_str("\ndecode flatness: — (need ≥2 context points)\n\n"),
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
pub fn render_scaling_json(curves: &[ScalingCurve]) -> Result<String> {
|
||||
Ok(serde_json::to_string_pretty(curves)?)
|
||||
}
|
||||
|
||||
/// Cold-load / model-swap cost view (#90): reload latency + cold
|
||||
/// first-request per model.
|
||||
pub fn render_swap_markdown(costs: &[SwapCost]) -> String {
|
||||
let mut out = String::new();
|
||||
out.push_str(
|
||||
"| engine | model | unload (s) | reload (s) | cold TTFT (s) | cold total (s) | build | n |\n",
|
||||
);
|
||||
out.push_str("|---|---|---:|---:|---:|---:|---|---:|\n");
|
||||
for c in costs {
|
||||
out.push_str(&format!(
|
||||
"| {} | {} | {} | {} | {} | {} | `{}` | {} |\n",
|
||||
c.target_name,
|
||||
c.model_id,
|
||||
fmt_ms_as_s(c.unload_ms_median),
|
||||
fmt_ms_as_s(c.load_ms_median),
|
||||
fmt_opt(c.cold_ttft_s_median, 3),
|
||||
fmt_opt(c.cold_total_s_median, 3),
|
||||
c.git_sha,
|
||||
c.samples,
|
||||
));
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
pub fn render_swap_json(costs: &[SwapCost]) -> Result<String> {
|
||||
Ok(serde_json::to_string_pretty(costs)?)
|
||||
}
|
||||
|
||||
/// Capability-probe view (#91): per (model, probe) the median quality score
|
||||
/// (the A/B number), then each run's id, score, and an artifact snippet so
|
||||
/// unscored runs can be located and scored (`helexa-bench score --id …`).
|
||||
pub fn render_capability_markdown(runs: &[CapabilityRun]) -> String {
|
||||
use std::collections::BTreeMap;
|
||||
let mut groups: BTreeMap<(String, String, String), Vec<&CapabilityRun>> = BTreeMap::new();
|
||||
for r in runs {
|
||||
groups
|
||||
.entry((
|
||||
r.target_name.clone(),
|
||||
r.model_id.clone(),
|
||||
r.scenario_id.clone(),
|
||||
))
|
||||
.or_default()
|
||||
.push(r);
|
||||
}
|
||||
let mut out = String::new();
|
||||
for ((target, model, scenario), rs) in groups {
|
||||
let scores: Vec<f64> = rs.iter().filter_map(|r| r.quality_score).collect();
|
||||
let median = median_slice(&scores);
|
||||
out.push_str(&format!(
|
||||
"### {target} · {model} · {scenario} — median score {} ({}/{} scored)\n\n",
|
||||
median
|
||||
.map(|m| format!("{m:.1}"))
|
||||
.unwrap_or_else(|| "—".into()),
|
||||
scores.len(),
|
||||
rs.len(),
|
||||
));
|
||||
out.push_str("| run | score | scorer | build | artifact (snippet) |\n");
|
||||
out.push_str("|---:|---:|---|---|---|\n");
|
||||
for r in rs {
|
||||
out.push_str(&format!(
|
||||
"| {} | {} | {} | `{}` | {} |\n",
|
||||
r.id,
|
||||
r.quality_score
|
||||
.map(|s| format!("{s:.1}"))
|
||||
.unwrap_or_else(|| "—".into()),
|
||||
r.scorer.as_deref().unwrap_or("—"),
|
||||
r.git_sha,
|
||||
snippet(r.artifact.as_deref()),
|
||||
));
|
||||
}
|
||||
out.push('\n');
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
pub fn render_capability_json(runs: &[CapabilityRun]) -> Result<String> {
|
||||
Ok(serde_json::to_string_pretty(runs)?)
|
||||
}
|
||||
|
||||
/// First ~80 chars of an artifact on one line, for the table cell.
|
||||
fn snippet(artifact: Option<&str>) -> String {
|
||||
match artifact {
|
||||
Some(a) => {
|
||||
let one_line: String = a.split_whitespace().collect::<Vec<_>>().join(" ");
|
||||
let trimmed: String = one_line.chars().take(80).collect();
|
||||
if one_line.chars().count() > 80 {
|
||||
format!("{trimmed}…")
|
||||
} else {
|
||||
trimmed
|
||||
}
|
||||
}
|
||||
None => "—".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn median_slice(v: &[f64]) -> Option<f64> {
|
||||
if v.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let mut s = v.to_vec();
|
||||
s.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||
let lo = (s.len() - 1) / 2;
|
||||
let hi = s.len() / 2;
|
||||
Some((s[lo] + s[hi]) / 2.0)
|
||||
}
|
||||
|
||||
/// Milliseconds rendered as seconds (reload costs read naturally in s).
|
||||
fn fmt_ms_as_s(ms: Option<f64>) -> String {
|
||||
match ms {
|
||||
Some(x) => format!("{:.2}", x / 1000.0),
|
||||
None => "—".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn fmt_opt(v: Option<f64>, places: usize) -> String {
|
||||
match v {
|
||||
Some(x) => format!("{x:.places$}"),
|
||||
@@ -61,9 +257,113 @@ fn fmt_opt(v: Option<f64>, places: usize) -> String {
|
||||
}
|
||||
}
|
||||
|
||||
/// Integer cell (concurrency width); `—` when unset (non-concurrency rows).
|
||||
fn fmt_u64(v: Option<u64>) -> String {
|
||||
match v {
|
||||
Some(x) => x.to_string(),
|
||||
None => "—".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// `used/total` in GB (e.g. `42.0/64.0`) — the headroom-at-a-glance cell.
|
||||
/// `used` alone if the node total is unknown; `—` if no telemetry.
|
||||
fn fmt_vram(used_mb: Option<f64>, total_mb: Option<u64>) -> String {
|
||||
match (used_mb, total_mb) {
|
||||
(Some(u), Some(t)) => format!("{:.1}/{:.1}", u / 1024.0, t as f64 / 1024.0),
|
||||
(Some(u), None) => format!("{:.1}", u / 1024.0),
|
||||
_ => "—".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::store::{ScalingCurve, ScalingPoint};
|
||||
|
||||
#[test]
|
||||
fn capability_markdown_groups_with_median_and_snippet() {
|
||||
let runs = vec![
|
||||
CapabilityRun {
|
||||
id: 7,
|
||||
ts: "t".into(),
|
||||
target_name: "beast".into(),
|
||||
model_id: "m".into(),
|
||||
scenario_id: "capability:plan".into(),
|
||||
git_sha: "abc".into(),
|
||||
quality_score: Some(8.0),
|
||||
scorer: Some("manual".into()),
|
||||
artifact: Some("A detailed plan with trade-offs and sequencing.".into()),
|
||||
},
|
||||
CapabilityRun {
|
||||
id: 8,
|
||||
ts: "t".into(),
|
||||
target_name: "beast".into(),
|
||||
model_id: "m".into(),
|
||||
scenario_id: "capability:plan".into(),
|
||||
git_sha: "abc".into(),
|
||||
quality_score: Some(6.0),
|
||||
scorer: Some("manual".into()),
|
||||
artifact: Some("Shorter plan.".into()),
|
||||
},
|
||||
];
|
||||
let md = render_capability_markdown(&runs);
|
||||
assert!(md.contains("capability:plan"));
|
||||
assert!(md.contains("median score 7.0")); // median(8,6)
|
||||
assert!(md.contains("trade-offs"));
|
||||
assert!(md.contains("| 7 |"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn swap_markdown_renders_reload_and_cold_costs() {
|
||||
let costs = vec![SwapCost {
|
||||
target_name: "beast".into(),
|
||||
model_id: "Qwen/Qwen3.6-27B".into(),
|
||||
git_sha: "abc1234".into(),
|
||||
gpu: Some("2× RTX 5090".into()),
|
||||
unload_ms_median: Some(320.0),
|
||||
load_ms_median: Some(25000.0),
|
||||
cold_ttft_s_median: Some(2.5),
|
||||
cold_total_s_median: Some(5.0),
|
||||
samples: 3,
|
||||
}];
|
||||
let md = render_swap_markdown(&costs);
|
||||
assert!(md.contains("reload (s)"));
|
||||
assert!(md.contains("beast"));
|
||||
assert!(md.contains("25.00")); // 25000 ms → 25.00 s
|
||||
assert!(md.contains("2.500"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scaling_markdown_renders_curve_and_flatness() {
|
||||
let curves = vec![ScalingCurve {
|
||||
target_name: "beast".into(),
|
||||
model_id: "Qwen/Qwen3.6-27B".into(),
|
||||
git_sha: "abc1234".into(),
|
||||
gpu: Some("2× RTX 5090".into()),
|
||||
points: vec![
|
||||
ScalingPoint {
|
||||
prompt_size: 128,
|
||||
prompt_tokens: Some(130),
|
||||
prefill_tps: Some(900.0),
|
||||
decode_tps: Some(50.0),
|
||||
samples: 5,
|
||||
},
|
||||
ScalingPoint {
|
||||
prompt_size: 4096,
|
||||
prompt_tokens: Some(4100),
|
||||
prefill_tps: Some(2800.0),
|
||||
decode_tps: Some(48.0),
|
||||
samples: 5,
|
||||
},
|
||||
],
|
||||
decode_flatness: Some(0.96),
|
||||
}];
|
||||
let md = render_scaling_markdown(&curves);
|
||||
assert!(md.contains("### beast · Qwen/Qwen3.6-27B"));
|
||||
assert!(md.contains("ctx tok"));
|
||||
assert!(md.contains("decode flatness: 0.96"));
|
||||
assert!(md.contains("holds"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn markdown_has_header_and_row() {
|
||||
@@ -77,14 +377,36 @@ mod tests {
|
||||
ttft_s_median: Some(0.123),
|
||||
decode_tps_median: Some(45.6),
|
||||
total_s_median: Some(1.234),
|
||||
ttft_s_p95: Some(0.222),
|
||||
ttft_s_p99: Some(0.250),
|
||||
total_s_p95: Some(1.5),
|
||||
total_s_p99: Some(1.6),
|
||||
prefill_ms_median: Some(120.0),
|
||||
decode_ms_median: Some(1100.0),
|
||||
prefill_tps_median: Some(1066.7),
|
||||
vram_used_mb_median: Some(43008.0),
|
||||
vram_total_mb: Some(65536),
|
||||
gpu_util_pct_median: Some(89.0),
|
||||
gpu_temp_c_median: Some(64.0),
|
||||
concurrency: None,
|
||||
ttft_p95_load_s: None,
|
||||
queue_wait_ms_median: None,
|
||||
rejected_median: None,
|
||||
samples: 5,
|
||||
gpu: Some("2× RTX 5090".into()),
|
||||
}];
|
||||
let md = render_markdown(&rows);
|
||||
assert!(md.contains("| engine |"));
|
||||
assert!(md.contains("prefill tok/s"));
|
||||
assert!(md.contains("VRAM (GB)"));
|
||||
assert!(md.contains("conc"));
|
||||
assert!(md.contains("beast"));
|
||||
assert!(md.contains("`30d50d6`"));
|
||||
assert!(md.contains("0.123"));
|
||||
// p95 column rendered.
|
||||
assert!(md.contains("0.222"));
|
||||
// VRAM used/total in GB (43008/65536 MiB → 42.0/64.0).
|
||||
assert!(md.contains("42.0/64.0"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -99,6 +421,21 @@ mod tests {
|
||||
ttft_s_median: Some(0.1),
|
||||
decode_tps_median: None,
|
||||
total_s_median: Some(0.5),
|
||||
ttft_s_p95: Some(0.1),
|
||||
ttft_s_p99: Some(0.1),
|
||||
total_s_p95: Some(0.5),
|
||||
total_s_p99: Some(0.5),
|
||||
prefill_ms_median: None,
|
||||
decode_ms_median: None,
|
||||
prefill_tps_median: None,
|
||||
vram_used_mb_median: None,
|
||||
vram_total_mb: None,
|
||||
gpu_util_pct_median: None,
|
||||
gpu_temp_c_median: None,
|
||||
concurrency: None,
|
||||
ttft_p95_load_s: None,
|
||||
queue_wait_ms_median: None,
|
||||
rejected_median: None,
|
||||
samples: 1,
|
||||
gpu: None,
|
||||
}];
|
||||
|
||||
@@ -62,6 +62,41 @@ pub struct ScenarioMetrics {
|
||||
pub prompt_tokens: Option<u64>,
|
||||
/// Completion tokens: from `usage` when present, else content-chunk count.
|
||||
pub completion_tokens: u64,
|
||||
/// Server-measured prefill duration (ms), from the `usage.helexa_timing`
|
||||
/// extension (#85). `None` when the server didn't emit it (external
|
||||
/// engines, non-instrumented paths). The honest prefill-phase number,
|
||||
/// distinct from client-observed `ttft_s` which also includes request
|
||||
/// setup + first-byte network latency.
|
||||
pub prefill_ms: Option<u64>,
|
||||
/// Server-measured decode duration (ms), from `usage.helexa_timing`.
|
||||
pub decode_ms: Option<u64>,
|
||||
/// Tokens submitted to prefill — the denominator for prefill tok/s.
|
||||
pub prefill_tokens: Option<u64>,
|
||||
// ── Concurrency / agentic-load fields (#89) ──────────────────────────
|
||||
// Set only by the concurrency scenario, which fans out N simultaneous
|
||||
// streams to characterize the real a0/hermes/opencode workload that
|
||||
// batch-1 single-request measurement can't see. `None` for single
|
||||
// requests. For a concurrency burst, the inherited fields carry the
|
||||
// aggregate: `ttft_s` = median TTFT across streams, `decode_tps` = node
|
||||
// throughput (total tokens / burst window), `total_s` = burst wall-clock,
|
||||
// `completion_tokens` = total across streams.
|
||||
/// Number of simultaneous streams in the burst (the cell dimension).
|
||||
pub concurrency: Option<u32>,
|
||||
/// p95 of per-stream TTFT within the burst — the tail under simultaneous
|
||||
/// load, where batch-1 serialization actually hurts.
|
||||
pub ttft_p95_s: Option<f64>,
|
||||
/// Median per-stream admission queue-wait (ms), approximated as
|
||||
/// `ttft − prefill_ms` (#85): on a batch-1 server, later streams wait for
|
||||
/// earlier ones, so TTFT inflates while server prefill stays constant —
|
||||
/// the gap is the wait. `None` if streams didn't report `helexa_timing`.
|
||||
pub queue_wait_ms_median: Option<f64>,
|
||||
/// Streams shed by admission control (HTTP 429/503) during the burst —
|
||||
/// honest backpressure, not silent failures.
|
||||
pub rejected: Option<u32>,
|
||||
/// Full generated text, captured only by the capability probe (#91) so
|
||||
/// the output can be quality-scored later (manual or LLM-judge). `None`
|
||||
/// for latency/throughput scenarios, which discard the text.
|
||||
pub artifact: Option<String>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -85,10 +120,13 @@ pub trait Scenario: Send + Sync {
|
||||
async fn run(&self, ctx: &RunCtx) -> Result<ScenarioMetrics>;
|
||||
}
|
||||
|
||||
/// Build the active scenario set from config. One chat-latency scenario
|
||||
/// per configured prompt size.
|
||||
/// Build the active scenario set from config: one chat-latency scenario per
|
||||
/// prompt size, plus one concurrency scenario per configured level (#89).
|
||||
/// Concurrency levels default to empty (opt-in), since a burst puts real
|
||||
/// simultaneous load on a serving fleet — operators enable it deliberately.
|
||||
pub fn build_scenarios(cfg: &ScenarioConfig) -> Vec<Box<dyn Scenario>> {
|
||||
cfg.prompt_sizes
|
||||
let mut scenarios: Vec<Box<dyn Scenario>> = cfg
|
||||
.prompt_sizes
|
||||
.iter()
|
||||
.map(|&size| {
|
||||
Box::new(ChatLatencyScenario {
|
||||
@@ -96,7 +134,46 @@ pub fn build_scenarios(cfg: &ScenarioConfig) -> Vec<Box<dyn Scenario>> {
|
||||
approx_prompt_tokens: size,
|
||||
}) as Box<dyn Scenario>
|
||||
})
|
||||
.collect()
|
||||
.collect();
|
||||
for &n in &cfg.concurrency_levels {
|
||||
scenarios.push(Box::new(ConcurrencyScenario {
|
||||
id: format!("concurrency:{n}"),
|
||||
concurrency: n,
|
||||
approx_prompt_tokens: cfg.concurrency_prompt_tokens,
|
||||
}) as Box<dyn Scenario>);
|
||||
}
|
||||
for probe in &cfg.capability_probes {
|
||||
scenarios.push(Box::new(CapabilityScenario {
|
||||
id: format!("capability:{}", probe.name),
|
||||
prompt: probe.prompt.clone(),
|
||||
max_tokens: probe.max_tokens,
|
||||
}) as Box<dyn Scenario>);
|
||||
}
|
||||
scenarios
|
||||
}
|
||||
|
||||
/// A single small streamed request, timed like a chat-latency run. Used by
|
||||
/// the swap-cost measurement (#90) to capture the cold first-request latency
|
||||
/// straight after a reload. Reuses the shared SSE-timing core.
|
||||
pub async fn cold_probe(ctx: &RunCtx<'_>) -> Result<ScenarioMetrics> {
|
||||
let prompt = build_prompt(128);
|
||||
let payload = chat_payload(ctx, &prompt);
|
||||
tokio::time::timeout(ctx.timeout, stream_and_measure(ctx, &payload))
|
||||
.await
|
||||
.map_err(|_| anyhow!("cold probe timed out after {:?}", ctx.timeout))?
|
||||
}
|
||||
|
||||
/// The chat-completions request body shared by the latency and concurrency
|
||||
/// scenarios — streamed, deterministic (temperature 0), usage included.
|
||||
fn chat_payload(ctx: &RunCtx, prompt: &str) -> serde_json::Value {
|
||||
json!({
|
||||
"model": ctx.model_id,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"max_tokens": ctx.max_tokens,
|
||||
"temperature": 0,
|
||||
"stream": true,
|
||||
"stream_options": {"include_usage": true},
|
||||
})
|
||||
}
|
||||
|
||||
/// Streamed single-request chat-completions latency probe — the batch-1
|
||||
@@ -118,15 +195,7 @@ impl Scenario for ChatLatencyScenario {
|
||||
|
||||
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 payload = chat_payload(ctx, &prompt);
|
||||
let fut = stream_and_measure(ctx, &payload);
|
||||
tokio::time::timeout(ctx.timeout, fut)
|
||||
.await
|
||||
@@ -134,11 +203,182 @@ impl Scenario for ChatLatencyScenario {
|
||||
}
|
||||
}
|
||||
|
||||
/// Fan-out load probe: fire `concurrency` identical streams at once and
|
||||
/// measure how the fleet behaves under simultaneous pressure (#89). This is
|
||||
/// the only scenario that exercises the real a0/hermes/opencode pattern —
|
||||
/// many agentic requests per user turn — which batch-1 single-request
|
||||
/// timing cannot characterize. On a batch-1 serialized server, aggregate
|
||||
/// throughput stays ~flat while TTFT/queue-wait inflate with `concurrency`;
|
||||
/// that gap is the evidence for/against continuous batching.
|
||||
pub struct ConcurrencyScenario {
|
||||
id: String,
|
||||
concurrency: u32,
|
||||
approx_prompt_tokens: u32,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Scenario for ConcurrencyScenario {
|
||||
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 = chat_payload(ctx, &prompt);
|
||||
|
||||
// Fire all streams at once; each is independently timed and capped by
|
||||
// the per-request timeout so one hung stream can't stall the burst.
|
||||
let burst_start = Instant::now();
|
||||
let futs = (0..self.concurrency).map(|_| async {
|
||||
tokio::time::timeout(ctx.timeout, stream_and_measure(ctx, &payload)).await
|
||||
});
|
||||
let results = futures::future::join_all(futs).await;
|
||||
let burst_window = burst_start.elapsed().as_secs_f64();
|
||||
|
||||
let mut streams: Vec<ScenarioMetrics> = Vec::new();
|
||||
let mut rejected: u32 = 0;
|
||||
for r in results {
|
||||
match r {
|
||||
Ok(Ok(m)) => streams.push(m),
|
||||
// Admission backpressure (429/503) is shed load, counted
|
||||
// separately from genuine failures/timeouts.
|
||||
Ok(Err(e)) if is_admission_reject(&e) => rejected += 1,
|
||||
Ok(Err(_)) | Err(_) => {}
|
||||
}
|
||||
}
|
||||
if streams.is_empty() {
|
||||
return Err(anyhow!(
|
||||
"all {} concurrent streams failed ({rejected} shed by admission)",
|
||||
self.concurrency
|
||||
));
|
||||
}
|
||||
|
||||
let total_tokens: u64 = streams.iter().map(|m| m.completion_tokens).sum();
|
||||
let ttfts: Vec<f64> = streams.iter().map(|m| m.ttft_s).collect();
|
||||
// queue-wait ≈ TTFT − server prefill (#85); only for streams that
|
||||
// reported helexa_timing.
|
||||
let queue_waits: Vec<f64> = streams
|
||||
.iter()
|
||||
.filter_map(|m| {
|
||||
m.prefill_ms
|
||||
.map(|p| (m.ttft_s * 1000.0 - p as f64).max(0.0))
|
||||
})
|
||||
.collect();
|
||||
// Aggregate decode throughput across the whole node for the burst.
|
||||
let aggregate_tps = if burst_window > 0.0 {
|
||||
Some(total_tokens as f64 / burst_window)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
Ok(ScenarioMetrics {
|
||||
ttft_s: median(&ttfts).unwrap_or(0.0),
|
||||
decode_tps: aggregate_tps,
|
||||
total_s: burst_window,
|
||||
prompt_tokens: streams.iter().find_map(|m| m.prompt_tokens),
|
||||
completion_tokens: total_tokens,
|
||||
prefill_ms: None,
|
||||
decode_ms: None,
|
||||
prefill_tokens: None,
|
||||
concurrency: Some(self.concurrency),
|
||||
ttft_p95_s: percentile(&ttfts, 95.0),
|
||||
queue_wait_ms_median: median(&queue_waits),
|
||||
rejected: Some(rejected),
|
||||
artifact: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Quality probe (#91): runs a fixed prompt and stores the full generated
|
||||
/// text as an artifact for later scoring (manual now, LLM-judge later). The
|
||||
/// point is to compare reasoning/planning quality across models — the axis
|
||||
/// speed-only scenarios miss — so the frontier A/B (F3) picks on capability,
|
||||
/// not just throughput.
|
||||
pub struct CapabilityScenario {
|
||||
id: String,
|
||||
prompt: String,
|
||||
max_tokens: u64,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Scenario for CapabilityScenario {
|
||||
fn id(&self) -> &str {
|
||||
&self.id
|
||||
}
|
||||
|
||||
/// Capability probes have no synthetic prompt-token target; the cell is
|
||||
/// keyed by the scenario id alone.
|
||||
fn prompt_size(&self) -> u32 {
|
||||
0
|
||||
}
|
||||
|
||||
async fn run(&self, ctx: &RunCtx) -> Result<ScenarioMetrics> {
|
||||
let payload = json!({
|
||||
"model": ctx.model_id,
|
||||
"messages": [{"role": "user", "content": self.prompt}],
|
||||
"max_tokens": self.max_tokens,
|
||||
"temperature": 0,
|
||||
"stream": true,
|
||||
"stream_options": {"include_usage": true},
|
||||
});
|
||||
let fut = stream_and_measure_inner(ctx, &payload, true);
|
||||
tokio::time::timeout(ctx.timeout, fut)
|
||||
.await
|
||||
.map_err(|_| anyhow!("capability probe timed out after {:?}", ctx.timeout))?
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether a stream error was admission backpressure (HTTP 429/503) rather
|
||||
/// than a genuine failure. `stream_and_measure` renders the upstream status
|
||||
/// into the error string, so a substring check is sufficient.
|
||||
fn is_admission_reject(e: &anyhow::Error) -> bool {
|
||||
let s = e.to_string();
|
||||
s.contains("429") || s.contains("503")
|
||||
}
|
||||
|
||||
/// Median of a slice (sorted copy). `None` if empty.
|
||||
fn median(values: &[f64]) -> Option<f64> {
|
||||
if values.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let mut v = values.to_vec();
|
||||
v.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||
let lo = (v.len() - 1) / 2;
|
||||
let hi = v.len() / 2;
|
||||
Some((v[lo] + v[hi]) / 2.0)
|
||||
}
|
||||
|
||||
/// Nearest-rank percentile of a slice (`p` in 0..=100). `None` if empty.
|
||||
fn percentile(values: &[f64], p: f64) -> Option<f64> {
|
||||
if values.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let mut v = values.to_vec();
|
||||
v.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||
let rank = (p / 100.0 * v.len() as f64).ceil() as usize;
|
||||
Some(v[rank.clamp(1, v.len()) - 1])
|
||||
}
|
||||
|
||||
/// 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> {
|
||||
stream_and_measure_inner(ctx, payload, false).await
|
||||
}
|
||||
|
||||
/// As [`stream_and_measure`] but accumulates the full visible text when
|
||||
/// `capture_text` is set — used by the capability probe (#91) to store the
|
||||
/// generated artifact for later quality scoring.
|
||||
async fn stream_and_measure_inner(
|
||||
ctx: &RunCtx<'_>,
|
||||
payload: &serde_json::Value,
|
||||
capture_text: bool,
|
||||
) -> Result<ScenarioMetrics> {
|
||||
let start = Instant::now();
|
||||
let resp = ctx
|
||||
@@ -160,6 +400,10 @@ async fn stream_and_measure(
|
||||
let mut chunk_count: u64 = 0;
|
||||
let mut prompt_tokens: Option<u64> = None;
|
||||
let mut completion_tokens: Option<u64> = None;
|
||||
let mut prefill_ms: Option<u64> = None;
|
||||
let mut decode_ms: Option<u64> = None;
|
||||
let mut prefill_tokens: Option<u64> = None;
|
||||
let mut captured = String::new();
|
||||
|
||||
while let Some(event) = stream.next().await {
|
||||
let event = event.context("reading SSE stream")?;
|
||||
@@ -172,46 +416,86 @@ async fn stream_and_measure(
|
||||
Ok(c) => c,
|
||||
Err(_) => continue, // tolerate non-JSON keepalive frames
|
||||
};
|
||||
if let Some(choice) = chunk.choices.first()
|
||||
&& choice
|
||||
if let Some(choice) = chunk.choices.first() {
|
||||
// Liveness counts ANY generated delta (#117). Thinking
|
||||
// models (Qwen3-Next-Thinking, Qwen3 with thinking on)
|
||||
// stream `reasoning_content` first — sometimes for their
|
||||
// entire budget — and a content-only view misread that as
|
||||
// a dead stream ("no content chunks received") while also
|
||||
// producing impossible client-side rates (reasoning-
|
||||
// inclusive token counts over a visible-content-only
|
||||
// window; observed: "244 tok/s" on a 3060). For
|
||||
// non-thinking models the first delta IS content, so
|
||||
// `ttft_s` semantics are unchanged for them.
|
||||
let content = choice
|
||||
.delta
|
||||
.get("content")
|
||||
.and_then(|c| c.as_str())
|
||||
.is_some_and(|s| !s.is_empty())
|
||||
{
|
||||
if first.is_none() {
|
||||
first = Some(now);
|
||||
.filter(|c| !c.is_empty());
|
||||
let reasoning = choice
|
||||
.delta
|
||||
.get("reasoning_content")
|
||||
.and_then(|c| c.as_str())
|
||||
.filter(|c| !c.is_empty());
|
||||
if content.is_some() || reasoning.is_some() {
|
||||
if first.is_none() {
|
||||
first = Some(now);
|
||||
}
|
||||
last = Some(now);
|
||||
chunk_count += 1;
|
||||
}
|
||||
if capture_text && let Some(text) = content {
|
||||
captured.push_str(text);
|
||||
}
|
||||
last = Some(now);
|
||||
chunk_count += 1;
|
||||
}
|
||||
if let Some(usage) = chunk.usage {
|
||||
prompt_tokens = Some(usage.prompt_tokens);
|
||||
completion_tokens = Some(usage.completion_tokens);
|
||||
if let Some(t) = usage.helexa_timing {
|
||||
prefill_ms = Some(t.prefill_ms);
|
||||
decode_ms = Some(t.decode_ms);
|
||||
prefill_tokens = Some(t.prefill_tokens);
|
||||
}
|
||||
}
|
||||
}
|
||||
let end = Instant::now();
|
||||
|
||||
let first = first.ok_or_else(|| anyhow!("no content chunks received"))?;
|
||||
let first = first.ok_or_else(|| anyhow!("no generated chunks received"))?;
|
||||
|
||||
// neuron emits one SSE chunk per visible token, so chunk_count is an
|
||||
// engine-truth count when no usage frame is sent.
|
||||
// neuron emits one SSE chunk per generated 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.
|
||||
// Decode rate: prefer the server-measured split (#85) — it counts
|
||||
// every generated token over the actual decode window, immune to
|
||||
// reasoning-suppression frame mismatches. Fall back to the client
|
||||
// inter-chunk window with the CHUNK count (same frame) — never
|
||||
// usage.completion_tokens over the chunk window, which mixes a
|
||||
// reasoning-inclusive numerator with a visible-only denominator.
|
||||
let window = last
|
||||
.filter(|&l| l > first)
|
||||
.map(|l| (l - first).as_secs_f64())
|
||||
.unwrap_or(0.0);
|
||||
let decode_tps = match decode_ms {
|
||||
Some(ms) if ms > 200 && tokens > 0 => Some(tokens as f64 / (ms as f64 / 1000.0)),
|
||||
_ if window > 0.2 => Some(chunk_count as f64 / window),
|
||||
_ => None,
|
||||
};
|
||||
Ok(ScenarioMetrics {
|
||||
ttft_s: (first - start).as_secs_f64(),
|
||||
decode_tps: if window > 0.2 {
|
||||
Some(tokens as f64 / window)
|
||||
} else {
|
||||
None
|
||||
},
|
||||
decode_tps,
|
||||
total_s: (end - start).as_secs_f64(),
|
||||
prompt_tokens,
|
||||
completion_tokens: tokens,
|
||||
prefill_ms,
|
||||
decode_ms,
|
||||
prefill_tokens,
|
||||
// Concurrency fields unset on the single-request path; the
|
||||
// concurrency scenario builds its own aggregate (#89).
|
||||
concurrency: None,
|
||||
ttft_p95_s: None,
|
||||
queue_wait_ms_median: None,
|
||||
rejected: None,
|
||||
artifact: if capture_text { Some(captured) } else { None },
|
||||
})
|
||||
}
|
||||
|
||||
@@ -229,6 +513,54 @@ mod tests {
|
||||
assert!(small.ends_with("/no_think"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn median_and_percentile_basics() {
|
||||
assert_eq!(median(&[3.0, 1.0, 2.0]), Some(2.0));
|
||||
assert_eq!(median(&[]), None);
|
||||
let v = [1.0, 2.0, 3.0, 4.0, 5.0];
|
||||
assert_eq!(percentile(&v, 50.0), Some(3.0));
|
||||
assert_eq!(percentile(&v, 95.0), Some(5.0)); // nearest-rank → max with n=5
|
||||
assert_eq!(percentile(&[], 95.0), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn admission_rejects_detected_by_status() {
|
||||
assert!(is_admission_reject(&anyhow!(
|
||||
"upstream returned 429 Too Many Requests"
|
||||
)));
|
||||
assert!(is_admission_reject(&anyhow!(
|
||||
"upstream returned 503 Service Unavailable"
|
||||
)));
|
||||
assert!(!is_admission_reject(&anyhow!(
|
||||
"upstream returned 500 Internal"
|
||||
)));
|
||||
assert!(!is_admission_reject(&anyhow!("connection refused")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn concurrency_scenarios_built_from_config() {
|
||||
use crate::config::{CapabilityProbe, ScenarioConfig};
|
||||
let cfg = ScenarioConfig {
|
||||
prompt_sizes: vec![128],
|
||||
max_tokens: 64,
|
||||
concurrency_levels: vec![2, 8],
|
||||
concurrency_prompt_tokens: 512,
|
||||
capability_probes: vec![CapabilityProbe {
|
||||
name: "plan".into(),
|
||||
prompt: "Write a plan.".into(),
|
||||
max_tokens: 2048,
|
||||
}],
|
||||
};
|
||||
let ids: Vec<String> = build_scenarios(&cfg)
|
||||
.iter()
|
||||
.map(|s| s.id().to_string())
|
||||
.collect();
|
||||
assert!(ids.contains(&"chat:128".to_string()));
|
||||
assert!(ids.contains(&"concurrency:2".to_string()));
|
||||
assert!(ids.contains(&"concurrency:8".to_string()));
|
||||
assert!(ids.contains(&"capability:plan".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prompt_floor_for_tiny_targets() {
|
||||
// max(approx,16) floor means even 0 yields a non-trivial prompt.
|
||||
|
||||
@@ -51,6 +51,35 @@ pub struct RunRecord {
|
||||
pub decode_tps: Option<f64>,
|
||||
pub total_s: Option<f64>,
|
||||
pub completion_tokens: Option<u64>,
|
||||
// server-measured prefill/decode split (#85), null on engines/paths
|
||||
// that don't emit `usage.helexa_timing`.
|
||||
pub prefill_ms: Option<u64>,
|
||||
pub decode_ms: Option<u64>,
|
||||
pub prefill_tokens: Option<u64>,
|
||||
// GPU telemetry sampled from /health around the run (#87), null for
|
||||
// non-neuron targets or when /health was unreachable. vram_used_mb is
|
||||
// the node sum; util/temp are the hottest single device.
|
||||
pub vram_used_mb: Option<u64>,
|
||||
pub gpu_util_pct: Option<u32>,
|
||||
pub gpu_temp_c: Option<u32>,
|
||||
// concurrency / agentic-load burst metrics (#89), null for single-request
|
||||
// scenarios. For a burst, ttft_s/decode_tps/total_s carry the aggregate.
|
||||
pub concurrency: Option<u32>,
|
||||
pub ttft_p95_s: Option<f64>,
|
||||
pub queue_wait_ms: Option<f64>,
|
||||
pub rejected: Option<u32>,
|
||||
// cold-load / model-swap cost (#90), set only by the deliberate
|
||||
// `swap-cost` measurement (scenario_id = "swap"). The other metric
|
||||
// fields carry the cold first-request after reload.
|
||||
pub swap_unload_ms: Option<u64>,
|
||||
pub swap_load_ms: Option<u64>,
|
||||
// capability probe (#91): the full generated text (scenario_id =
|
||||
// "capability:<name>"), scored for quality later. `quality_score` /
|
||||
// `scorer` are null at insert — set by `score` (manual) or a future
|
||||
// LLM-judge.
|
||||
pub artifact: Option<String>,
|
||||
pub quality_score: Option<f64>,
|
||||
pub scorer: Option<String>,
|
||||
// outcome
|
||||
pub ok: bool,
|
||||
pub error: Option<String>,
|
||||
@@ -123,6 +152,21 @@ impl Store {
|
||||
decode_tps REAL,
|
||||
total_s REAL,
|
||||
completion_tokens INTEGER,
|
||||
prefill_ms INTEGER,
|
||||
decode_ms INTEGER,
|
||||
prefill_tokens INTEGER,
|
||||
vram_used_mb INTEGER,
|
||||
gpu_util_pct INTEGER,
|
||||
gpu_temp_c INTEGER,
|
||||
concurrency INTEGER,
|
||||
ttft_p95_s REAL,
|
||||
queue_wait_ms REAL,
|
||||
rejected INTEGER,
|
||||
swap_unload_ms INTEGER,
|
||||
swap_load_ms INTEGER,
|
||||
artifact TEXT,
|
||||
quality_score REAL,
|
||||
scorer TEXT,
|
||||
ok INTEGER NOT NULL,
|
||||
error TEXT
|
||||
);
|
||||
@@ -133,6 +177,51 @@ impl Store {
|
||||
"#,
|
||||
)
|
||||
.context("initialising sqlite schema")?;
|
||||
// Additive migrations for DBs created before a column existed.
|
||||
// `CREATE TABLE IF NOT EXISTS` above only seeds fresh DBs; existing
|
||||
// ones need the columns backfilled (as NULL) so older rows coexist
|
||||
// with new metrics. There is no migration framework — each entry is
|
||||
// an idempotent "add if missing".
|
||||
Self::ensure_columns(
|
||||
conn,
|
||||
"runs",
|
||||
&[
|
||||
("prefill_ms", "INTEGER"),
|
||||
("decode_ms", "INTEGER"),
|
||||
("prefill_tokens", "INTEGER"),
|
||||
("vram_used_mb", "INTEGER"),
|
||||
("gpu_util_pct", "INTEGER"),
|
||||
("gpu_temp_c", "INTEGER"),
|
||||
("concurrency", "INTEGER"),
|
||||
("ttft_p95_s", "REAL"),
|
||||
("queue_wait_ms", "REAL"),
|
||||
("rejected", "INTEGER"),
|
||||
("swap_unload_ms", "INTEGER"),
|
||||
("swap_load_ms", "INTEGER"),
|
||||
("artifact", "TEXT"),
|
||||
("quality_score", "REAL"),
|
||||
("scorer", "TEXT"),
|
||||
],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Add any of `columns` that the table is missing (`ALTER TABLE ADD
|
||||
/// COLUMN`). Idempotent: existing columns are read from
|
||||
/// `PRAGMA table_info` and skipped, so this is safe to run on every open.
|
||||
fn ensure_columns(conn: &Connection, table: &str, columns: &[(&str, &str)]) -> Result<()> {
|
||||
let mut existing = std::collections::HashSet::new();
|
||||
let mut stmt = conn.prepare(&format!("PRAGMA table_info({table})"))?;
|
||||
let names = stmt.query_map([], |row| row.get::<_, String>(1))?;
|
||||
for name in names {
|
||||
existing.insert(name?);
|
||||
}
|
||||
for (name, ty) in columns {
|
||||
if !existing.contains(*name) {
|
||||
conn.execute_batch(&format!("ALTER TABLE {table} ADD COLUMN {name} {ty};"))
|
||||
.with_context(|| format!("adding column {table}.{name}"))?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -166,6 +255,11 @@ impl Store {
|
||||
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,
|
||||
prefill_ms, decode_ms, prefill_tokens,
|
||||
vram_used_mb, gpu_util_pct, gpu_temp_c,
|
||||
concurrency, ttft_p95_s, queue_wait_ms, rejected,
|
||||
swap_unload_ms, swap_load_ms,
|
||||
artifact, quality_score, scorer,
|
||||
ok, error
|
||||
) VALUES (
|
||||
?1, ?2, ?3, ?4,
|
||||
@@ -176,7 +270,12 @@ impl Store {
|
||||
?20, ?21, ?22, ?23,
|
||||
?24, ?25, ?26, ?27,
|
||||
?28, ?29, ?30, ?31,
|
||||
?32, ?33
|
||||
?32, ?33, ?34,
|
||||
?35, ?36, ?37,
|
||||
?38, ?39, ?40, ?41,
|
||||
?42, ?43,
|
||||
?44, ?45, ?46,
|
||||
?47, ?48
|
||||
)",
|
||||
params![
|
||||
r.ts,
|
||||
@@ -210,6 +309,21 @@ impl Store {
|
||||
r.decode_tps,
|
||||
r.total_s,
|
||||
r.completion_tokens,
|
||||
r.prefill_ms,
|
||||
r.decode_ms,
|
||||
r.prefill_tokens,
|
||||
r.vram_used_mb,
|
||||
r.gpu_util_pct,
|
||||
r.gpu_temp_c,
|
||||
r.concurrency,
|
||||
r.ttft_p95_s,
|
||||
r.queue_wait_ms,
|
||||
r.rejected,
|
||||
r.swap_unload_ms,
|
||||
r.swap_load_ms,
|
||||
r.artifact,
|
||||
r.quality_score,
|
||||
r.scorer,
|
||||
r.ok as i64,
|
||||
r.error,
|
||||
],
|
||||
@@ -224,7 +338,10 @@ impl Store {
|
||||
// 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, gpus_json
|
||||
ttft_s, decode_tps, total_s, prompt_tokens_actual, gpus_json,
|
||||
prefill_ms, decode_ms, prefill_tokens,
|
||||
vram_used_mb, gpu_util_pct, gpu_temp_c,
|
||||
concurrency, ttft_p95_s, queue_wait_ms, rejected
|
||||
FROM runs
|
||||
WHERE ok=1
|
||||
ORDER BY target_name, model_id, scenario_id, id",
|
||||
@@ -241,12 +358,188 @@ impl Store {
|
||||
total_s: row.get(7)?,
|
||||
prompt_tokens_actual: row.get(8)?,
|
||||
gpus_json: row.get(9)?,
|
||||
prefill_ms: row.get(10)?,
|
||||
decode_ms: row.get(11)?,
|
||||
prefill_tokens: row.get(12)?,
|
||||
vram_used_mb: row.get(13)?,
|
||||
gpu_util_pct: row.get(14)?,
|
||||
gpu_temp_c: row.get(15)?,
|
||||
concurrency: row.get(16)?,
|
||||
ttft_p95_s: row.get(17)?,
|
||||
queue_wait_ms: row.get(18)?,
|
||||
rejected: row.get(19)?,
|
||||
})
|
||||
})?;
|
||||
let raws: Vec<RawRow> = rows.collect::<rusqlite::Result<_>>()?;
|
||||
Ok(aggregate(raws))
|
||||
}
|
||||
|
||||
/// Context-length scaling curves (#88): per (target, model), the latest
|
||||
/// build's `chat:<n>` cells pivoted by prompt size into prefill & decode
|
||||
/// tok/s vs context. The headline is `decode_flatness` — decode tok/s at
|
||||
/// the largest context divided by the smallest. Near 1.0 confirms the
|
||||
/// Gated-DeltaNet O(1)-in-sequence-length decode; a sharp drop locates
|
||||
/// where the model stops scaling for free.
|
||||
pub fn scaling(&self) -> Result<Vec<ScalingCurve>> {
|
||||
use std::collections::BTreeMap;
|
||||
// Reuse the already-aggregated report cells; the chat:<n> rows are
|
||||
// exactly the per-context measurement points.
|
||||
let mut by_model: BTreeMap<(String, String), Vec<ReportRow>> = BTreeMap::new();
|
||||
for r in self.report_rows()? {
|
||||
if r.scenario_id.starts_with("chat:") {
|
||||
by_model
|
||||
.entry((r.target_name.clone(), r.model_id.clone()))
|
||||
.or_default()
|
||||
.push(r);
|
||||
}
|
||||
}
|
||||
let mut out = Vec::new();
|
||||
for ((target_name, model_id), mut rows) in by_model {
|
||||
rows.sort_by_key(|r| r.prompt_size_approx);
|
||||
let points: Vec<ScalingPoint> = rows
|
||||
.iter()
|
||||
.map(|r| ScalingPoint {
|
||||
prompt_size: r.prompt_size_approx,
|
||||
prompt_tokens: r.prompt_tokens,
|
||||
prefill_tps: r.prefill_tps_median,
|
||||
decode_tps: r.decode_tps_median,
|
||||
samples: r.samples,
|
||||
})
|
||||
.collect();
|
||||
// Flatness across the smallest→largest points that both have a
|
||||
// decode rate (skips cells where the decode window was too short).
|
||||
let with_decode: Vec<&ScalingPoint> =
|
||||
points.iter().filter(|p| p.decode_tps.is_some()).collect();
|
||||
let decode_flatness = match (with_decode.first(), with_decode.last()) {
|
||||
(Some(lo), Some(hi)) if with_decode.len() >= 2 => {
|
||||
match (lo.decode_tps, hi.decode_tps) {
|
||||
(Some(a), Some(b)) if a > 0.0 => Some(b / a),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
out.push(ScalingCurve {
|
||||
target_name,
|
||||
model_id,
|
||||
git_sha: rows.first().map(|r| r.git_sha.clone()).unwrap_or_default(),
|
||||
gpu: rows.iter().find_map(|r| r.gpu.clone()),
|
||||
points,
|
||||
decode_flatness,
|
||||
});
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Cold-load / model-swap costs (#90): per (target, model) at the latest
|
||||
/// build, the median unload→reload latency and the cold first-request
|
||||
/// latency after reload (the `scenario_id = "swap"` rows).
|
||||
pub fn swap_costs(&self) -> Result<Vec<SwapCost>> {
|
||||
use std::collections::BTreeMap;
|
||||
let mut stmt = self.conn.prepare(
|
||||
"SELECT target_name, model_id, git_sha, gpus_json,
|
||||
swap_unload_ms, swap_load_ms, ttft_s, total_s
|
||||
FROM runs
|
||||
WHERE ok=1 AND scenario_id='swap'
|
||||
ORDER BY target_name, model_id, id",
|
||||
)?;
|
||||
struct Raw {
|
||||
target: String,
|
||||
model: String,
|
||||
sha: String,
|
||||
gpus_json: Option<String>,
|
||||
unload_ms: Option<f64>,
|
||||
load_ms: Option<f64>,
|
||||
ttft_s: Option<f64>,
|
||||
total_s: Option<f64>,
|
||||
}
|
||||
let raws: Vec<Raw> = stmt
|
||||
.query_map([], |r| {
|
||||
Ok(Raw {
|
||||
target: r.get(0)?,
|
||||
model: r.get(1)?,
|
||||
sha: r.get(2)?,
|
||||
gpus_json: r.get(3)?,
|
||||
unload_ms: r.get::<_, Option<i64>>(4)?.map(|v| v as f64),
|
||||
load_ms: r.get::<_, Option<i64>>(5)?.map(|v| v as f64),
|
||||
ttft_s: r.get(6)?,
|
||||
total_s: r.get(7)?,
|
||||
})
|
||||
})?
|
||||
.collect::<rusqlite::Result<_>>()?;
|
||||
let mut by: BTreeMap<(String, String), Vec<Raw>> = BTreeMap::new();
|
||||
for r in raws {
|
||||
by.entry((r.target.clone(), r.model.clone()))
|
||||
.or_default()
|
||||
.push(r);
|
||||
}
|
||||
let mut out = Vec::new();
|
||||
for ((target, model), rows) in by {
|
||||
let latest = rows.last().map(|r| r.sha.clone()).unwrap_or_default();
|
||||
let cell: Vec<&Raw> = rows.iter().filter(|r| r.sha == latest).collect();
|
||||
out.push(SwapCost {
|
||||
target_name: target,
|
||||
model_id: model,
|
||||
git_sha: latest,
|
||||
gpu: cell
|
||||
.iter()
|
||||
.find_map(|r| r.gpus_json.as_deref().and_then(gpu_label)),
|
||||
unload_ms_median: median(cell.iter().filter_map(|r| r.unload_ms)),
|
||||
load_ms_median: median(cell.iter().filter_map(|r| r.load_ms)),
|
||||
cold_ttft_s_median: median(cell.iter().filter_map(|r| r.ttft_s)),
|
||||
cold_total_s_median: median(cell.iter().filter_map(|r| r.total_s)),
|
||||
samples: cell.len(),
|
||||
});
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Capability-probe runs (#91): the stored output artifacts and their
|
||||
/// quality scores. `unscored_only` filters to runs awaiting a score —
|
||||
/// the worklist for manual scoring or a future LLM-judge.
|
||||
pub fn capability_runs(&self, unscored_only: bool) -> Result<Vec<CapabilityRun>> {
|
||||
let sql = format!(
|
||||
"SELECT id, ts, target_name, model_id, scenario_id, git_sha,
|
||||
quality_score, scorer, artifact
|
||||
FROM runs
|
||||
WHERE ok=1 AND scenario_id LIKE 'capability:%'{}
|
||||
ORDER BY id DESC",
|
||||
if unscored_only {
|
||||
" AND quality_score IS NULL"
|
||||
} else {
|
||||
""
|
||||
},
|
||||
);
|
||||
let mut stmt = self.conn.prepare(&sql)?;
|
||||
let rows = stmt
|
||||
.query_map([], |r| {
|
||||
Ok(CapabilityRun {
|
||||
id: r.get(0)?,
|
||||
ts: r.get(1)?,
|
||||
target_name: r.get(2)?,
|
||||
model_id: r.get(3)?,
|
||||
scenario_id: r.get(4)?,
|
||||
git_sha: r.get(5)?,
|
||||
quality_score: r.get(6)?,
|
||||
scorer: r.get(7)?,
|
||||
artifact: r.get(8)?,
|
||||
})
|
||||
})?
|
||||
.collect::<rusqlite::Result<_>>()?;
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
/// Attach a quality score to one capability run (#91). `scorer` records
|
||||
/// who/what scored it ("manual", "llm:claude-…"). Returns the row count
|
||||
/// updated (0 if the id doesn't exist).
|
||||
pub fn set_score(&self, run_id: i64, score: f64, scorer: &str) -> Result<usize> {
|
||||
let n = self.conn.execute(
|
||||
"UPDATE runs SET quality_score=?1, scorer=?2 WHERE id=?3",
|
||||
params![score, scorer, run_id],
|
||||
)?;
|
||||
Ok(n)
|
||||
}
|
||||
|
||||
// ── Read API surface (consumed by api.rs) ─────────────────────────
|
||||
|
||||
/// Total recorded runs (for `/api/health`).
|
||||
@@ -379,7 +672,10 @@ impl Store {
|
||||
"SELECT id, ts, target_name, hostname, git_sha, build_timestamp, package_version,
|
||||
model_id, harness, scenario_id, prompt_size_approx, prompt_tokens_actual,
|
||||
max_tokens, ttft_s, decode_tps, total_s, completion_tokens, ok, error,
|
||||
gpus_json
|
||||
gpus_json, prefill_ms, decode_ms, prefill_tokens,
|
||||
vram_used_mb, gpu_util_pct, gpu_temp_c,
|
||||
concurrency, ttft_p95_s, queue_wait_ms, rejected,
|
||||
swap_unload_ms, swap_load_ms
|
||||
FROM runs",
|
||||
);
|
||||
let mut conds: Vec<String> = Vec::new();
|
||||
@@ -435,6 +731,18 @@ impl Store {
|
||||
completion_tokens: r.get(16)?,
|
||||
ok: r.get::<_, i64>(17)? != 0,
|
||||
error: r.get(18)?,
|
||||
prefill_ms: r.get(20)?,
|
||||
decode_ms: r.get(21)?,
|
||||
prefill_tokens: r.get(22)?,
|
||||
vram_used_mb: r.get(23)?,
|
||||
gpu_util_pct: r.get(24)?,
|
||||
gpu_temp_c: r.get(25)?,
|
||||
concurrency: r.get(26)?,
|
||||
ttft_p95_s: r.get(27)?,
|
||||
queue_wait_ms: r.get(28)?,
|
||||
rejected: r.get(29)?,
|
||||
swap_unload_ms: r.get(30)?,
|
||||
swap_load_ms: r.get(31)?,
|
||||
})
|
||||
})?
|
||||
.collect::<rusqlite::Result<_>>()?;
|
||||
@@ -554,6 +862,18 @@ pub struct RunRow {
|
||||
pub decode_tps: Option<f64>,
|
||||
pub total_s: Option<f64>,
|
||||
pub completion_tokens: Option<u64>,
|
||||
pub prefill_ms: Option<u64>,
|
||||
pub decode_ms: Option<u64>,
|
||||
pub prefill_tokens: Option<u64>,
|
||||
pub vram_used_mb: Option<u64>,
|
||||
pub gpu_util_pct: Option<u64>,
|
||||
pub gpu_temp_c: Option<u64>,
|
||||
pub concurrency: Option<u64>,
|
||||
pub ttft_p95_s: Option<f64>,
|
||||
pub queue_wait_ms: Option<f64>,
|
||||
pub rejected: Option<u64>,
|
||||
pub swap_unload_ms: Option<u64>,
|
||||
pub swap_load_ms: Option<u64>,
|
||||
pub ok: bool,
|
||||
pub error: Option<String>,
|
||||
}
|
||||
@@ -569,6 +889,16 @@ struct RawRow {
|
||||
total_s: Option<f64>,
|
||||
prompt_tokens_actual: Option<u64>,
|
||||
gpus_json: Option<String>,
|
||||
prefill_ms: Option<u64>,
|
||||
decode_ms: Option<u64>,
|
||||
prefill_tokens: Option<u64>,
|
||||
vram_used_mb: Option<u64>,
|
||||
gpu_util_pct: Option<u64>,
|
||||
gpu_temp_c: Option<u64>,
|
||||
concurrency: Option<u64>,
|
||||
ttft_p95_s: Option<f64>,
|
||||
queue_wait_ms: Option<f64>,
|
||||
rejected: Option<u64>,
|
||||
}
|
||||
|
||||
/// An aggregated cell ready for the report table.
|
||||
@@ -583,11 +913,100 @@ pub struct ReportRow {
|
||||
pub ttft_s_median: Option<f64>,
|
||||
pub decode_tps_median: Option<f64>,
|
||||
pub total_s_median: Option<f64>,
|
||||
/// Latency tail percentiles — where batch-1 pain actually shows up, and
|
||||
/// invisible behind a bare median. p95/p99 nearest-rank; with few
|
||||
/// samples they collapse toward the max (honest, not interpolated).
|
||||
pub ttft_s_p95: Option<f64>,
|
||||
pub ttft_s_p99: Option<f64>,
|
||||
pub total_s_p95: Option<f64>,
|
||||
pub total_s_p99: Option<f64>,
|
||||
/// Server-measured prefill/decode split (#85). `prefill_tps_median` is
|
||||
/// the true prompt-encoding rate (prefill_tokens / prefill_ms),
|
||||
/// complementing `decode_tps_median` (the generation rate).
|
||||
pub prefill_ms_median: Option<f64>,
|
||||
pub decode_ms_median: Option<f64>,
|
||||
pub prefill_tps_median: Option<f64>,
|
||||
/// GPU telemetry sampled from /health around the run (#87). `vram_used_mb`
|
||||
/// is the node sum; `vram_total_mb` (from discovery) lets the report show
|
||||
/// real headroom — the "2/3 used" hunch as a number. util/temp are the
|
||||
/// hottest device. All `None` for non-neuron targets.
|
||||
pub vram_used_mb_median: Option<f64>,
|
||||
pub vram_total_mb: Option<u64>,
|
||||
pub gpu_util_pct_median: Option<f64>,
|
||||
pub gpu_temp_c_median: Option<f64>,
|
||||
/// Concurrency / agentic-load burst metrics (#89). `concurrency` is the
|
||||
/// burst width (constant per cell). `ttft_p95_load_s` is the within-burst
|
||||
/// TTFT tail; `queue_wait_ms_median` the admission wait; `rejected_median`
|
||||
/// the per-burst shed count. All `None` for non-concurrency scenarios.
|
||||
pub concurrency: Option<u64>,
|
||||
pub ttft_p95_load_s: Option<f64>,
|
||||
pub queue_wait_ms_median: Option<f64>,
|
||||
pub rejected_median: Option<f64>,
|
||||
pub samples: usize,
|
||||
/// Public-facing resource name (the host's GPU(s)), e.g. "2× RTX 5090".
|
||||
pub gpu: Option<String>,
|
||||
}
|
||||
|
||||
/// One context-length scaling curve for a (target, model) at its latest
|
||||
/// build — the points ordered by prompt size, plus the decode-flatness
|
||||
/// summary (#88).
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct ScalingCurve {
|
||||
pub target_name: String,
|
||||
pub model_id: String,
|
||||
pub git_sha: String,
|
||||
pub gpu: Option<String>,
|
||||
pub points: Vec<ScalingPoint>,
|
||||
/// decode tok/s at the largest context ÷ at the smallest. ~1.0 = flat
|
||||
/// (GDN O(1) decode); <1 quantifies the drop-off. `None` with <2 points.
|
||||
pub decode_flatness: Option<f64>,
|
||||
}
|
||||
|
||||
/// Cold-load / model-swap cost for a (target, model) at its latest build
|
||||
/// (#90): the reload latency and the cold first-request after it.
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct SwapCost {
|
||||
pub target_name: String,
|
||||
pub model_id: String,
|
||||
pub git_sha: String,
|
||||
pub gpu: Option<String>,
|
||||
/// Time to free the model (`POST /models/unload`), ms.
|
||||
pub unload_ms_median: Option<f64>,
|
||||
/// Time to reload it (`POST /models/load`, synchronous), ms — the
|
||||
/// headline swap cost feeding the vision cold-swap policy (F4e).
|
||||
pub load_ms_median: Option<f64>,
|
||||
/// TTFT of the first request after reload (cold caches), seconds.
|
||||
pub cold_ttft_s_median: Option<f64>,
|
||||
/// Total wall-clock of that cold first request, seconds.
|
||||
pub cold_total_s_median: Option<f64>,
|
||||
pub samples: usize,
|
||||
}
|
||||
|
||||
/// One capability-probe run (#91): the stored output artifact plus its
|
||||
/// quality score (null until scored manually or by a future LLM-judge).
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct CapabilityRun {
|
||||
pub id: i64,
|
||||
pub ts: String,
|
||||
pub target_name: String,
|
||||
pub model_id: String,
|
||||
pub scenario_id: String,
|
||||
pub git_sha: String,
|
||||
pub quality_score: Option<f64>,
|
||||
pub scorer: Option<String>,
|
||||
pub artifact: Option<String>,
|
||||
}
|
||||
|
||||
/// One point on a [`ScalingCurve`]: the throughput at a given context size.
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct ScalingPoint {
|
||||
pub prompt_size: u32,
|
||||
pub prompt_tokens: Option<u64>,
|
||||
pub prefill_tps: Option<f64>,
|
||||
pub decode_tps: 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.
|
||||
@@ -611,6 +1030,11 @@ fn aggregate(raws: Vec<RawRow>) -> Vec<ReportRow> {
|
||||
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);
|
||||
// Per-row prefill tok/s, derived from the server-measured split.
|
||||
let prefill_tps = |r: &&RawRow| match (r.prefill_tokens, r.prefill_ms) {
|
||||
(Some(tok), Some(ms)) if ms > 0 => Some(tok as f64 * 1000.0 / ms as f64),
|
||||
_ => None,
|
||||
};
|
||||
out.push(ReportRow {
|
||||
target_name,
|
||||
model_id,
|
||||
@@ -621,6 +1045,27 @@ fn aggregate(raws: Vec<RawRow>) -> Vec<ReportRow> {
|
||||
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)),
|
||||
ttft_s_p95: percentile(cell.iter().filter_map(|r| r.ttft_s), 95.0),
|
||||
ttft_s_p99: percentile(cell.iter().filter_map(|r| r.ttft_s), 99.0),
|
||||
total_s_p95: percentile(cell.iter().filter_map(|r| r.total_s), 95.0),
|
||||
total_s_p99: percentile(cell.iter().filter_map(|r| r.total_s), 99.0),
|
||||
prefill_ms_median: median(cell.iter().filter_map(|r| r.prefill_ms.map(|m| m as f64))),
|
||||
decode_ms_median: median(cell.iter().filter_map(|r| r.decode_ms.map(|m| m as f64))),
|
||||
prefill_tps_median: median(cell.iter().filter_map(prefill_tps)),
|
||||
vram_used_mb_median: median(
|
||||
cell.iter().filter_map(|r| r.vram_used_mb.map(|v| v as f64)),
|
||||
),
|
||||
vram_total_mb: cell
|
||||
.iter()
|
||||
.find_map(|r| r.gpus_json.as_deref().and_then(gpu_total_vram_mb)),
|
||||
gpu_util_pct_median: median(
|
||||
cell.iter().filter_map(|r| r.gpu_util_pct.map(|v| v as f64)),
|
||||
),
|
||||
gpu_temp_c_median: median(cell.iter().filter_map(|r| r.gpu_temp_c.map(|v| v as f64))),
|
||||
concurrency: cell.iter().find_map(|r| r.concurrency),
|
||||
ttft_p95_load_s: median(cell.iter().filter_map(|r| r.ttft_p95_s)),
|
||||
queue_wait_ms_median: median(cell.iter().filter_map(|r| r.queue_wait_ms)),
|
||||
rejected_median: median(cell.iter().filter_map(|r| r.rejected.map(|v| v as f64))),
|
||||
samples: cell.len(),
|
||||
gpu: cell
|
||||
.iter()
|
||||
@@ -630,6 +1075,19 @@ fn aggregate(raws: Vec<RawRow>) -> Vec<ReportRow> {
|
||||
out
|
||||
}
|
||||
|
||||
/// Node total VRAM in MB, summed across the devices in a run's stored
|
||||
/// `gpus_json` (the discovery `DeviceInfo` list, each with `vram_total_mb`).
|
||||
/// Pairs with the sampled `vram_used_mb` to report real headroom (#87).
|
||||
/// `None` when empty/absent or no device declares a total.
|
||||
fn gpu_total_vram_mb(gpus_json: &str) -> Option<u64> {
|
||||
let devices: Vec<serde_json::Value> = serde_json::from_str(gpus_json).ok()?;
|
||||
let total: u64 = devices
|
||||
.iter()
|
||||
.filter_map(|d| d.get("vram_total_mb").and_then(|v| v.as_u64()))
|
||||
.sum();
|
||||
(total > 0).then_some(total)
|
||||
}
|
||||
|
||||
/// Compact GPU label from a run's stored `gpus_json` (the discovery device
|
||||
/// list) — e.g. "2× RTX 5090", "RTX 4090". `None` when empty/absent. Used
|
||||
/// as the public-facing resource name in place of internal hostnames.
|
||||
@@ -680,6 +1138,22 @@ fn median(values: impl Iterator<Item = f64>) -> Option<f64> {
|
||||
Some((v[lo] + v[hi]) / 2.0)
|
||||
}
|
||||
|
||||
/// Nearest-rank percentile (`p` in 0..=100). Chosen over interpolation
|
||||
/// because bench cells hold only a handful of samples: with n=5, p95/p99
|
||||
/// resolve to the max, which honestly says "this is the worst we saw"
|
||||
/// rather than inventing a value between samples we never observed.
|
||||
fn percentile(values: impl Iterator<Item = f64>, p: 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));
|
||||
// rank = ceil(p/100 * n), clamped to [1, n]; index is rank-1.
|
||||
let rank = (p / 100.0 * v.len() as f64).ceil() as usize;
|
||||
let idx = rank.clamp(1, v.len()) - 1;
|
||||
Some(v[idx])
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -717,6 +1191,21 @@ mod tests {
|
||||
decode_tps: Some(50.0),
|
||||
total_s: Some(1.0),
|
||||
completion_tokens: Some(50),
|
||||
prefill_ms: Some(200),
|
||||
decode_ms: Some(1000),
|
||||
prefill_tokens: Some(130),
|
||||
vram_used_mb: Some(42000),
|
||||
gpu_util_pct: Some(88),
|
||||
gpu_temp_c: Some(64),
|
||||
concurrency: None,
|
||||
ttft_p95_s: None,
|
||||
queue_wait_ms: None,
|
||||
rejected: None,
|
||||
swap_unload_ms: None,
|
||||
swap_load_ms: None,
|
||||
artifact: None,
|
||||
quality_score: None,
|
||||
scorer: None,
|
||||
ok,
|
||||
error: if ok { None } else { Some("boom".into()) },
|
||||
}
|
||||
@@ -755,6 +1244,197 @@ mod tests {
|
||||
assert!((rows[0].ttft_s_median.unwrap() - 0.3).abs() < 1e-9);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn report_surfaces_percentiles_and_prefill_split() {
|
||||
let s = Store::open_in_memory().unwrap();
|
||||
// Five samples on one cell with spread TTFT so percentiles differ
|
||||
// from the median, plus a server-measured prefill/decode split.
|
||||
for (i, ttft) in [0.10, 0.12, 0.14, 0.16, 0.50].iter().enumerate() {
|
||||
let mut r = rec("beast", "sha", "m", "chat:128", true);
|
||||
r.ttft_s = Some(*ttft);
|
||||
r.total_s = Some(ttft + 1.0);
|
||||
r.prefill_ms = Some(200 + i as u64);
|
||||
r.prefill_tokens = Some(400);
|
||||
s.insert_run(&r).unwrap();
|
||||
}
|
||||
let rows = s.report_rows().unwrap();
|
||||
assert_eq!(rows.len(), 1);
|
||||
let row = &rows[0];
|
||||
assert_eq!(row.samples, 5);
|
||||
// p50 is the middle value; p95/p99 (nearest-rank, n=5) hit the max.
|
||||
assert!((row.ttft_s_median.unwrap() - 0.14).abs() < 1e-9);
|
||||
assert!((row.ttft_s_p95.unwrap() - 0.50).abs() < 1e-9);
|
||||
assert!((row.ttft_s_p99.unwrap() - 0.50).abs() < 1e-9);
|
||||
// prefill tok/s = 400 tok / ~0.2 s ≈ 2000 tok/s.
|
||||
assert!(row.prefill_tps_median.unwrap() > 1900.0);
|
||||
assert!(row.prefill_ms_median.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn report_surfaces_vram_and_gpu_telemetry() {
|
||||
let s = Store::open_in_memory().unwrap();
|
||||
let mut r = rec("beast", "sha", "m", "chat:128", true);
|
||||
// Node total VRAM from discovery devices → headroom denominator.
|
||||
r.gpus_json =
|
||||
Some(r#"[{"name":"RTX 5090","vram_total_mb":32000},{"name":"RTX 5090","vram_total_mb":32000}]"#.into());
|
||||
r.vram_used_mb = Some(42000);
|
||||
r.gpu_util_pct = Some(90);
|
||||
r.gpu_temp_c = Some(66);
|
||||
s.insert_run(&r).unwrap();
|
||||
let rows = s.report_rows().unwrap();
|
||||
let row = &rows[0];
|
||||
assert_eq!(row.vram_used_mb_median, Some(42000.0));
|
||||
assert_eq!(row.vram_total_mb, Some(64000)); // 2× 32000
|
||||
assert_eq!(row.gpu_util_pct_median, Some(90.0));
|
||||
assert_eq!(row.gpu_temp_c_median, Some(66.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn report_surfaces_concurrency_burst_metrics() {
|
||||
let s = Store::open_in_memory().unwrap();
|
||||
// Two concurrency:8 burst-runs with shed load and a queue-wait tail.
|
||||
for (qw, rej) in [(120.0, 1u32), (180.0, 3u32)] {
|
||||
let mut r = rec("beast", "sha", "m", "concurrency:8", true);
|
||||
r.concurrency = Some(8);
|
||||
r.ttft_p95_s = Some(0.9);
|
||||
r.queue_wait_ms = Some(qw);
|
||||
r.rejected = Some(rej);
|
||||
s.insert_run(&r).unwrap();
|
||||
}
|
||||
let row = s
|
||||
.report_rows()
|
||||
.unwrap()
|
||||
.into_iter()
|
||||
.find(|r| r.scenario_id == "concurrency:8")
|
||||
.unwrap();
|
||||
assert_eq!(row.concurrency, Some(8));
|
||||
assert_eq!(row.queue_wait_ms_median, Some(150.0)); // median(120,180)
|
||||
assert_eq!(row.rejected_median, Some(2.0)); // median(1,3)
|
||||
assert_eq!(row.ttft_p95_load_s, Some(0.9));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capability_runs_store_artifact_and_accept_scores() {
|
||||
let s = Store::open_in_memory().unwrap();
|
||||
let mut r = rec("beast", "sha", "m", "capability:plan", true);
|
||||
r.artifact = Some("a thorough implementation plan...".into());
|
||||
s.insert_run(&r).unwrap();
|
||||
// A non-capability row must not appear in capability_runs.
|
||||
s.insert_run(&rec("beast", "sha", "m", "chat:128", true))
|
||||
.unwrap();
|
||||
|
||||
let unscored = s.capability_runs(true).unwrap();
|
||||
assert_eq!(unscored.len(), 1);
|
||||
let id = unscored[0].id;
|
||||
assert!(unscored[0].artifact.as_deref().unwrap().contains("plan"));
|
||||
assert!(unscored[0].quality_score.is_none());
|
||||
|
||||
// Score it; it then drops off the unscored worklist.
|
||||
assert_eq!(s.set_score(id, 7.5, "manual").unwrap(), 1);
|
||||
assert!(s.capability_runs(true).unwrap().is_empty());
|
||||
let all = s.capability_runs(false).unwrap();
|
||||
assert_eq!(all[0].quality_score, Some(7.5));
|
||||
assert_eq!(all[0].scorer.as_deref(), Some("manual"));
|
||||
// Unknown id updates nothing.
|
||||
assert_eq!(s.set_score(999_999, 1.0, "manual").unwrap(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn swap_costs_pivots_swap_rows() {
|
||||
let s = Store::open_in_memory().unwrap();
|
||||
for (unload, load) in [(300u64, 24000u64), (340, 26000)] {
|
||||
let mut r = rec("beast", "sha", "m", "swap", true);
|
||||
r.swap_unload_ms = Some(unload);
|
||||
r.swap_load_ms = Some(load);
|
||||
r.ttft_s = Some(2.5); // cold first-request
|
||||
r.total_s = Some(5.0);
|
||||
s.insert_run(&r).unwrap();
|
||||
}
|
||||
// A non-swap row must be ignored by swap_costs.
|
||||
s.insert_run(&rec("beast", "sha", "m", "chat:128", true))
|
||||
.unwrap();
|
||||
let costs = s.swap_costs().unwrap();
|
||||
assert_eq!(costs.len(), 1);
|
||||
let c = &costs[0];
|
||||
assert_eq!(c.unload_ms_median, Some(320.0)); // median(300,340)
|
||||
assert_eq!(c.load_ms_median, Some(25000.0)); // median(24000,26000)
|
||||
assert_eq!(c.cold_ttft_s_median, Some(2.5));
|
||||
assert_eq!(c.samples, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scaling_pivots_chat_cells_and_computes_flatness() {
|
||||
let s = Store::open_in_memory().unwrap();
|
||||
// Two context points for one model: decode tok/s 50 @128, 45 @4096.
|
||||
let mut small = rec("beast", "sha", "m", "chat:128", true);
|
||||
small.prompt_size_approx = 128;
|
||||
small.decode_tps = Some(50.0);
|
||||
small.prefill_ms = Some(100);
|
||||
small.prefill_tokens = Some(128);
|
||||
s.insert_run(&small).unwrap();
|
||||
let mut big = rec("beast", "sha", "m", "chat:4096", true);
|
||||
big.prompt_size_approx = 4096;
|
||||
big.decode_tps = Some(45.0);
|
||||
big.prefill_ms = Some(1000);
|
||||
big.prefill_tokens = Some(4096);
|
||||
s.insert_run(&big).unwrap();
|
||||
// A concurrency cell must NOT leak into the scaling curve.
|
||||
let mut conc = rec("beast", "sha", "m", "concurrency:8", true);
|
||||
conc.concurrency = Some(8);
|
||||
s.insert_run(&conc).unwrap();
|
||||
|
||||
let curves = s.scaling().unwrap();
|
||||
assert_eq!(curves.len(), 1);
|
||||
let c = &curves[0];
|
||||
assert_eq!(c.points.len(), 2); // only the two chat:<n> points
|
||||
assert_eq!(c.points[0].prompt_size, 128); // ordered ascending
|
||||
assert_eq!(c.points[1].prompt_size, 4096);
|
||||
// flatness = decode@largest / decode@smallest = 45/50 = 0.9
|
||||
assert!((c.decode_flatness.unwrap() - 0.9).abs() < 1e-9);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gpu_total_vram_sums_devices() {
|
||||
let j = r#"[{"vram_total_mb":32000},{"vram_total_mb":32000}]"#;
|
||||
assert_eq!(gpu_total_vram_mb(j), Some(64000));
|
||||
assert_eq!(gpu_total_vram_mb("[]"), None);
|
||||
assert_eq!(gpu_total_vram_mb("not json"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn percentile_nearest_rank() {
|
||||
let vals = || [1.0, 2.0, 3.0, 4.0, 5.0].into_iter();
|
||||
assert_eq!(percentile(vals(), 50.0), Some(3.0));
|
||||
assert_eq!(percentile(vals(), 95.0), Some(5.0));
|
||||
assert_eq!(percentile(vals(), 99.0), Some(5.0));
|
||||
assert_eq!(percentile(std::iter::empty(), 95.0), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn migration_is_idempotent_and_backfills() {
|
||||
// A DB whose `runs` table predates the prefill columns: create the
|
||||
// pre-#85 shape, insert a row, then run ensure_columns twice.
|
||||
let conn = Connection::open_in_memory().unwrap();
|
||||
conn.execute_batch(
|
||||
"CREATE TABLE runs (id INTEGER PRIMARY KEY, ttft_s REAL);
|
||||
INSERT INTO runs (ttft_s) VALUES (0.1);",
|
||||
)
|
||||
.unwrap();
|
||||
for _ in 0..2 {
|
||||
Store::ensure_columns(
|
||||
&conn,
|
||||
"runs",
|
||||
&[("prefill_ms", "INTEGER"), ("decode_ms", "INTEGER")],
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
// Columns now exist and the old row reads them back as NULL.
|
||||
let got: Option<i64> = conn
|
||||
.query_row("SELECT prefill_ms FROM runs", [], |r| r.get(0))
|
||||
.unwrap();
|
||||
assert_eq!(got, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gpu_label_formats() {
|
||||
let two = r#"[{"name":"NVIDIA GeForce RTX 5090"},{"name":"NVIDIA GeForce RTX 5090"}]"#;
|
||||
|
||||
@@ -9,11 +9,11 @@
|
||||
|
||||
use crate::client::TargetClient;
|
||||
use crate::config::{BenchConfig, TargetConfig, TargetKind};
|
||||
use crate::scenario::{RunCtx, build_scenarios};
|
||||
use crate::scenario::{RunCtx, ScenarioMetrics, build_scenarios};
|
||||
use crate::store::{RunRecord, Store};
|
||||
use anyhow::Result;
|
||||
use cortex_core::build_info::BuildInfo;
|
||||
use cortex_core::discovery::DiscoveryResponse;
|
||||
use cortex_core::discovery::{DiscoveryResponse, HealthResponse};
|
||||
use cortex_core::harness::ModelInfo;
|
||||
|
||||
/// helexa-bench's own build version.
|
||||
@@ -38,6 +38,38 @@ pub struct SweepSummary {
|
||||
pub targets_unreachable: usize,
|
||||
}
|
||||
|
||||
/// Node-level GPU telemetry folded from one `/health` snapshot (#87):
|
||||
/// VRAM used summed across the node's devices, and the hottest/busiest
|
||||
/// single device for utilization and temperature.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct HealthAgg {
|
||||
vram_used_mb: u64,
|
||||
gpu_util_pct: u32,
|
||||
gpu_temp_c: u32,
|
||||
}
|
||||
|
||||
/// Cold-load / model-swap timing for one measure_swap cycle (#90).
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct SwapTiming {
|
||||
unload_ms: u64,
|
||||
load_ms: u64,
|
||||
}
|
||||
|
||||
impl HealthAgg {
|
||||
fn from_health(h: &HealthResponse) -> Self {
|
||||
HealthAgg {
|
||||
vram_used_mb: h.devices.iter().map(|d| d.vram_used_mb).sum(),
|
||||
gpu_util_pct: h
|
||||
.devices
|
||||
.iter()
|
||||
.map(|d| d.utilization_pct)
|
||||
.max()
|
||||
.unwrap_or(0),
|
||||
gpu_temp_c: h.devices.iter().map(|d| d.temp_c).max().unwrap_or(0),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Sweeper {
|
||||
cfg: BenchConfig,
|
||||
client: TargetClient,
|
||||
@@ -72,6 +104,105 @@ impl Sweeper {
|
||||
}
|
||||
}
|
||||
|
||||
/// Deliberate cold-load / model-swap cost measurement (#90), invoked by
|
||||
/// the `swap-cost` subcommand — **never** the continuous sweep. For each
|
||||
/// neuron target and each currently-warm model: unload it, time the
|
||||
/// reload, then time a cold first request. This takes the model offline
|
||||
/// for the reload, so it is an explicit operator action (maintenance
|
||||
/// window), recorded under `scenario_id = "swap"`.
|
||||
pub async fn swap_cost_once(&self) -> Result<SweepSummary> {
|
||||
let mut summary = SweepSummary::default();
|
||||
for target in &self.cfg.targets {
|
||||
if target.kind != TargetKind::Neuron {
|
||||
continue; // load/unload is a neuron-native operation
|
||||
}
|
||||
let build = match self.client.fetch_version(target).await {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
summary.targets_unreachable += 1;
|
||||
tracing::warn!(target = %target.name, error = %format!("{e:#}"), "swap: target unreachable");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let discovery = self.client.fetch_discovery(target).await.unwrap_or(None);
|
||||
let models = self.client.warm_models(target).await.unwrap_or_default();
|
||||
for model in &models {
|
||||
match self
|
||||
.measure_swap(target, &build, discovery.as_ref(), model)
|
||||
.await
|
||||
{
|
||||
Ok(()) => summary.measured += 1,
|
||||
Err(e) => {
|
||||
summary.failed += 1;
|
||||
tracing::warn!(target = %target.name, model = %model.id, error = %format!("{e:#}"), "swap: measurement failed");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
/// Unload → timed reload → timed cold first request for one model.
|
||||
async fn measure_swap(
|
||||
&self,
|
||||
target: &TargetConfig,
|
||||
build: &BuildInfo,
|
||||
discovery: Option<&DiscoveryResponse>,
|
||||
model: &ModelInfo,
|
||||
) -> Result<()> {
|
||||
let spec = TargetClient::spec_from_info(model)?;
|
||||
tracing::warn!(target = %target.name, model = %model.id, "swap: unloading (model goes offline until reload)");
|
||||
|
||||
let t0 = std::time::Instant::now();
|
||||
self.client.unload_model(target, &model.id).await?;
|
||||
let unload_ms = t0.elapsed().as_millis() as u64;
|
||||
|
||||
let t1 = std::time::Instant::now();
|
||||
self.client.load_model(target, &spec).await?;
|
||||
let load_ms = t1.elapsed().as_millis() as u64;
|
||||
tracing::info!(target = %target.name, model = %model.id, unload_ms, load_ms, "swap: reloaded; measuring cold first request");
|
||||
|
||||
// Cold first request — caches empty straight after the load.
|
||||
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(),
|
||||
};
|
||||
let cold = crate::scenario::cold_probe(&ctx).await;
|
||||
let swap = SwapTiming { unload_ms, load_ms };
|
||||
let rec = match &cold {
|
||||
Ok(m) => self.build_record(
|
||||
target,
|
||||
build,
|
||||
discovery,
|
||||
model,
|
||||
"swap",
|
||||
0,
|
||||
Ok(m),
|
||||
None,
|
||||
Some(swap),
|
||||
),
|
||||
Err(e) => {
|
||||
let msg = format!("{e:#}");
|
||||
self.build_record(
|
||||
target,
|
||||
build,
|
||||
discovery,
|
||||
model,
|
||||
"swap",
|
||||
0,
|
||||
Err(&msg),
|
||||
None,
|
||||
Some(swap),
|
||||
)
|
||||
}
|
||||
};
|
||||
self.store.insert_run(&rec)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// One full pass over all targets.
|
||||
pub async fn run_once(&self) -> Result<SweepSummary> {
|
||||
let mut summary = SweepSummary::default();
|
||||
@@ -131,7 +262,21 @@ impl Sweeper {
|
||||
}
|
||||
|
||||
for i in 0..need {
|
||||
match scenario.run(&ctx).await {
|
||||
let result = scenario.run(&ctx).await;
|
||||
// Sample GPU telemetry right after the run, while the
|
||||
// model is loaded and decode VRAM is at its recent peak
|
||||
// (#87). neuron's /health is ~5s-cached, so this is a
|
||||
// coarse high-water proxy, not an instantaneous peak — but
|
||||
// it's the headroom signal we can read over the wire. A
|
||||
// flaky /health degrades to None, never a failed run.
|
||||
let health = self
|
||||
.client
|
||||
.fetch_health(target)
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
.map(|h| HealthAgg::from_health(&h));
|
||||
match result {
|
||||
Ok(m) => {
|
||||
let rec = self.build_record(
|
||||
target,
|
||||
@@ -141,6 +286,8 @@ impl Sweeper {
|
||||
scenario.id(),
|
||||
scenario.prompt_size(),
|
||||
Ok(&m),
|
||||
health,
|
||||
None,
|
||||
);
|
||||
self.store.insert_run(&rec)?;
|
||||
summary.measured += 1;
|
||||
@@ -160,6 +307,8 @@ impl Sweeper {
|
||||
scenario.id(),
|
||||
scenario.prompt_size(),
|
||||
Err(&msg),
|
||||
health,
|
||||
None,
|
||||
);
|
||||
self.store.insert_run(&rec)?;
|
||||
summary.failed += 1;
|
||||
@@ -186,19 +335,14 @@ impl Sweeper {
|
||||
scenario_id: &str,
|
||||
prompt_size: u32,
|
||||
result: Result<&crate::scenario::ScenarioMetrics, &str>,
|
||||
health: Option<HealthAgg>,
|
||||
swap: Option<SwapTiming>,
|
||||
) -> 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),
|
||||
let (m, error): (Option<&ScenarioMetrics>, Option<String>) = match result {
|
||||
Ok(m) => (Some(m), None),
|
||||
Err(e) => (None, Some(e.to_string())),
|
||||
};
|
||||
let ok = m.is_some();
|
||||
|
||||
RunRecord {
|
||||
ts: chrono::Utc::now().to_rfc3339(),
|
||||
@@ -230,12 +374,29 @@ impl Sweeper {
|
||||
.unwrap_or_else(|_| "[]".to_string()),
|
||||
scenario_id: scenario_id.to_string(),
|
||||
prompt_size_approx: prompt_size,
|
||||
prompt_tokens_actual: prompt_tokens,
|
||||
prompt_tokens_actual: m.and_then(|m| m.prompt_tokens),
|
||||
max_tokens: self.cfg.scenarios.max_tokens,
|
||||
ttft_s: ttft,
|
||||
decode_tps: decode,
|
||||
total_s: total,
|
||||
completion_tokens: completion,
|
||||
ttft_s: m.map(|m| m.ttft_s),
|
||||
decode_tps: m.and_then(|m| m.decode_tps),
|
||||
total_s: m.map(|m| m.total_s),
|
||||
completion_tokens: m.map(|m| m.completion_tokens),
|
||||
prefill_ms: m.and_then(|m| m.prefill_ms),
|
||||
decode_ms: m.and_then(|m| m.decode_ms),
|
||||
prefill_tokens: m.and_then(|m| m.prefill_tokens),
|
||||
vram_used_mb: health.map(|h| h.vram_used_mb),
|
||||
gpu_util_pct: health.map(|h| h.gpu_util_pct),
|
||||
gpu_temp_c: health.map(|h| h.gpu_temp_c),
|
||||
concurrency: m.and_then(|m| m.concurrency),
|
||||
ttft_p95_s: m.and_then(|m| m.ttft_p95_s),
|
||||
queue_wait_ms: m.and_then(|m| m.queue_wait_ms_median),
|
||||
rejected: m.and_then(|m| m.rejected),
|
||||
swap_unload_ms: swap.map(|s| s.unload_ms),
|
||||
swap_load_ms: swap.map(|s| s.load_ms),
|
||||
// Capability artifact (#91); score/scorer are attached later by
|
||||
// the `score` subcommand or a future LLM-judge.
|
||||
artifact: m.and_then(|m| m.artifact.clone()),
|
||||
quality_score: None,
|
||||
scorer: None,
|
||||
ok,
|
||||
error,
|
||||
}
|
||||
|
||||
@@ -46,6 +46,21 @@ fn rec(
|
||||
decode_tps: if ok { Some(30.0) } else { None },
|
||||
total_s: if ok { Some(2.0) } else { None },
|
||||
completion_tokens: if ok { Some(60) } else { None },
|
||||
prefill_ms: if ok { Some(150) } else { None },
|
||||
decode_ms: if ok { Some(1800) } else { None },
|
||||
prefill_tokens: if ok { Some(130) } else { None },
|
||||
vram_used_mb: if ok { Some(42000) } else { None },
|
||||
gpu_util_pct: if ok { Some(85) } else { None },
|
||||
gpu_temp_c: if ok { Some(63) } else { None },
|
||||
concurrency: None,
|
||||
ttft_p95_s: None,
|
||||
queue_wait_ms: None,
|
||||
rejected: None,
|
||||
swap_unload_ms: None,
|
||||
swap_load_ms: None,
|
||||
artifact: None,
|
||||
quality_score: None,
|
||||
scorer: None,
|
||||
ok,
|
||||
error: if ok { None } else { Some("boom".into()) },
|
||||
}
|
||||
|
||||
@@ -89,6 +89,9 @@ fn config_for(endpoint: String, db_path: String) -> BenchConfig {
|
||||
scenarios: ScenarioConfig {
|
||||
prompt_sizes: vec![128], // single scenario keeps assertions simple
|
||||
max_tokens: 16,
|
||||
concurrency_levels: Vec::new(),
|
||||
concurrency_prompt_tokens: 512,
|
||||
capability_probes: Vec::new(),
|
||||
},
|
||||
api: Default::default(),
|
||||
targets: vec![TargetConfig {
|
||||
|
||||
@@ -55,6 +55,12 @@ pub fn aggregate_models(topology: &HashMap<String, CortexTopology>) -> Vec<Corte
|
||||
|
||||
let mut out: Vec<CortexModelEntry> = merged.into_values().collect();
|
||||
out.sort_by(|a, b| a.id.cmp(&b.id));
|
||||
// Re-derive the flat ecosystem fields (#78) from the merged (tightest)
|
||||
// limit — the values deserialized from each cortex are per-operator and
|
||||
// may not match the federation-wide merge.
|
||||
for e in &mut out {
|
||||
e.sync_flat_limit();
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
@@ -74,6 +80,10 @@ fn router_entry(cortex: &str, e: &CortexModelEntry) -> CortexModelEntry {
|
||||
cost: e.cost.clone(),
|
||||
tool_call: e.tool_call,
|
||||
reasoning: e.reasoning,
|
||||
// Derived from `limit` by the final sync pass in aggregate_models.
|
||||
max_model_len: None,
|
||||
max_input_tokens: None,
|
||||
max_output_tokens: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -151,6 +161,9 @@ mod tests {
|
||||
cost: None,
|
||||
tool_call: false,
|
||||
reasoning: false,
|
||||
max_model_len: None,
|
||||
max_input_tokens: None,
|
||||
max_output_tokens: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -239,5 +252,9 @@ mod tests {
|
||||
assert_eq!(out.len(), 1);
|
||||
assert_eq!(out[0].limit.as_ref().unwrap().context, 16_384);
|
||||
assert_eq!(out[0].cost.as_ref().unwrap().input, 0.20);
|
||||
// Flat #78 fields re-derived from the merged (tightest) limit.
|
||||
assert_eq!(out[0].max_model_len, Some(16_384));
|
||||
assert_eq!(out[0].max_input_tokens, None);
|
||||
assert_eq!(out[0].max_output_tokens, Some(4096));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -39,6 +39,9 @@ fn model_entry(loaded: bool, feasible: bool) -> CortexModelEntry {
|
||||
cost: None,
|
||||
tool_call: false,
|
||||
reasoning: false,
|
||||
max_model_len: None,
|
||||
max_input_tokens: None,
|
||||
max_output_tokens: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -22,7 +22,7 @@ use axum::http::{StatusCode, header};
|
||||
use axum::middleware::Next;
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::routing::post;
|
||||
use axum::{Json, Router};
|
||||
use axum::{Extension, Json, Router};
|
||||
use cortex_core::error_envelope::OpenAiError;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use subtle::ConstantTimeEq;
|
||||
@@ -41,6 +41,7 @@ pub fn router(state: &AppState) -> Router<AppState> {
|
||||
.route("/authz/v1/settle", post(settle))
|
||||
.route("/authz/v1/release", post(release))
|
||||
.route("/authz/v1/snapshot", post(snapshot))
|
||||
.route("/authz/v1/served-usage", post(served_usage))
|
||||
.layer(axum::middleware::from_fn_with_state(
|
||||
state.clone(),
|
||||
client_auth,
|
||||
@@ -274,6 +275,61 @@ async fn snapshot(State(state): State<AppState>, Json(req): Json<SnapshotReq>) -
|
||||
}
|
||||
}
|
||||
|
||||
// ── served-usage report (#58) ───────────────────────────────────────
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ServedUsageReport {
|
||||
rows: Vec<ServedUsageRow>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ServedUsageRow {
|
||||
account_id: String,
|
||||
key_id: String,
|
||||
period: String, // YYYY-MM-DD
|
||||
served_tokens: i64,
|
||||
}
|
||||
|
||||
/// `POST /authz/v1/served-usage` — a cortex reports the absolute served-token
|
||||
/// counters it has accrued for the current period. Upsert is monotonic
|
||||
/// (`GREATEST`) so re-sends and races are idempotent and never regress.
|
||||
/// `operator_id` comes from the validated client bearer (request extension).
|
||||
async fn served_usage(
|
||||
State(state): State<AppState>,
|
||||
Extension(operator): Extension<OperatorId>,
|
||||
Json(req): Json<ServedUsageReport>,
|
||||
) -> Response {
|
||||
for row in &req.rows {
|
||||
let (Ok(account_id), Ok(key_id)) = (
|
||||
Uuid::parse_str(&row.account_id),
|
||||
Uuid::parse_str(&row.key_id),
|
||||
) else {
|
||||
continue; // skip malformed ids rather than fail the whole batch
|
||||
};
|
||||
let Ok(period) = chrono::NaiveDate::parse_from_str(&row.period, "%Y-%m-%d") else {
|
||||
continue;
|
||||
};
|
||||
let res = sqlx::query(
|
||||
"INSERT INTO served_usage (operator_id, account_id, key_id, period, served_tokens) \
|
||||
VALUES ($1, $2, $3, $4, $5) \
|
||||
ON CONFLICT (operator_id, account_id, key_id, period) \
|
||||
DO UPDATE SET served_tokens = GREATEST(served_usage.served_tokens, EXCLUDED.served_tokens)",
|
||||
)
|
||||
.bind(&operator.0)
|
||||
.bind(account_id)
|
||||
.bind(key_id)
|
||||
.bind(period)
|
||||
.bind(row.served_tokens.max(0))
|
||||
.execute(&state.pool)
|
||||
.await;
|
||||
if let Err(e) = res {
|
||||
tracing::error!(error = %e, "served-usage upsert failed");
|
||||
return envelope_response(OpenAiError::service_unavailable("authority error", Some(5)));
|
||||
}
|
||||
}
|
||||
StatusCode::NO_CONTENT.into_response()
|
||||
}
|
||||
|
||||
fn bad_request(msg: &str) -> Response {
|
||||
envelope_response(OpenAiError::new(
|
||||
400,
|
||||
|
||||
@@ -20,7 +20,9 @@ pub mod email;
|
||||
pub mod error;
|
||||
pub mod handlers;
|
||||
pub mod ledger;
|
||||
pub mod reconcile;
|
||||
pub mod state;
|
||||
pub mod topup;
|
||||
pub mod web;
|
||||
|
||||
use anyhow::Result;
|
||||
|
||||
@@ -20,6 +20,28 @@ enum Commands {
|
||||
#[arg(short, long, default_value = "helexa-upstream.toml")]
|
||||
config: String,
|
||||
},
|
||||
/// Mint single-use top-up codes and print them (one per line). The raw
|
||||
/// codes are shown only here — only their hash is stored. (The future
|
||||
/// faucet bot calls the same path.)
|
||||
Mint {
|
||||
#[arg(short, long, default_value = "helexa-upstream.toml")]
|
||||
config: String,
|
||||
/// Tokens each code grants.
|
||||
#[arg(long)]
|
||||
value: i64,
|
||||
/// How many codes to mint.
|
||||
#[arg(long, default_value_t = 1)]
|
||||
count: u32,
|
||||
/// Optional human label (e.g. "small", "beta-launch").
|
||||
#[arg(long)]
|
||||
denomination: Option<String>,
|
||||
},
|
||||
/// Roll up not-yet-reconciled served usage per operator/period (#58),
|
||||
/// stamp it reconciled, and print the totals. Payout is out of scope.
|
||||
Reconcile {
|
||||
#[arg(short, long, default_value = "helexa-upstream.toml")]
|
||||
config: String,
|
||||
},
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
@@ -40,6 +62,37 @@ async fn main() -> Result<()> {
|
||||
tracing::info!(listen = %cfg.server.listen, "starting helexa-upstream");
|
||||
helexa_upstream::run(cfg).await?;
|
||||
}
|
||||
Commands::Mint {
|
||||
config,
|
||||
value,
|
||||
count,
|
||||
denomination,
|
||||
} => {
|
||||
let cfg = UpstreamConfig::load(&config)
|
||||
.map_err(|e| anyhow::anyhow!("failed to load config from '{config}': {e}"))?;
|
||||
let pool =
|
||||
helexa_upstream::db::connect_and_migrate(&cfg.db.url, cfg.db.max_connections)
|
||||
.await?;
|
||||
let codes =
|
||||
helexa_upstream::topup::mint(&pool, value, count, denomination.as_deref()).await?;
|
||||
// Raw codes to stdout (one per line) for the operator to distribute;
|
||||
// logs/diagnostics go to stderr via tracing.
|
||||
for code in codes {
|
||||
println!("{code}");
|
||||
}
|
||||
}
|
||||
Commands::Reconcile { config } => {
|
||||
let cfg = UpstreamConfig::load(&config)
|
||||
.map_err(|e| anyhow::anyhow!("failed to load config from '{config}': {e}"))?;
|
||||
let pool =
|
||||
helexa_upstream::db::connect_and_migrate(&cfg.db.url, cfg.db.max_connections)
|
||||
.await?;
|
||||
let rollup = helexa_upstream::reconcile::reconcile(&pool).await?;
|
||||
for r in &rollup {
|
||||
println!("{}\t{}\t{}", r.operator_id, r.period, r.total_served_tokens);
|
||||
}
|
||||
tracing::info!(operators_periods = rollup.len(), "reconciliation complete");
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
|
||||
43
crates/helexa-upstream/src/reconcile.rs
Normal file
43
crates/helexa-upstream/src/reconcile.rs
Normal file
@@ -0,0 +1,43 @@
|
||||
//! Reconciliation rollup (#58): aggregate the served-usage ledger per
|
||||
//! operator and period for operator compensation, stamping rows
|
||||
//! `reconciled_at` so each window is settled once. The payout mechanism
|
||||
//! itself is out of scope — this produces the authoritative per-operator
|
||||
//! totals a settlement process consumes.
|
||||
|
||||
use sqlx::Row;
|
||||
use sqlx::postgres::PgPool;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct RollupRow {
|
||||
pub operator_id: String,
|
||||
pub period: chrono::NaiveDate,
|
||||
pub total_served_tokens: i64,
|
||||
}
|
||||
|
||||
/// Roll up all not-yet-reconciled served-usage into per-(operator, period)
|
||||
/// totals, then stamp those rows `reconciled_at`. Returns the rollup.
|
||||
/// Idempotent: a second run finds nothing unreconciled and returns empty.
|
||||
pub async fn reconcile(pool: &PgPool) -> Result<Vec<RollupRow>, sqlx::Error> {
|
||||
let mut tx = pool.begin().await?;
|
||||
let rows = sqlx::query(
|
||||
// SUM(bigint) is numeric in Postgres — cast back to bigint for i64.
|
||||
"SELECT operator_id, period, SUM(served_tokens)::bigint AS total \
|
||||
FROM served_usage WHERE reconciled_at IS NULL \
|
||||
GROUP BY operator_id, period ORDER BY operator_id, period",
|
||||
)
|
||||
.fetch_all(&mut *tx)
|
||||
.await?;
|
||||
let rollup: Vec<RollupRow> = rows
|
||||
.iter()
|
||||
.map(|r| RollupRow {
|
||||
operator_id: r.get("operator_id"),
|
||||
period: r.get("period"),
|
||||
total_served_tokens: r.get::<i64, _>("total"),
|
||||
})
|
||||
.collect();
|
||||
sqlx::query("UPDATE served_usage SET reconciled_at = now() WHERE reconciled_at IS NULL")
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
tx.commit().await?;
|
||||
Ok(rollup)
|
||||
}
|
||||
82
crates/helexa-upstream/src/topup.rs
Normal file
82
crates/helexa-upstream/src/topup.rs
Normal file
@@ -0,0 +1,82 @@
|
||||
//! Single-use top-up codes (#B5) — the second half of the hybrid allocation
|
||||
//! model. Each code grants `value` tokens to the account that redeems it,
|
||||
//! raising `accounts.allocation_total`. Minting codes is operator/CLI side
|
||||
//! (the future faucet bot calls the same `mint` path); redemption is a
|
||||
//! `/web/v1` action.
|
||||
//!
|
||||
//! Security: only `sha256(code)` is stored. Redemption is **timing-safe and
|
||||
//! single-use** — a conditional `UPDATE … WHERE redeemed_by IS NULL` does
|
||||
//! the claim atomically (concurrent double-redeem → exactly one winner), and
|
||||
//! a not-found code and an already-redeemed code return the **same** generic
|
||||
//! failure with the same code path (no oracle for "valid but spent").
|
||||
|
||||
use crate::crypto::{random_token, sha256};
|
||||
use sqlx::Row;
|
||||
use sqlx::postgres::PgPool;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum TopUpError {
|
||||
/// Code unknown OR already redeemed — deliberately indistinguishable.
|
||||
#[error("invalid or already-redeemed code")]
|
||||
Invalid,
|
||||
#[error(transparent)]
|
||||
Db(#[from] sqlx::Error),
|
||||
}
|
||||
|
||||
/// Redeem `raw_code` for `account_id`, raising the account's
|
||||
/// `allocation_total` by the code's value. Returns the new total.
|
||||
pub async fn redeem(pool: &PgPool, account_id: Uuid, raw_code: &str) -> Result<i64, TopUpError> {
|
||||
let mut tx = pool.begin().await?;
|
||||
// Atomic single-use claim. `redeemed_by IS NULL` is the guarantee: under
|
||||
// concurrent redemption exactly one UPDATE touches the row.
|
||||
let claimed = sqlx::query(
|
||||
"UPDATE top_up_codes SET redeemed_by = $1, redeemed_at = now() \
|
||||
WHERE code_hash = $2 AND redeemed_by IS NULL RETURNING value",
|
||||
)
|
||||
.bind(account_id)
|
||||
.bind(sha256(raw_code))
|
||||
.fetch_optional(&mut *tx)
|
||||
.await?;
|
||||
let Some(row) = claimed else {
|
||||
// Not found or already redeemed — same path, same error.
|
||||
return Err(TopUpError::Invalid);
|
||||
};
|
||||
let value: i64 = row.get("value");
|
||||
let new_total: i64 = sqlx::query(
|
||||
"UPDATE accounts SET allocation_total = allocation_total + $1 WHERE id = $2 \
|
||||
RETURNING allocation_total",
|
||||
)
|
||||
.bind(value)
|
||||
.bind(account_id)
|
||||
.fetch_one(&mut *tx)
|
||||
.await?
|
||||
.get("allocation_total");
|
||||
tx.commit().await?;
|
||||
Ok(new_total)
|
||||
}
|
||||
|
||||
/// Mint `count` codes each worth `value` tokens, optionally tagged with a
|
||||
/// `denomination` label. Returns the raw codes (shown once — only their
|
||||
/// hash is stored). The CLI prints these; the future faucet bot calls this.
|
||||
pub async fn mint(
|
||||
pool: &PgPool,
|
||||
value: i64,
|
||||
count: u32,
|
||||
denomination: Option<&str>,
|
||||
) -> Result<Vec<String>, sqlx::Error> {
|
||||
let mut codes = Vec::with_capacity(count as usize);
|
||||
for _ in 0..count {
|
||||
let raw = format!("helexa-topup-{}", random_token());
|
||||
sqlx::query(
|
||||
"INSERT INTO top_up_codes (code_hash, value, denomination) VALUES ($1, $2, $3)",
|
||||
)
|
||||
.bind(sha256(&raw))
|
||||
.bind(value)
|
||||
.bind(denomination)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
codes.push(raw);
|
||||
}
|
||||
Ok(codes)
|
||||
}
|
||||
@@ -35,6 +35,7 @@ pub fn router(state: &AppState) -> Router<AppState> {
|
||||
"/web/v1/keys/{id}/limit",
|
||||
axum::routing::patch(update_key_limit),
|
||||
)
|
||||
.route("/web/v1/redeem", post(redeem))
|
||||
.layer(axum::middleware::from_fn_with_state(
|
||||
state.clone(),
|
||||
require_session,
|
||||
@@ -565,3 +566,29 @@ async fn update_key_limit(
|
||||
}
|
||||
Ok(StatusCode::NO_CONTENT.into_response())
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct RedeemReq {
|
||||
code: String,
|
||||
}
|
||||
|
||||
/// `POST /web/v1/redeem` — redeem a single-use top-up code, raising the
|
||||
/// account's allocation. Returns the new total. Generic 400 for an invalid
|
||||
/// or already-redeemed code (no oracle).
|
||||
async fn redeem(
|
||||
State(state): State<AppState>,
|
||||
Extension(user): Extension<AuthUser>,
|
||||
Json(req): Json<RedeemReq>,
|
||||
) -> WebResult<Response> {
|
||||
let acct = account_id_for(&state, user.0).await?;
|
||||
match crate::topup::redeem(&state.pool, acct, &req.code).await {
|
||||
Ok(new_total) => Ok(Json(json!({ "allocation_total": new_total })).into_response()),
|
||||
Err(crate::topup::TopUpError::Invalid) => {
|
||||
Err(WebError::BadRequest("invalid or already-redeemed code"))
|
||||
}
|
||||
Err(crate::topup::TopUpError::Db(e)) => {
|
||||
tracing::error!(error = %e, "redeem db error");
|
||||
Err(WebError::Internal)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
116
crates/helexa-upstream/tests/served_usage_pg.rs
Normal file
116
crates/helexa-upstream/tests/served_usage_pg.rs
Normal file
@@ -0,0 +1,116 @@
|
||||
//! Integration test for the served-usage report (#58): the idempotent,
|
||||
//! monotonic upsert and the reconcile rollup. Gated on
|
||||
//! UPSTREAM_TEST_DATABASE_URL (skips cleanly when unset).
|
||||
|
||||
use helexa_upstream::config::{ClientToken, UpstreamConfig};
|
||||
use helexa_upstream::db::connect_and_migrate;
|
||||
use helexa_upstream::email::EmailSender;
|
||||
use helexa_upstream::reconcile::reconcile;
|
||||
use helexa_upstream::state::AppState;
|
||||
use serde_json::{Value, json};
|
||||
use sqlx::Row;
|
||||
use sqlx::postgres::PgPool;
|
||||
use uuid::Uuid;
|
||||
|
||||
const CLIENT_TOKEN: &str = "su-test-token";
|
||||
const OPERATOR: &str = "op-su-test";
|
||||
|
||||
async fn spawn_or_skip(test: &str) -> Option<(String, PgPool)> {
|
||||
let Ok(url) = std::env::var("UPSTREAM_TEST_DATABASE_URL") else {
|
||||
eprintln!("skipping {test}: UPSTREAM_TEST_DATABASE_URL not set");
|
||||
return None;
|
||||
};
|
||||
let pool = connect_and_migrate(&url, 16).await.expect("migrate");
|
||||
let mut config = UpstreamConfig {
|
||||
server: Default::default(),
|
||||
db: helexa_upstream::config::DbSettings {
|
||||
url,
|
||||
max_connections: 16,
|
||||
},
|
||||
grant: Default::default(),
|
||||
abuse: Default::default(),
|
||||
client_auth: Default::default(),
|
||||
authz: Default::default(),
|
||||
auth: Default::default(),
|
||||
email: Default::default(),
|
||||
};
|
||||
config.client_auth.tokens.push(ClientToken {
|
||||
token: CLIENT_TOKEN.into(),
|
||||
operator_id: OPERATOR.into(),
|
||||
});
|
||||
let email = EmailSender::from_config(&config.email).unwrap();
|
||||
let state = AppState::new(pool.clone(), config, email);
|
||||
let app = helexa_upstream::build_app(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();
|
||||
});
|
||||
Some((format!("http://{addr}"), pool))
|
||||
}
|
||||
|
||||
async fn report(base: &str, rows: Value) -> u16 {
|
||||
reqwest::Client::new()
|
||||
.post(format!("{base}/authz/v1/served-usage"))
|
||||
.bearer_auth(CLIENT_TOKEN)
|
||||
.json(&json!({ "rows": rows }))
|
||||
.send()
|
||||
.await
|
||||
.unwrap()
|
||||
.status()
|
||||
.as_u16()
|
||||
}
|
||||
|
||||
async fn stored(pool: &PgPool, account: Uuid, key: Uuid) -> i64 {
|
||||
sqlx::query(
|
||||
"SELECT served_tokens FROM served_usage WHERE operator_id = $1 AND account_id = $2 AND key_id = $3",
|
||||
)
|
||||
.bind(OPERATOR)
|
||||
.bind(account)
|
||||
.bind(key)
|
||||
.fetch_one(pool)
|
||||
.await
|
||||
.unwrap()
|
||||
.get("served_tokens")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn served_usage_upsert_is_monotonic_and_reconciles() {
|
||||
let Some((base, pool)) = spawn_or_skip("served_usage_upsert_is_monotonic_and_reconciles").await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let account = Uuid::new_v4();
|
||||
let key = Uuid::new_v4();
|
||||
let period = "2026-06-23";
|
||||
let row = |n: i64| json!([{"account_id": account, "key_id": key, "period": period, "served_tokens": n}]);
|
||||
|
||||
// First report.
|
||||
assert_eq!(report(&base, row(100)).await, 204);
|
||||
assert_eq!(stored(&pool, account, key).await, 100);
|
||||
|
||||
// Re-send a higher absolute value → advances.
|
||||
assert_eq!(report(&base, row(250)).await, 204);
|
||||
assert_eq!(stored(&pool, account, key).await, 250);
|
||||
|
||||
// A lower value (e.g. a restarted cortex) must NOT regress (GREATEST).
|
||||
assert_eq!(report(&base, row(50)).await, 204);
|
||||
assert_eq!(stored(&pool, account, key).await, 250);
|
||||
|
||||
// Re-sending the same value is idempotent.
|
||||
assert_eq!(report(&base, row(250)).await, 204);
|
||||
assert_eq!(stored(&pool, account, key).await, 250);
|
||||
|
||||
// Reconcile rolls it up and stamps reconciled_at; a second run is empty.
|
||||
let rollup = reconcile(&pool).await.unwrap();
|
||||
let mine = rollup
|
||||
.iter()
|
||||
.find(|r| r.operator_id == OPERATOR)
|
||||
.expect("operator in rollup");
|
||||
assert!(mine.total_served_tokens >= 250);
|
||||
let again = reconcile(&pool).await.unwrap();
|
||||
assert!(
|
||||
again.iter().all(|r| r.operator_id != OPERATOR),
|
||||
"already reconciled"
|
||||
);
|
||||
}
|
||||
@@ -294,3 +294,132 @@ async fn fingerprint_abuse_silently_deactivates_all_no_clue() {
|
||||
"deactivated account's key looks like any invalid key"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn topup_redeem_raises_allocation_single_use() {
|
||||
let Some((base, pool)) = spawn_or_skip("topup_redeem_raises_allocation_single_use").await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let email = unique_email();
|
||||
post(
|
||||
format!("{base}/web/v1/register"),
|
||||
json!({"email": email, "password": "password123"}),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
pool.execute(
|
||||
sqlx::query("UPDATE users SET email_verified = true WHERE email = $1").bind(&email),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let token = post(
|
||||
format!("{base}/web/v1/login"),
|
||||
json!({"email": email, "password": "password123"}),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.json::<Value>()
|
||||
.await
|
||||
.unwrap()["token"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.to_string();
|
||||
|
||||
// Mint a code worth 500_000 (mint path used by the CLI/faucet).
|
||||
let codes = helexa_upstream::topup::mint(&pool, 500_000, 1, Some("test"))
|
||||
.await
|
||||
.unwrap();
|
||||
let code = &codes[0];
|
||||
|
||||
// Redeem → allocation_total rises from the 1_000_000 free grant.
|
||||
let r = post(
|
||||
format!("{base}/web/v1/redeem"),
|
||||
json!({"code": code}),
|
||||
Some(&token),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(r.status(), 200);
|
||||
assert_eq!(
|
||||
r.json::<Value>().await.unwrap()["allocation_total"],
|
||||
1_500_000
|
||||
);
|
||||
|
||||
// Single-use: a second redemption fails generically (no oracle).
|
||||
let r = post(
|
||||
format!("{base}/web/v1/redeem"),
|
||||
json!({"code": code}),
|
||||
Some(&token),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(r.status(), 400);
|
||||
|
||||
// Unknown code: same generic 400.
|
||||
let r = post(
|
||||
format!("{base}/web/v1/redeem"),
|
||||
json!({"code": "helexa-topup-does-not-exist"}),
|
||||
Some(&token),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(r.status(), 400);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn topup_concurrent_double_redeem_one_winner() {
|
||||
let Some((base, pool)) = spawn_or_skip("topup_concurrent_double_redeem_one_winner").await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
// Two verified accounts.
|
||||
let mut tokens = Vec::new();
|
||||
for _ in 0..2 {
|
||||
let email = unique_email();
|
||||
post(
|
||||
format!("{base}/web/v1/register"),
|
||||
json!({"email": email, "password": "password123"}),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
pool.execute(
|
||||
sqlx::query("UPDATE users SET email_verified = true WHERE email = $1").bind(&email),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let t = post(
|
||||
format!("{base}/web/v1/login"),
|
||||
json!({"email": email, "password": "password123"}),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.json::<Value>()
|
||||
.await
|
||||
.unwrap()["token"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.to_string();
|
||||
tokens.push(t);
|
||||
}
|
||||
let code = helexa_upstream::topup::mint(&pool, 100, 1, None)
|
||||
.await
|
||||
.unwrap()
|
||||
.remove(0);
|
||||
|
||||
// Both accounts race to redeem the same code; exactly one wins.
|
||||
let (a, b) = tokio::join!(
|
||||
post(
|
||||
format!("{base}/web/v1/redeem"),
|
||||
json!({"code": code}),
|
||||
Some(&tokens[0])
|
||||
),
|
||||
post(
|
||||
format!("{base}/web/v1/redeem"),
|
||||
json!({"code": code}),
|
||||
Some(&tokens[1])
|
||||
),
|
||||
);
|
||||
let wins = [a.status(), b.status()]
|
||||
.iter()
|
||||
.filter(|s| s.as_u16() == 200)
|
||||
.count();
|
||||
assert_eq!(wins, 1, "exactly one redemption wins the single-use code");
|
||||
}
|
||||
|
||||
@@ -38,6 +38,7 @@ cudnn = [
|
||||
flash-attn = [
|
||||
"cuda",
|
||||
"candle-transformers/flash-attn",
|
||||
"dep:candle-flash-attn",
|
||||
]
|
||||
# Reserved for GPU-only integration tests in later stages.
|
||||
cuda-integration = ["cuda"]
|
||||
@@ -71,6 +72,10 @@ rayon = "1"
|
||||
candle-core = "0.10.2"
|
||||
candle-nn = "0.10.2"
|
||||
candle-transformers = "0.10.2"
|
||||
# Direct dependency so `candle_flash_attn::flash_attn` is nameable in
|
||||
# the attention core (#95); the candle-transformers feature alone only
|
||||
# enables flash inside ITS models.
|
||||
candle-flash-attn = { version = "0.10.2", optional = true }
|
||||
# Direct dep on cudarc (matching candle's transitive version) so the
|
||||
# TP worker pool can call cudarc::nccl::{Comm, Id} directly. Gated on
|
||||
# the `cuda` feature; same toolchain requirement as candle's CUDA path.
|
||||
|
||||
@@ -16,7 +16,7 @@ use cortex_core::discovery::{DiscoveryResponse, HealthResponse};
|
||||
use cortex_core::entitlements::{HEADER_ACCOUNT_ID, HEADER_KEY_ID};
|
||||
use cortex_core::harness::ModelSpec;
|
||||
use cortex_core::openai::{ChatCompletionRequest, MessageContent};
|
||||
use cortex_core::responses::{ResponsesRequest, ResponsesUsage};
|
||||
use cortex_core::responses::{OutputTokensDetails, ResponsesRequest, ResponsesUsage};
|
||||
use futures::stream::{self, StreamExt};
|
||||
use serde_json::{Value, json};
|
||||
use std::convert::Infallible;
|
||||
@@ -418,8 +418,14 @@ async fn responses(
|
||||
input_tokens: u.prompt_tokens,
|
||||
output_tokens: u.completion_tokens,
|
||||
total_tokens: u.prompt_tokens + u.completion_tokens,
|
||||
// Non-streaming reasoning accounting deferred (#64).
|
||||
output_tokens_details: None,
|
||||
// Carry the reasoning sub-count through from the chat
|
||||
// usage — the non-streaming path now splits off the
|
||||
// `<think>` span and counts it (see `split_off_reasoning`).
|
||||
output_tokens_details: u.completion_tokens_details.as_ref().map(|d| {
|
||||
OutputTokensDetails {
|
||||
reasoning_tokens: d.reasoning_tokens,
|
||||
}
|
||||
}),
|
||||
input_tokens_details: None,
|
||||
});
|
||||
let meta = openai_responses::ResponseMeta {
|
||||
|
||||
@@ -94,13 +94,24 @@ impl AdmissionController {
|
||||
/// overall queue is full or the principal is over its fair-share cap —
|
||||
/// then waits up to `max_wait` for an in-flight slot. The returned permit
|
||||
/// must be held for the request's lifetime; dropping it frees the slots.
|
||||
///
|
||||
/// CANCELLATION SAFETY: the semaphore wait below is where a client
|
||||
/// disconnect lands — axum drops the request future mid-await. The
|
||||
/// reservation therefore lives in a RAII [`PendingReservation`] taken
|
||||
/// BEFORE the await: if this future is dropped while queued, the
|
||||
/// guard's Drop rolls the counts back. (The original version
|
||||
/// incremented raw counters and only decremented on the timeout
|
||||
/// branch — every abandoned wait leaked a `pending` + per-principal
|
||||
/// slot, ratcheting the model into a permanent instant-429 state
|
||||
/// under client retry storms. Observed live 2026-07-02:
|
||||
/// `queue_depth: 1` pinned on an idle model.)
|
||||
pub async fn enter(
|
||||
&self,
|
||||
principal: Option<&str>,
|
||||
) -> Result<AdmissionPermit, AdmissionRejection> {
|
||||
// Decision + reservation under one brief lock so concurrent callers
|
||||
// can't both slip past the thresholds. No await is held here.
|
||||
{
|
||||
let reservation = {
|
||||
let mut st = self.state.lock().expect("admission state poisoned");
|
||||
if st.pending >= self.max_pending {
|
||||
return Err(AdmissionRejection::QueueFull {
|
||||
@@ -119,31 +130,25 @@ impl AdmissionController {
|
||||
if let Some(p) = principal {
|
||||
*st.per_principal.entry(p.to_string()).or_insert(0) += 1;
|
||||
}
|
||||
}
|
||||
PendingReservation {
|
||||
state: Arc::clone(&self.state),
|
||||
principal: principal.map(str::to_string),
|
||||
}
|
||||
};
|
||||
|
||||
match tokio::time::timeout(self.max_wait, Arc::clone(&self.slots).acquire_owned()).await {
|
||||
Ok(Ok(permit)) => Ok(AdmissionPermit {
|
||||
_permit: permit,
|
||||
state: Arc::clone(&self.state),
|
||||
principal: principal.map(str::to_string),
|
||||
_reservation: reservation,
|
||||
}),
|
||||
// Semaphore is never closed; treat a closed/elapsed wait the
|
||||
// same. `reservation` drops here, rolling back the counts.
|
||||
Ok(Err(_)) | Err(_) => Err(AdmissionRejection::Timeout {
|
||||
retry_after_secs: self.retry_hint(self.max_pending),
|
||||
}),
|
||||
// Semaphore is never closed; treat a closed/elapsed wait the same.
|
||||
Ok(Err(_)) | Err(_) => {
|
||||
self.release(principal);
|
||||
Err(AdmissionRejection::Timeout {
|
||||
retry_after_secs: self.retry_hint(self.max_pending),
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Roll back a reserved-but-not-admitted slot (wait timed out).
|
||||
fn release(&self, principal: Option<&str>) {
|
||||
let mut st = self.state.lock().expect("admission state poisoned");
|
||||
st.pending = st.pending.saturating_sub(1);
|
||||
decrement_principal(&mut st.per_principal, principal);
|
||||
}
|
||||
|
||||
/// Requests currently running (holding an in-flight slot).
|
||||
pub fn in_flight(&self) -> usize {
|
||||
self.max_in_flight
|
||||
@@ -177,16 +182,17 @@ fn decrement_principal(map: &mut HashMap<String, usize>, principal: Option<&str>
|
||||
}
|
||||
}
|
||||
|
||||
/// Held for a request's lifetime; frees the in-flight + queue slot (and the
|
||||
/// principal's fair-share slot) on drop.
|
||||
/// RAII accounting for one reserved slot (queued or in-flight): decrements
|
||||
/// `pending` and the principal's fair-share count on drop, whichever way
|
||||
/// the reservation ends — admitted-and-finished, wait timeout, or the
|
||||
/// caller's future being dropped mid-queue (client disconnect).
|
||||
#[derive(Debug)]
|
||||
pub struct AdmissionPermit {
|
||||
_permit: OwnedSemaphorePermit,
|
||||
struct PendingReservation {
|
||||
state: Arc<Mutex<AdmissionState>>,
|
||||
principal: Option<String>,
|
||||
}
|
||||
|
||||
impl Drop for AdmissionPermit {
|
||||
impl Drop for PendingReservation {
|
||||
fn drop(&mut self) {
|
||||
let mut st = self.state.lock().expect("admission state poisoned");
|
||||
st.pending = st.pending.saturating_sub(1);
|
||||
@@ -194,6 +200,14 @@ impl Drop for AdmissionPermit {
|
||||
}
|
||||
}
|
||||
|
||||
/// Held for a request's lifetime; frees the in-flight slot (semaphore
|
||||
/// permit) and the queue + fair-share accounting (reservation) on drop.
|
||||
#[derive(Debug)]
|
||||
pub struct AdmissionPermit {
|
||||
_permit: OwnedSemaphorePermit,
|
||||
_reservation: PendingReservation,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -295,4 +309,65 @@ mod tests {
|
||||
drop(_a1);
|
||||
b.await.unwrap().expect("B is served after A releases");
|
||||
}
|
||||
|
||||
/// Regression for the 2026-07-02 retry-storm incident: a client that
|
||||
/// disconnects while QUEUED drops the `enter()` future mid-await.
|
||||
/// The reservation must roll back — the original implementation
|
||||
/// leaked `pending` + the per-principal count on this path, pinning
|
||||
/// the model in a permanent instant-429 state.
|
||||
#[tokio::test]
|
||||
async fn cancelled_queued_waiter_rolls_back_accounting() {
|
||||
let cfg = AdmissionConfig {
|
||||
max_in_flight: 1,
|
||||
max_queue_depth: 2,
|
||||
max_wait_secs: 30,
|
||||
// Cap 3 lets the runner + both waiters coexist; if the two
|
||||
// cancelled waiters leaked their counts, the principal would
|
||||
// sit at 3 == cap and the post-cancel enter below would hit
|
||||
// PrincipalCap instead of queueing.
|
||||
max_per_principal: 3,
|
||||
};
|
||||
let ctrl = Arc::new(AdmissionController::new(&cfg));
|
||||
let running = ctrl.enter(Some("acct/key")).await.expect("admit running");
|
||||
|
||||
// Two waiters from the same principal park in the queue…
|
||||
let mut waiters = Vec::new();
|
||||
for _ in 0..2 {
|
||||
let c = Arc::clone(&ctrl);
|
||||
waiters.push(tokio::spawn(async move {
|
||||
c.enter(Some("acct/key")).await.map(drop)
|
||||
}));
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
assert_eq!(ctrl.queue_depth(), 2);
|
||||
|
||||
// …and both clients vanish (abort = the dropped request future).
|
||||
for w in &waiters {
|
||||
w.abort();
|
||||
}
|
||||
for w in waiters {
|
||||
let _ = w.await;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
|
||||
assert_eq!(
|
||||
ctrl.queue_depth(),
|
||||
0,
|
||||
"cancelled waiters must not leak queue slots"
|
||||
);
|
||||
|
||||
// The principal's fair-share count must also be clean: with the
|
||||
// runner still holding 1 of its cap of 3, a new request from the
|
||||
// same principal queues instead of hitting PrincipalCap (which a
|
||||
// leak of the two cancelled counts would trigger).
|
||||
let c = Arc::clone(&ctrl);
|
||||
let retry = tokio::spawn(async move { c.enter(Some("acct/key")).await.map(drop) });
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
assert_eq!(ctrl.queue_depth(), 1, "post-cancel request queues normally");
|
||||
drop(running);
|
||||
retry
|
||||
.await
|
||||
.unwrap()
|
||||
.expect("post-cancel request is served — no leaked principal count");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -22,6 +22,7 @@ use super::TextConfig;
|
||||
use super::full_attn::Qwen3_5Attention;
|
||||
use super::linear_attn::GatedDeltaNet;
|
||||
use super::mlp::Qwen3_5MLP;
|
||||
use super::moe::Qwen3_5MoeBlock;
|
||||
use super::rmsnorm::Qwen3_5RmsNorm;
|
||||
use super::rope::RotaryEmbedding;
|
||||
use super::snapshot::LayerKvSnapshot;
|
||||
@@ -35,10 +36,27 @@ enum AttentionKind {
|
||||
Linear(GatedDeltaNet),
|
||||
}
|
||||
|
||||
/// The FFN slot: dense SwiGLU (Qwen3.6) or the high-sparsity MoE block
|
||||
/// (qwen3_next 80B-A3B family, #92), selected per layer by
|
||||
/// [`TextConfig::layer_uses_moe`].
|
||||
enum MlpKind {
|
||||
Dense(Qwen3_5MLP),
|
||||
Moe(Qwen3_5MoeBlock),
|
||||
}
|
||||
|
||||
impl Module for MlpKind {
|
||||
fn forward(&self, x: &Tensor) -> candle_core::Result<Tensor> {
|
||||
match self {
|
||||
MlpKind::Dense(mlp) => mlp.forward(x),
|
||||
MlpKind::Moe(moe) => moe.forward(x),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Qwen3_5DecoderLayer {
|
||||
input_layernorm: Qwen3_5RmsNorm,
|
||||
post_attention_layernorm: Qwen3_5RmsNorm,
|
||||
mlp: Qwen3_5MLP,
|
||||
mlp: MlpKind,
|
||||
attention: AttentionKind,
|
||||
}
|
||||
|
||||
@@ -73,7 +91,11 @@ impl Qwen3_5DecoderLayer {
|
||||
),
|
||||
};
|
||||
|
||||
let mlp = Qwen3_5MLP::load(cfg, &vb.pp("mlp"))?;
|
||||
let mlp = if cfg.layer_uses_moe(layer_idx) {
|
||||
MlpKind::Moe(Qwen3_5MoeBlock::load(cfg, &vb.pp("mlp"))?)
|
||||
} else {
|
||||
MlpKind::Dense(Qwen3_5MLP::load(cfg, &vb.pp("mlp"))?)
|
||||
};
|
||||
let input_layernorm =
|
||||
Qwen3_5RmsNorm::load(&vb.pp("input_layernorm"), cfg.hidden_size, cfg.rms_norm_eps)?;
|
||||
let post_attention_layernorm = Qwen3_5RmsNorm::load(
|
||||
|
||||
@@ -29,6 +29,86 @@ use super::TextConfig;
|
||||
use super::rmsnorm::Qwen3_5RmsNorm;
|
||||
use super::rope::RotaryEmbedding;
|
||||
|
||||
/// Runtime kill-switch for the FlashAttention path (#95):
|
||||
/// `NEURON_FLASH_ATTN=0` (or `false`) forces the eager fallback
|
||||
/// without a rebuild — the A/B lever and the rollback if the kernels
|
||||
/// misbehave on some device. Read once.
|
||||
#[cfg(feature = "flash-attn")]
|
||||
fn flash_attn_enabled() -> bool {
|
||||
static ENABLED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
|
||||
*ENABLED.get_or_init(|| {
|
||||
let on = !std::env::var("NEURON_FLASH_ATTN").is_ok_and(|v| v == "0" || v == "false");
|
||||
tracing::info!(enabled = on, "FlashAttention path (#95)");
|
||||
on
|
||||
})
|
||||
}
|
||||
|
||||
/// Attention core shared by the single-GPU and TP full-attention
|
||||
/// layers (#95): `(B, H, L, D)` query and `(B, H_kv, S, D)` key/value
|
||||
/// (post-KV-cache, NOT GQA-repeated) → `(B, H, L, D)` context.
|
||||
///
|
||||
/// With the `flash-attn` feature on a CUDA device in f16/bf16, this
|
||||
/// dispatches to the FlashAttention kernel: GQA is native (no
|
||||
/// repeated-K/V materialisation) and causality is a kernel flag, so
|
||||
/// the O(L²) mask/score tensors never exist. The kernels align the
|
||||
/// causal mask to the BOTTOM-RIGHT when `seqlen_q != seqlen_k`
|
||||
/// (flash-attention v2.1+ semantics), which is exactly what chunked
|
||||
/// prefill continuation needs: a chunk of L new queries against
|
||||
/// `offset + L` cached keys masks correctly.
|
||||
///
|
||||
/// INVARIANT: `attn_mask` is either `None` (decode / single position)
|
||||
/// or the standard causal mask — the only mask the qwen3_5 forward
|
||||
/// constructs. The flash path encodes it as `causal = attn_mask
|
||||
/// .is_some()`; a future non-causal mask must extend this signature,
|
||||
/// not silently pass through.
|
||||
///
|
||||
/// Falls back to the eager matmul→softmax→matmul everywhere else
|
||||
/// (CPU, f32, feature off, or `NEURON_FLASH_ATTN=0`).
|
||||
pub(crate) fn attention_context(
|
||||
q: &Tensor,
|
||||
k: &Tensor,
|
||||
v: &Tensor,
|
||||
attn_mask: Option<&Tensor>,
|
||||
num_kv_groups: usize,
|
||||
scale: f64,
|
||||
) -> candle_core::Result<Tensor> {
|
||||
#[cfg(feature = "flash-attn")]
|
||||
{
|
||||
use candle_core::DType;
|
||||
let dtype = q.dtype();
|
||||
// Prefill only (q_len > 1): measured on beast (27B, 30k-token
|
||||
// prompt, 2x RTX 5090), flash cuts prefill 24.8s → 22.1s but
|
||||
// REGRESSES decode ~20% (50 → 60 ms/token at 30k KV) — FA2
|
||||
// without flash-decoding is weak at query-length 1 and the
|
||||
// per-step layout transposes add overhead. Greedy outputs are
|
||||
// byte-identical either way, so this split is purely a
|
||||
// performance routing decision.
|
||||
if flash_attn_enabled()
|
||||
&& q.dim(2)? > 1
|
||||
&& q.device().is_cuda()
|
||||
&& (dtype == DType::F16 || dtype == DType::BF16)
|
||||
{
|
||||
// flash_attn wants (B, L, H, D); the callers carry (B, H, L, D).
|
||||
let qf = q.transpose(1, 2)?.contiguous()?;
|
||||
let kf = k.transpose(1, 2)?.contiguous()?;
|
||||
let vf = v.transpose(1, 2)?.contiguous()?;
|
||||
let causal = attn_mask.is_some();
|
||||
let ctx = candle_flash_attn::flash_attn(&qf, &kf, &vf, scale as f32, causal)?;
|
||||
return ctx.transpose(1, 2)?.contiguous();
|
||||
}
|
||||
}
|
||||
|
||||
// Eager fallback: materialise GQA-repeated K/V and the score matrix.
|
||||
let k = repeat_kv(k.clone(), num_kv_groups)?.contiguous()?;
|
||||
let v = repeat_kv(v.clone(), num_kv_groups)?.contiguous()?;
|
||||
let mut scores = (q.matmul(&k.transpose(2, 3)?)? * scale)?;
|
||||
if let Some(m) = attn_mask {
|
||||
scores = scores.broadcast_add(m)?;
|
||||
}
|
||||
let probs = candle_nn::ops::softmax_last_dim(&scores)?;
|
||||
probs.matmul(&v)
|
||||
}
|
||||
|
||||
pub struct Qwen3_5Attention {
|
||||
q_proj: Linear,
|
||||
k_proj: Linear,
|
||||
@@ -139,18 +219,10 @@ impl Qwen3_5Attention {
|
||||
// 4. KV cache.
|
||||
let (k, v) = self.kv_cache.append(&k, &v)?;
|
||||
|
||||
// 5. GQA repeat (cheap shape op).
|
||||
let k = repeat_kv(k, self.num_kv_groups)?.contiguous()?;
|
||||
let v = repeat_kv(v, self.num_kv_groups)?.contiguous()?;
|
||||
|
||||
// 6. Scaled dot-product + causal mask.
|
||||
// 5+6. Attention core — FlashAttention when available, eager
|
||||
// GQA-repeat + masked softmax otherwise (#95).
|
||||
let scale = 1.0_f64 / (self.head_dim as f64).sqrt();
|
||||
let mut scores = (q.matmul(&k.transpose(2, 3)?)? * scale)?;
|
||||
if let Some(m) = attn_mask {
|
||||
scores = scores.broadcast_add(m)?;
|
||||
}
|
||||
let probs = candle_nn::ops::softmax_last_dim(&scores)?;
|
||||
let ctx = probs.matmul(&v)?; // (B, H, L, D)
|
||||
let ctx = attention_context(&q, &k, &v, attn_mask, self.num_kv_groups, scale)?; // (B, H, L, D)
|
||||
|
||||
// 7. Reshape back, apply the output gate, project.
|
||||
let ctx = ctx
|
||||
|
||||
@@ -139,10 +139,42 @@ impl GatedDeltaNet {
|
||||
let conv_dim = key_dim * 2 + value_dim;
|
||||
|
||||
// ----- Linear projections (all `bias=False` in the reference). -----
|
||||
let in_proj_qkv = load_linear_no_bias(vb, "in_proj_qkv", cfg.hidden_size, conv_dim)?;
|
||||
let in_proj_z = load_linear_no_bias(vb, "in_proj_z", cfg.hidden_size, value_dim)?;
|
||||
let in_proj_b = load_linear_no_bias(vb, "in_proj_b", cfg.hidden_size, num_v_heads)?;
|
||||
let in_proj_a = load_linear_no_bias(vb, "in_proj_a", cfg.hidden_size, num_v_heads)?;
|
||||
// Two checkpoint layouts exist for the input projections:
|
||||
// - Qwen3.6 (qwen3_5): separate `in_proj_qkv` / `in_proj_z` /
|
||||
// `in_proj_b` / `in_proj_a`, with qkv stored as contiguous
|
||||
// [Q | K | V] blocks — loads directly.
|
||||
// - Qwen3-Next 80B-A3B (qwen3_next, #92): fused `in_proj_qkvz`
|
||||
// + `in_proj_ba`, **interleaved per key-head group** (see
|
||||
// `split_fused_qkvz`/`split_fused_ba`) — de-interleaved once
|
||||
// at load into the same contiguous layout, so the forward
|
||||
// path (incl. the conv over [Q|K|V] channels) is unchanged.
|
||||
let (in_proj_qkv, in_proj_z, in_proj_b, in_proj_a) =
|
||||
if vb.contains_tensor("in_proj_qkvz.weight") {
|
||||
let qkvz = vb
|
||||
.pp("in_proj_qkvz")
|
||||
.get((2 * key_dim + 2 * value_dim, cfg.hidden_size), "weight")
|
||||
.with_context(|| format!("load '{}/in_proj_qkvz/weight'", vb.prefix()))?;
|
||||
let ba = vb
|
||||
.pp("in_proj_ba")
|
||||
.get((2 * num_v_heads, cfg.hidden_size), "weight")
|
||||
.with_context(|| format!("load '{}/in_proj_ba/weight'", vb.prefix()))?;
|
||||
let (qkv_w, z_w) =
|
||||
split_fused_qkvz(&qkvz, num_k_heads, num_v_heads, head_k_dim, head_v_dim)?;
|
||||
let (b_w, a_w) = split_fused_ba(&ba, num_k_heads, num_v_heads)?;
|
||||
(
|
||||
Linear::new(qkv_w, None),
|
||||
Linear::new(z_w, None),
|
||||
Linear::new(b_w, None),
|
||||
Linear::new(a_w, None),
|
||||
)
|
||||
} else {
|
||||
(
|
||||
load_linear_no_bias(vb, "in_proj_qkv", cfg.hidden_size, conv_dim)?,
|
||||
load_linear_no_bias(vb, "in_proj_z", cfg.hidden_size, value_dim)?,
|
||||
load_linear_no_bias(vb, "in_proj_b", cfg.hidden_size, num_v_heads)?,
|
||||
load_linear_no_bias(vb, "in_proj_a", cfg.hidden_size, num_v_heads)?,
|
||||
)
|
||||
};
|
||||
let out_proj = load_linear_no_bias(vb, "out_proj", value_dim, cfg.hidden_size)?;
|
||||
|
||||
// ----- Conv1d weight (depthwise, bias=False). -----
|
||||
@@ -889,6 +921,61 @@ fn load_linear_no_bias(
|
||||
Ok(Linear::new(weight, None))
|
||||
}
|
||||
|
||||
/// De-interleave a fused `in_proj_qkvz.weight` (qwen3_next layout, #92)
|
||||
/// into a contiguous `[Q | K | V]` qkv weight plus a `Z` weight.
|
||||
///
|
||||
/// The fused rows are grouped **per key head**: for each of the
|
||||
/// `num_k_heads` groups (`r = num_v_heads / num_k_heads`, group stride
|
||||
/// `s = 2*head_k + 2*head_v*r`), the group holds
|
||||
/// `[q (head_k) | k (head_k) | v (head_v*r) | z (head_v*r)]` — the
|
||||
/// reshape in upstream `fix_query_key_value_ordering`
|
||||
/// `(num_k_heads, 2*head_k + 2*head_v*num_v/num_k)`. Concatenating the
|
||||
/// per-group regions restores the global-contiguous layout the rest of
|
||||
/// this module (incl. the conv over `[Q|K|V]` channels) expects.
|
||||
pub(crate) fn split_fused_qkvz(
|
||||
qkvz: &Tensor,
|
||||
num_k_heads: usize,
|
||||
num_v_heads: usize,
|
||||
head_k_dim: usize,
|
||||
head_v_dim: usize,
|
||||
) -> Result<(Tensor, Tensor)> {
|
||||
let r = num_v_heads / num_k_heads;
|
||||
let stride = 2 * head_k_dim + 2 * head_v_dim * r;
|
||||
let (mut qs, mut ks, mut vs, mut zs) = (Vec::new(), Vec::new(), Vec::new(), Vec::new());
|
||||
for g in 0..num_k_heads {
|
||||
let base = g * stride;
|
||||
qs.push(qkvz.narrow(0, base, head_k_dim)?);
|
||||
ks.push(qkvz.narrow(0, base + head_k_dim, head_k_dim)?);
|
||||
vs.push(qkvz.narrow(0, base + 2 * head_k_dim, head_v_dim * r)?);
|
||||
zs.push(qkvz.narrow(0, base + 2 * head_k_dim + head_v_dim * r, head_v_dim * r)?);
|
||||
}
|
||||
let parts: Vec<Tensor> = qs.into_iter().chain(ks).chain(vs).collect();
|
||||
let qkv = Tensor::cat(&parts, 0)?.contiguous()?;
|
||||
let z = Tensor::cat(&zs, 0)?.contiguous()?;
|
||||
Ok((qkv, z))
|
||||
}
|
||||
|
||||
/// De-interleave a fused `in_proj_ba.weight` (qwen3_next layout, #92)
|
||||
/// into per-v-head `b` (beta) and `a` (decay) weights. Same per-key-head
|
||||
/// grouping as [`split_fused_qkvz`]: each group holds `[b (r) | a (r)]`
|
||||
/// rows, `r = num_v_heads / num_k_heads`.
|
||||
pub(crate) fn split_fused_ba(
|
||||
ba: &Tensor,
|
||||
num_k_heads: usize,
|
||||
num_v_heads: usize,
|
||||
) -> Result<(Tensor, Tensor)> {
|
||||
let r = num_v_heads / num_k_heads;
|
||||
let (mut bs, mut r#as) = (Vec::new(), Vec::new());
|
||||
for g in 0..num_k_heads {
|
||||
let base = g * 2 * r;
|
||||
bs.push(ba.narrow(0, base, r)?);
|
||||
r#as.push(ba.narrow(0, base + r, r)?);
|
||||
}
|
||||
let b = Tensor::cat(&bs, 0)?.contiguous()?;
|
||||
let a = Tensor::cat(&r#as, 0)?.contiguous()?;
|
||||
Ok((b, a))
|
||||
}
|
||||
|
||||
/// Numerically-stable `softplus(x) = ln(1 + exp(x))`. Matches PyTorch's
|
||||
/// `F.softplus` default (beta=1, threshold=20: for large positive x,
|
||||
/// returns x as-is to avoid overflow in the exp).
|
||||
@@ -1138,6 +1225,13 @@ mod tests {
|
||||
linear_key_head_dim: 4,
|
||||
linear_value_head_dim: 4,
|
||||
linear_conv_kernel_dim: 4,
|
||||
num_experts: 0,
|
||||
num_experts_per_tok: 0,
|
||||
moe_intermediate_size: 0,
|
||||
shared_expert_intermediate_size: 0,
|
||||
decoder_sparse_step: 1,
|
||||
mlp_only_layers: Vec::new(),
|
||||
norm_topk_prob: false,
|
||||
};
|
||||
|
||||
// Build a synthetic VarBuilder with all-zeros weights.
|
||||
@@ -1179,4 +1273,115 @@ mod tests {
|
||||
let v: Vec<f32> = y.flatten_all().unwrap().to_vec1().unwrap();
|
||||
assert!(v.iter().all(|x| x.is_finite()));
|
||||
}
|
||||
|
||||
/// Interleave known per-head Q/K/V/Z (and B/A) rows into the fused
|
||||
/// qwen3_next layout, split, and expect the original contiguous
|
||||
/// blocks back. Layout under test: per key-head group g,
|
||||
/// `[q_g | k_g | v_g | z_g]` with r = num_v/num_k value heads per
|
||||
/// group (upstream `fix_query_key_value_ordering`).
|
||||
#[test]
|
||||
fn split_fused_qkvz_and_ba_roundtrip() {
|
||||
let dev = Device::Cpu;
|
||||
let (num_k, num_v, head_k, head_v, hidden) = (2usize, 4usize, 3usize, 2usize, 5usize);
|
||||
let r = num_v / num_k;
|
||||
|
||||
// Distinct constant per logical row so any mis-slicing shows.
|
||||
let row = |tag: f32| Tensor::full(tag, (1, hidden), &dev).unwrap();
|
||||
let mut fused_rows: Vec<Tensor> = Vec::new();
|
||||
let (mut q_rows, mut k_rows, mut v_rows, mut z_rows) =
|
||||
(Vec::new(), Vec::new(), Vec::new(), Vec::new());
|
||||
for g in 0..num_k {
|
||||
let base = 1000.0 * (g as f32 + 1.0);
|
||||
for i in 0..head_k {
|
||||
let t = row(base + i as f32);
|
||||
fused_rows.push(t.clone());
|
||||
q_rows.push(t);
|
||||
}
|
||||
for i in 0..head_k {
|
||||
let t = row(base + 100.0 + i as f32);
|
||||
fused_rows.push(t.clone());
|
||||
k_rows.push(t);
|
||||
}
|
||||
for i in 0..head_v * r {
|
||||
let t = row(base + 200.0 + i as f32);
|
||||
fused_rows.push(t.clone());
|
||||
v_rows.push(t);
|
||||
}
|
||||
for i in 0..head_v * r {
|
||||
let t = row(base + 300.0 + i as f32);
|
||||
fused_rows.push(t.clone());
|
||||
z_rows.push(t);
|
||||
}
|
||||
}
|
||||
let fused = Tensor::cat(&fused_rows, 0).unwrap();
|
||||
let expected_qkv = Tensor::cat(
|
||||
&q_rows
|
||||
.iter()
|
||||
.chain(k_rows.iter())
|
||||
.chain(v_rows.iter())
|
||||
.cloned()
|
||||
.collect::<Vec<_>>(),
|
||||
0,
|
||||
)
|
||||
.unwrap();
|
||||
let expected_z = Tensor::cat(&z_rows, 0).unwrap();
|
||||
|
||||
let (qkv, z) = split_fused_qkvz(&fused, num_k, num_v, head_k, head_v).unwrap();
|
||||
assert_eq!(qkv.dims(), &[2 * num_k * head_k + num_v * head_v, hidden]);
|
||||
let diff_qkv: f32 = (qkv - expected_qkv)
|
||||
.unwrap()
|
||||
.abs()
|
||||
.unwrap()
|
||||
.max_all()
|
||||
.unwrap()
|
||||
.to_scalar()
|
||||
.unwrap();
|
||||
let diff_z: f32 = (z - expected_z)
|
||||
.unwrap()
|
||||
.abs()
|
||||
.unwrap()
|
||||
.max_all()
|
||||
.unwrap()
|
||||
.to_scalar()
|
||||
.unwrap();
|
||||
assert_eq!(diff_qkv, 0.0);
|
||||
assert_eq!(diff_z, 0.0);
|
||||
|
||||
// ba: per group, [b (r rows) | a (r rows)].
|
||||
let mut ba_rows = Vec::new();
|
||||
let (mut b_rows, mut a_rows) = (Vec::new(), Vec::new());
|
||||
for g in 0..num_k {
|
||||
let base = 10.0 * (g as f32 + 1.0);
|
||||
for i in 0..r {
|
||||
let t = row(base + i as f32);
|
||||
ba_rows.push(t.clone());
|
||||
b_rows.push(t);
|
||||
}
|
||||
for i in 0..r {
|
||||
let t = row(base + 5.0 + i as f32);
|
||||
ba_rows.push(t.clone());
|
||||
a_rows.push(t);
|
||||
}
|
||||
}
|
||||
let ba = Tensor::cat(&ba_rows, 0).unwrap();
|
||||
let (b, a) = split_fused_ba(&ba, num_k, num_v).unwrap();
|
||||
let diff_b: f32 = (b - Tensor::cat(&b_rows, 0).unwrap())
|
||||
.unwrap()
|
||||
.abs()
|
||||
.unwrap()
|
||||
.max_all()
|
||||
.unwrap()
|
||||
.to_scalar()
|
||||
.unwrap();
|
||||
let diff_a: f32 = (a - Tensor::cat(&a_rows, 0).unwrap())
|
||||
.unwrap()
|
||||
.abs()
|
||||
.unwrap()
|
||||
.max_all()
|
||||
.unwrap()
|
||||
.to_scalar()
|
||||
.unwrap();
|
||||
assert_eq!(diff_b, 0.0);
|
||||
assert_eq!(diff_a, 0.0);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -17,12 +17,32 @@ pub struct Qwen3_5MLP {
|
||||
}
|
||||
|
||||
impl Qwen3_5MLP {
|
||||
/// Construct directly from pre-built projections (MoE-block tests).
|
||||
#[cfg(test)]
|
||||
pub(crate) fn from_weights(gate_proj: Linear, up_proj: Linear, down_proj: Linear) -> Self {
|
||||
Self {
|
||||
gate_proj,
|
||||
up_proj,
|
||||
down_proj,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn load(cfg: &TextConfig, vb: &ShardedVarBuilder) -> Result<Self> {
|
||||
let h = cfg.hidden_size;
|
||||
let i = cfg.intermediate_size;
|
||||
let gate_proj = load_linear_no_bias(vb, "gate_proj", h, i)?;
|
||||
let up_proj = load_linear_no_bias(vb, "up_proj", h, i)?;
|
||||
let down_proj = load_linear_no_bias(vb, "down_proj", i, h)?;
|
||||
Self::load_with_dims(vb, cfg.hidden_size, cfg.intermediate_size)
|
||||
}
|
||||
|
||||
/// Load with explicit dims — the MoE block (#92) reuses this SwiGLU
|
||||
/// shape for routed experts (`moe_intermediate_size`) and the shared
|
||||
/// expert (`shared_expert_intermediate_size`), both narrower than
|
||||
/// the dense `intermediate_size`.
|
||||
pub fn load_with_dims(
|
||||
vb: &ShardedVarBuilder,
|
||||
hidden: usize,
|
||||
intermediate: usize,
|
||||
) -> Result<Self> {
|
||||
let gate_proj = load_linear_no_bias(vb, "gate_proj", hidden, intermediate)?;
|
||||
let up_proj = load_linear_no_bias(vb, "up_proj", hidden, intermediate)?;
|
||||
let down_proj = load_linear_no_bias(vb, "down_proj", intermediate, hidden)?;
|
||||
Ok(Self {
|
||||
gate_proj,
|
||||
up_proj,
|
||||
|
||||
@@ -76,6 +76,7 @@ pub mod decoder;
|
||||
pub mod full_attn;
|
||||
pub mod linear_attn;
|
||||
pub mod mlp;
|
||||
pub mod moe;
|
||||
pub mod rmsnorm;
|
||||
pub mod rope;
|
||||
pub mod snapshot;
|
||||
@@ -90,6 +91,13 @@ use rope::RotaryEmbedding;
|
||||
/// magic strings.
|
||||
pub const MODEL_TYPE: &str = "qwen3_5";
|
||||
|
||||
/// `model_type` of the MoE sibling family (Qwen3-Next-80B-A3B /
|
||||
/// Qwen3-Coder-Next): the same Gated DeltaNet hybrid layer stack with
|
||||
/// a high-sparsity MoE FFN per layer. Served by this same arch module —
|
||||
/// [`Config::from_config_json`] normalises the flat qwen3_next
|
||||
/// `config.json` layout into the nested shape used here.
|
||||
pub const MODEL_TYPE_NEXT: &str = "qwen3_next";
|
||||
|
||||
/// Top-level shape of Qwen3-Next's `config.json`. The real
|
||||
/// hyperparameters live in `text_config`; the rest is multimodal /
|
||||
/// tokeniser glue we don't need for the language-model forward.
|
||||
@@ -117,6 +125,79 @@ pub struct Config {
|
||||
pub image_token_id: Option<u32>,
|
||||
}
|
||||
|
||||
impl Config {
|
||||
/// Parse a `config.json` for either family this arch serves,
|
||||
/// normalising layout differences (#92):
|
||||
///
|
||||
/// - `model_type == "qwen3_5"` (Qwen3.6): hyperparameters nested
|
||||
/// under `text_config`, RoPE nested under `rope_parameters` —
|
||||
/// deserialises directly.
|
||||
/// - `model_type == "qwen3_next"` (Qwen3-Next-80B-A3B family):
|
||||
/// **flat** layout — hyperparameters at the top level,
|
||||
/// `rope_theta`/`partial_rotary_factor` flat, no vision block.
|
||||
/// Wrapped into the nested shape here. The output gate on full
|
||||
/// attention is unconditional in the upstream qwen3_next
|
||||
/// implementation (the config carries no flag), so
|
||||
/// `attn_output_gate` is forced on.
|
||||
///
|
||||
/// Both variants may omit `layer_types` (qwen3_next always does);
|
||||
/// it is derived from `full_attention_interval` using the upstream
|
||||
/// convention: layer `i` is `full_attention` iff
|
||||
/// `(i + 1) % interval == 0`, else `linear_attention`.
|
||||
pub fn from_config_json(json: &str) -> Result<Self> {
|
||||
let v: serde_json::Value =
|
||||
serde_json::from_str(json).context("parse config.json as JSON")?;
|
||||
let model_type = v
|
||||
.get("model_type")
|
||||
.and_then(|m| m.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
|
||||
let mut cfg: Config = if model_type == MODEL_TYPE_NEXT {
|
||||
let mut text = v.clone();
|
||||
if text.get("rope_parameters").is_none() {
|
||||
let mut rope = serde_json::Map::new();
|
||||
for key in ["rope_theta", "partial_rotary_factor", "rope_type"] {
|
||||
if let Some(val) = v.get(key) {
|
||||
rope.insert(key.to_string(), val.clone());
|
||||
}
|
||||
}
|
||||
text["rope_parameters"] = serde_json::Value::Object(rope);
|
||||
}
|
||||
let mut text_config: TextConfig = serde_json::from_value(text)
|
||||
.context("parse flat qwen3_next config.json hyperparameters")?;
|
||||
text_config.attn_output_gate = true;
|
||||
Config {
|
||||
model_type,
|
||||
text_config,
|
||||
vision_config: None,
|
||||
image_token_id: None,
|
||||
}
|
||||
} else {
|
||||
serde_json::from_str(json).context("parse nested qwen3_5 config.json")?
|
||||
};
|
||||
|
||||
if cfg.text_config.layer_types.is_empty() {
|
||||
let interval = cfg.text_config.full_attention_interval.unwrap_or(4);
|
||||
anyhow::ensure!(
|
||||
interval > 0,
|
||||
"full_attention_interval must be >= 1 to derive layer_types"
|
||||
);
|
||||
cfg.text_config.layer_types = (0..cfg.text_config.num_hidden_layers)
|
||||
.map(|i| {
|
||||
if (i + 1).is_multiple_of(interval) {
|
||||
"full_attention".to_string()
|
||||
} else {
|
||||
"linear_attention".to_string()
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
}
|
||||
|
||||
Ok(cfg)
|
||||
}
|
||||
}
|
||||
|
||||
/// Inner config (the `text_config` block). Mirrors the Qwen3 layout
|
||||
/// but with the extras Qwen3-Next adds (`attn_output_gate`,
|
||||
/// `layer_types`, `full_attention_interval`, larger `head_dim`).
|
||||
@@ -185,6 +266,56 @@ pub struct TextConfig {
|
||||
/// (Qwen3.6-27B: 4).
|
||||
#[serde(default)]
|
||||
pub linear_conv_kernel_dim: usize,
|
||||
|
||||
// --- High-sparsity MoE FFN (Qwen3-Next 80B-A3B family, #92) --------
|
||||
// All default to the dense case (0 experts) so existing dense
|
||||
// configs (Qwen3.6-27B) deserialise unchanged. A layer gets the MoE
|
||||
// FFN iff `layer_uses_moe` says so; otherwise the dense SwiGLU.
|
||||
/// Total routed experts per MoE layer (80B-A3B: 512). `0` → dense
|
||||
/// model, no MoE anywhere.
|
||||
#[serde(default)]
|
||||
pub num_experts: usize,
|
||||
/// Experts activated per token (80B-A3B: 10).
|
||||
#[serde(default)]
|
||||
pub num_experts_per_tok: usize,
|
||||
/// Per-expert FFN width (80B-A3B: 512). Distinct from the dense
|
||||
/// `intermediate_size`.
|
||||
#[serde(default)]
|
||||
pub moe_intermediate_size: usize,
|
||||
/// Width of the always-on shared expert (80B-A3B: 512). `0` → no
|
||||
/// shared expert (Qwen3-30B-A3B style).
|
||||
#[serde(default)]
|
||||
pub shared_expert_intermediate_size: usize,
|
||||
/// Every `decoder_sparse_step`-th layer is MoE (1 → all layers,
|
||||
/// the 80B-A3B case). Follows the upstream `(i+1) % step == 0`
|
||||
/// convention.
|
||||
#[serde(default = "default_decoder_sparse_step")]
|
||||
pub decoder_sparse_step: usize,
|
||||
/// Layer indices forced to the dense MLP even when MoE is on.
|
||||
/// Empty for 80B-A3B.
|
||||
#[serde(default)]
|
||||
pub mlp_only_layers: Vec<usize>,
|
||||
/// Renormalise the top-k routing weights to sum to 1 (80B-A3B:
|
||||
/// true). Upstream selects top-k *after* softmax over all experts.
|
||||
#[serde(default)]
|
||||
pub norm_topk_prob: bool,
|
||||
}
|
||||
|
||||
impl TextConfig {
|
||||
/// Whether decoder layer `layer_idx` carries the MoE FFN (vs the
|
||||
/// dense SwiGLU). Mirrors upstream `Qwen3NextDecoderLayer`:
|
||||
/// experts configured, layer not in `mlp_only_layers`, and on the
|
||||
/// `decoder_sparse_step` grid.
|
||||
pub fn layer_uses_moe(&self, layer_idx: usize) -> bool {
|
||||
self.num_experts > 0
|
||||
&& self.decoder_sparse_step > 0
|
||||
&& !self.mlp_only_layers.contains(&layer_idx)
|
||||
&& (layer_idx + 1).is_multiple_of(self.decoder_sparse_step)
|
||||
}
|
||||
}
|
||||
|
||||
fn default_decoder_sparse_step() -> usize {
|
||||
1
|
||||
}
|
||||
|
||||
fn default_hidden_act() -> String {
|
||||
@@ -330,17 +461,20 @@ pub struct Qwen3_5Model {
|
||||
}
|
||||
|
||||
impl Qwen3_5Model {
|
||||
pub fn load(cfg: &TextConfig, vb: &ShardedVarBuilder) -> Result<Self> {
|
||||
/// `text_prefix` is where the text core lives in the checkpoint:
|
||||
/// - Qwen3.6 (multimodal, `model_type = "qwen3_5"`):
|
||||
/// `model.language_model` — sibling to `model.visual.*` (the
|
||||
/// vision tower) and top-level `lm_head` / `mtp.*`.
|
||||
/// - Qwen3-Next-80B-A3B (text-only, `model_type = "qwen3_next"`):
|
||||
/// plain `model`.
|
||||
///
|
||||
/// [`Qwen3_5ForCausalLM::new`] picks by `Config::model_type` via
|
||||
/// [`text_weight_prefix`].
|
||||
pub fn load(cfg: &TextConfig, vb: &ShardedVarBuilder, text_prefix: &str) -> Result<Self> {
|
||||
let dtype = vb.dtype();
|
||||
let device = vb.device().clone();
|
||||
|
||||
// Qwen3-Next is a multimodal architecture whose text core lives
|
||||
// under `model.language_model.*` — sibling to `model.visual.*`
|
||||
// (the vision tower) and to top-level `lm_head` / `mtp.*`.
|
||||
// Every text-side tensor in the safetensors files is under
|
||||
// this prefix; we ignore the vision and MTP weights for
|
||||
// language-model inference.
|
||||
let text_vb = vb.pp("model.language_model");
|
||||
let text_vb = vb.pp(text_prefix);
|
||||
|
||||
let embed_vb = text_vb.pp("embed_tokens");
|
||||
let embed_weight = embed_vb
|
||||
@@ -444,6 +578,67 @@ impl Qwen3_5Model {
|
||||
self.forward_inner(input, offset, None, None, &[], None)
|
||||
}
|
||||
|
||||
/// Lockstep batched decode step (#98): `input` is `(B, 1)` — one
|
||||
/// new token per batch row — with each row at its own sequence
|
||||
/// position `positions[i]` (typically `prefix_lens[i] + step`).
|
||||
/// The cache must hold batched state (see
|
||||
/// `snapshot::assemble_batch`); `attn_mask` is the padding mask
|
||||
/// from [`Self::batch_decode_mask`] (or `None` when no row is
|
||||
/// padded). Text-only: `rope_delta` is ignored — positions are
|
||||
/// explicit and vision requests never enter the batch path.
|
||||
pub fn forward_batch_decode(
|
||||
&mut self,
|
||||
input: &Tensor,
|
||||
positions: &[usize],
|
||||
attn_mask: Option<&Tensor>,
|
||||
) -> candle_core::Result<Tensor> {
|
||||
let (b, l) = input.dims2()?;
|
||||
if l != 1 {
|
||||
candle_core::bail!("forward_batch_decode: expected (B, 1) input, got (B, {l})");
|
||||
}
|
||||
if positions.len() != b {
|
||||
candle_core::bail!(
|
||||
"forward_batch_decode: {} positions for batch of {b}",
|
||||
positions.len()
|
||||
);
|
||||
}
|
||||
let mut h = self.embed_tokens.forward(input)?;
|
||||
let (cos, sin) = self.rotary.batch_cos_sin(positions)?;
|
||||
for layer in &mut self.layers {
|
||||
h = layer.forward(&h, attn_mask, &cos, &sin)?;
|
||||
}
|
||||
self.norm.forward(&h)
|
||||
}
|
||||
|
||||
/// Additive padding mask for a batched decode step: shape
|
||||
/// `(B, 1, 1, total_len)`, `-inf` on each row's padding gap
|
||||
/// `[prefix_lens[i], padded_len)`, zero elsewhere. `total_len` is
|
||||
/// the KV length *after* this step's append (`padded_len + step +
|
||||
/// 1`). Returns `None` when no row is padded (uniform prefix
|
||||
/// lengths) — the decode step then needs no mask at all, matching
|
||||
/// the single-sequence fast path.
|
||||
pub fn batch_decode_mask(
|
||||
&self,
|
||||
prefix_lens: &[usize],
|
||||
padded_len: usize,
|
||||
total_len: usize,
|
||||
) -> candle_core::Result<Option<Tensor>> {
|
||||
if prefix_lens.iter().all(|&len| len == padded_len) {
|
||||
return Ok(None);
|
||||
}
|
||||
let minf = f32::NEG_INFINITY;
|
||||
let b = prefix_lens.len();
|
||||
let mask: Vec<f32> = prefix_lens
|
||||
.iter()
|
||||
.flat_map(|&len| {
|
||||
(0..total_len).map(move |j| if j >= len && j < padded_len { minf } else { 0. })
|
||||
})
|
||||
.collect();
|
||||
Ok(Some(
|
||||
Tensor::from_vec(mask, (b, 1, 1, total_len), &self.device)?.to_dtype(self.dtype)?,
|
||||
))
|
||||
}
|
||||
|
||||
/// 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
|
||||
@@ -605,10 +800,20 @@ pub struct Qwen3_5ForCausalLM {
|
||||
image_token_id: Option<u32>,
|
||||
}
|
||||
|
||||
/// Checkpoint prefix of the text core for a given `model_type` — see
|
||||
/// [`Qwen3_5Model::load`].
|
||||
pub fn text_weight_prefix(model_type: &str) -> &'static str {
|
||||
if model_type == MODEL_TYPE_NEXT {
|
||||
"model"
|
||||
} else {
|
||||
"model.language_model"
|
||||
}
|
||||
}
|
||||
|
||||
impl Qwen3_5ForCausalLM {
|
||||
pub fn new(config: Config, vb: ShardedVarBuilder) -> Result<Self> {
|
||||
let cfg = &config.text_config;
|
||||
let base = Qwen3_5Model::load(cfg, &vb)?;
|
||||
let base = Qwen3_5Model::load(cfg, &vb, text_weight_prefix(&config.model_type))?;
|
||||
let lm_head = if cfg.tie_word_embeddings {
|
||||
Linear::new(base.embed_weight().clone(), None)
|
||||
} else {
|
||||
@@ -676,6 +881,34 @@ impl Qwen3_5ForCausalLM {
|
||||
hidden.i((.., l - 1.., ..))?.apply(&self.lm_head)
|
||||
}
|
||||
|
||||
/// Lockstep batched decode step (#98): `(B, 1)` input, per-row
|
||||
/// positions, padding mask from
|
||||
/// [`Qwen3_5Model::batch_decode_mask`]. Returns `(B, 1,
|
||||
/// vocab_size)` — one logits row per batch row.
|
||||
pub fn forward_batch_decode(
|
||||
&mut self,
|
||||
input: &Tensor,
|
||||
positions: &[usize],
|
||||
attn_mask: Option<&Tensor>,
|
||||
) -> candle_core::Result<Tensor> {
|
||||
let hidden = self
|
||||
.base
|
||||
.forward_batch_decode(input, positions, attn_mask)?;
|
||||
hidden.apply(&self.lm_head)
|
||||
}
|
||||
|
||||
/// Padding mask for a batched decode step — see
|
||||
/// [`Qwen3_5Model::batch_decode_mask`].
|
||||
pub fn batch_decode_mask(
|
||||
&self,
|
||||
prefix_lens: &[usize],
|
||||
padded_len: usize,
|
||||
total_len: usize,
|
||||
) -> candle_core::Result<Option<Tensor>> {
|
||||
self.base
|
||||
.batch_decode_mask(prefix_lens, padded_len, total_len)
|
||||
}
|
||||
|
||||
/// Stage B: forward with image-embedding splice. Mirrors `forward`
|
||||
/// but routes through `Qwen3_5Model::forward_with_vision` so the
|
||||
/// LM's input embeddings get the image patches spliced in at
|
||||
@@ -915,6 +1148,233 @@ mod tests {
|
||||
assert_eq!(cfg.text_config.layer_types.len(), 4);
|
||||
assert_eq!(cfg.text_config.rope_parameters.rope_theta, 10_000_000.0);
|
||||
assert!((cfg.text_config.rope_parameters.partial_rotary_factor - 0.25).abs() < 1e-6);
|
||||
// Dense config: no MoE anywhere.
|
||||
assert_eq!(cfg.text_config.num_experts, 0);
|
||||
assert!(!cfg.text_config.layer_uses_moe(0));
|
||||
|
||||
// The normalising entry point must agree with plain serde for
|
||||
// the nested shape (and leave the explicit layer_types alone).
|
||||
let via_norm = Config::from_config_json(raw).expect("normalised parse");
|
||||
assert_eq!(via_norm.text_config.layer_types.len(), 4);
|
||||
assert_eq!(via_norm.text_config.layer_types[3], "full_attention");
|
||||
}
|
||||
|
||||
/// The flat qwen3_next layout (Qwen3-Next-80B-A3B family): all
|
||||
/// hyperparameters top-level, flat rope fields, no `layer_types`,
|
||||
/// MoE fields present. Sample mirrors
|
||||
/// `Qwen/Qwen3-Next-80B-A3B-Instruct/config.json`.
|
||||
#[test]
|
||||
fn config_normalises_the_flat_qwen3_next_shape() {
|
||||
let raw = r#"{
|
||||
"architectures": ["Qwen3NextForCausalLM"],
|
||||
"model_type": "qwen3_next",
|
||||
"vocab_size": 151936,
|
||||
"hidden_size": 2048,
|
||||
"intermediate_size": 5120,
|
||||
"num_hidden_layers": 48,
|
||||
"num_attention_heads": 16,
|
||||
"num_key_value_heads": 2,
|
||||
"head_dim": 256,
|
||||
"max_position_embeddings": 262144,
|
||||
"partial_rotary_factor": 0.25,
|
||||
"rope_theta": 10000000,
|
||||
"rms_norm_eps": 1e-6,
|
||||
"tie_word_embeddings": false,
|
||||
"full_attention_interval": 4,
|
||||
"linear_conv_kernel_dim": 4,
|
||||
"linear_key_head_dim": 128,
|
||||
"linear_num_key_heads": 16,
|
||||
"linear_num_value_heads": 32,
|
||||
"linear_value_head_dim": 128,
|
||||
"decoder_sparse_step": 1,
|
||||
"mlp_only_layers": [],
|
||||
"moe_intermediate_size": 512,
|
||||
"norm_topk_prob": true,
|
||||
"num_experts": 512,
|
||||
"num_experts_per_tok": 10,
|
||||
"shared_expert_intermediate_size": 512
|
||||
}"#;
|
||||
let cfg = Config::from_config_json(raw).expect("parse qwen3_next config");
|
||||
assert_eq!(cfg.model_type, MODEL_TYPE_NEXT);
|
||||
assert!(cfg.vision_config.is_none());
|
||||
|
||||
let t = &cfg.text_config;
|
||||
assert_eq!(t.hidden_size, 2048);
|
||||
// Flat rope fields normalised into the nested block.
|
||||
assert_eq!(t.rope_parameters.rope_theta, 10_000_000.0);
|
||||
assert!((t.rope_parameters.partial_rotary_factor - 0.25).abs() < 1e-6);
|
||||
// Output-gated attention is unconditional for qwen3_next.
|
||||
assert!(t.attn_output_gate);
|
||||
// layer_types derived from the interval: (i+1) % 4 == 0 → full.
|
||||
assert_eq!(t.layer_types.len(), 48);
|
||||
assert_eq!(t.layer_types[3], "full_attention");
|
||||
assert_eq!(t.layer_types[47], "full_attention");
|
||||
assert_eq!(t.layer_types[0], "linear_attention");
|
||||
assert_eq!(t.layer_types[46], "linear_attention");
|
||||
assert_eq!(
|
||||
t.layer_types
|
||||
.iter()
|
||||
.filter(|s| *s == "full_attention")
|
||||
.count(),
|
||||
12
|
||||
);
|
||||
// MoE hyperparameters land.
|
||||
assert_eq!(t.num_experts, 512);
|
||||
assert_eq!(t.num_experts_per_tok, 10);
|
||||
assert_eq!(t.moe_intermediate_size, 512);
|
||||
assert_eq!(t.shared_expert_intermediate_size, 512);
|
||||
assert!(t.norm_topk_prob);
|
||||
// decoder_sparse_step 1 + empty mlp_only_layers → every layer MoE.
|
||||
assert!(t.layer_uses_moe(0));
|
||||
assert!(t.layer_uses_moe(47));
|
||||
}
|
||||
|
||||
/// End-to-end structural check for the qwen3_next path (#92): a
|
||||
/// tiny random-weight checkpoint in the **flat** layout (`model.*`
|
||||
/// prefix, fused `in_proj_qkvz`/`in_proj_ba`, per-expert MoE
|
||||
/// tensors, shared expert) loads through `Config::from_config_json`
|
||||
/// and `Qwen3_5ForCausalLM::new`, producing finite logits of the
|
||||
/// right shape. Numerical parity vs HF is pinned separately by the
|
||||
/// `qwen3_next_parity` fixture test.
|
||||
#[test]
|
||||
fn tiny_qwen3_next_checkpoint_loads_and_forwards() {
|
||||
use candle_core::Device;
|
||||
use std::collections::HashMap;
|
||||
|
||||
let raw = r#"{
|
||||
"model_type": "qwen3_next",
|
||||
"vocab_size": 32, "hidden_size": 8, "intermediate_size": 16,
|
||||
"num_hidden_layers": 2, "num_attention_heads": 2,
|
||||
"num_key_value_heads": 1, "head_dim": 4,
|
||||
"max_position_embeddings": 64, "rms_norm_eps": 1e-6,
|
||||
"full_attention_interval": 2,
|
||||
"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,
|
||||
"num_experts": 4, "num_experts_per_tok": 2,
|
||||
"moe_intermediate_size": 4,
|
||||
"shared_expert_intermediate_size": 4,
|
||||
"norm_topk_prob": true
|
||||
}"#;
|
||||
let cfg = Config::from_config_json(raw).expect("parse tiny qwen3_next config");
|
||||
assert_eq!(cfg.text_config.layer_types[0], "linear_attention");
|
||||
assert_eq!(cfg.text_config.layer_types[1], "full_attention");
|
||||
|
||||
let dev = Device::Cpu;
|
||||
let randn = |shape: &[usize]| Tensor::randn(0f32, 0.1f32, shape, &dev).unwrap();
|
||||
let ones = |shape: &[usize]| Tensor::ones(shape, DType::F32, &dev).unwrap();
|
||||
let mut t: HashMap<String, Tensor> = HashMap::new();
|
||||
|
||||
let (h, vocab) = (8usize, 32usize);
|
||||
t.insert("model.embed_tokens.weight".into(), randn(&[vocab, h]));
|
||||
t.insert("lm_head.weight".into(), randn(&[vocab, h]));
|
||||
t.insert("model.norm.weight".into(), ones(&[h]));
|
||||
|
||||
let moe = |t: &mut HashMap<String, Tensor>, p: &str| {
|
||||
t.insert(format!("{p}.gate.weight"), randn(&[4, h]));
|
||||
for e in 0..4 {
|
||||
t.insert(format!("{p}.experts.{e}.gate_proj.weight"), randn(&[4, h]));
|
||||
t.insert(format!("{p}.experts.{e}.up_proj.weight"), randn(&[4, h]));
|
||||
t.insert(format!("{p}.experts.{e}.down_proj.weight"), randn(&[h, 4]));
|
||||
}
|
||||
t.insert(
|
||||
format!("{p}.shared_expert.gate_proj.weight"),
|
||||
randn(&[4, h]),
|
||||
);
|
||||
t.insert(format!("{p}.shared_expert.up_proj.weight"), randn(&[4, h]));
|
||||
t.insert(
|
||||
format!("{p}.shared_expert.down_proj.weight"),
|
||||
randn(&[h, 4]),
|
||||
);
|
||||
t.insert(format!("{p}.shared_expert_gate.weight"), randn(&[1, h]));
|
||||
};
|
||||
|
||||
// Layer 0: linear_attention with the FUSED qwen3_next input
|
||||
// projections. key_dim = 2*4 = 8, value_dim = 4*4 = 16 →
|
||||
// qkvz rows = 2*8 + 2*16 = 48, ba rows = 2*4 = 8, conv_dim = 32.
|
||||
let l0 = "model.layers.0";
|
||||
t.insert(
|
||||
format!("{l0}.linear_attn.in_proj_qkvz.weight"),
|
||||
randn(&[48, h]),
|
||||
);
|
||||
t.insert(
|
||||
format!("{l0}.linear_attn.in_proj_ba.weight"),
|
||||
randn(&[8, h]),
|
||||
);
|
||||
t.insert(
|
||||
format!("{l0}.linear_attn.conv1d.weight"),
|
||||
randn(&[32, 1, 4]),
|
||||
);
|
||||
t.insert(format!("{l0}.linear_attn.dt_bias"), randn(&[4]));
|
||||
t.insert(format!("{l0}.linear_attn.A_log"), randn(&[4]));
|
||||
t.insert(format!("{l0}.linear_attn.norm.weight"), ones(&[4]));
|
||||
t.insert(format!("{l0}.linear_attn.out_proj.weight"), randn(&[h, 16]));
|
||||
t.insert(format!("{l0}.input_layernorm.weight"), ones(&[h]));
|
||||
t.insert(format!("{l0}.post_attention_layernorm.weight"), ones(&[h]));
|
||||
moe(&mut t, &format!("{l0}.mlp"));
|
||||
|
||||
// Layer 1: full_attention (output-gated: q_proj is 2×).
|
||||
let l1 = "model.layers.1";
|
||||
t.insert(
|
||||
format!("{l1}.self_attn.q_proj.weight"),
|
||||
randn(&[2 * 2 * 4, h]),
|
||||
);
|
||||
t.insert(format!("{l1}.self_attn.k_proj.weight"), randn(&[4, h]));
|
||||
t.insert(format!("{l1}.self_attn.v_proj.weight"), randn(&[4, h]));
|
||||
t.insert(format!("{l1}.self_attn.o_proj.weight"), randn(&[h, 8]));
|
||||
t.insert(format!("{l1}.self_attn.q_norm.weight"), ones(&[4]));
|
||||
t.insert(format!("{l1}.self_attn.k_norm.weight"), ones(&[4]));
|
||||
t.insert(format!("{l1}.input_layernorm.weight"), ones(&[h]));
|
||||
t.insert(format!("{l1}.post_attention_layernorm.weight"), ones(&[h]));
|
||||
moe(&mut t, &format!("{l1}.mlp"));
|
||||
|
||||
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; nothing mutates it.
|
||||
let vb = unsafe {
|
||||
candle_nn::var_builder::ShardedSafeTensors::var_builder(
|
||||
std::slice::from_ref(&path),
|
||||
DType::F32,
|
||||
&dev,
|
||||
)
|
||||
.expect("build ShardedVarBuilder")
|
||||
};
|
||||
|
||||
let mut model = Qwen3_5ForCausalLM::new(cfg, vb).expect("load tiny qwen3_next checkpoint");
|
||||
let input = Tensor::new(&[1u32, 5, 9], &dev)
|
||||
.unwrap()
|
||||
.unsqueeze(0)
|
||||
.unwrap();
|
||||
let logits = model.forward(&input, 0).expect("forward");
|
||||
assert_eq!(logits.dims(), &[1, 1, vocab]);
|
||||
let v: Vec<f32> = logits.flatten_all().unwrap().to_vec1().unwrap();
|
||||
assert!(v.iter().all(|x| x.is_finite()), "logits must be finite");
|
||||
}
|
||||
|
||||
/// `mlp_only_layers` and `decoder_sparse_step` gate `layer_uses_moe`
|
||||
/// per the upstream convention.
|
||||
#[test]
|
||||
fn layer_uses_moe_respects_step_and_exclusions() {
|
||||
let raw = r#"{
|
||||
"model_type": "qwen3_next",
|
||||
"vocab_size": 8, "hidden_size": 8, "intermediate_size": 8,
|
||||
"num_hidden_layers": 8, "num_attention_heads": 2,
|
||||
"num_key_value_heads": 1, "head_dim": 4,
|
||||
"max_position_embeddings": 128, "rms_norm_eps": 1e-6,
|
||||
"num_experts": 4, "num_experts_per_tok": 2,
|
||||
"moe_intermediate_size": 8,
|
||||
"decoder_sparse_step": 2,
|
||||
"mlp_only_layers": [3]
|
||||
}"#;
|
||||
let cfg = Config::from_config_json(raw).expect("parse");
|
||||
let t = &cfg.text_config;
|
||||
// step 2 → layers 1, 3, 5, 7 are on the sparse grid…
|
||||
assert!(!t.layer_uses_moe(0));
|
||||
assert!(t.layer_uses_moe(1));
|
||||
// …but 3 is excluded by mlp_only_layers.
|
||||
assert!(!t.layer_uses_moe(3));
|
||||
assert!(t.layer_uses_moe(5));
|
||||
}
|
||||
|
||||
/// `splice_runs` replaces (1, L, H) embedding rows at the given
|
||||
|
||||
360
crates/neuron/src/harness/arch/qwen3_5/moe.rs
Normal file
360
crates/neuron/src/harness/arch/qwen3_5/moe.rs
Normal file
@@ -0,0 +1,360 @@
|
||||
//! High-sparsity MoE FFN block for the qwen3_next family (#92).
|
||||
//!
|
||||
//! Qwen3-Next-80B-A3B replaces the dense SwiGLU in (almost) every
|
||||
//! decoder layer with `Qwen3NextSparseMoeBlock`: a top-k router over
|
||||
//! `num_experts` small SwiGLU experts, plus an always-on **shared
|
||||
//! expert** mixed in through a per-token sigmoid gate:
|
||||
//!
|
||||
//! ```text
|
||||
//! probs = softmax(gate(x)) # over ALL experts, f32
|
||||
//! w, idx = topk(probs, num_experts_per_tok)
|
||||
//! w = w / sum(w) # iff norm_topk_prob
|
||||
//! routed = Σ_j w_j · expert_{idx_j}(x)
|
||||
//! shared = sigmoid(shared_expert_gate(x)) · shared_expert(x)
|
||||
//! y = routed + shared
|
||||
//! ```
|
||||
//!
|
||||
//! Routing follows the upstream softmax-then-topk order (NOT
|
||||
//! topk-then-softmax — the renormalisation only equals softmax over
|
||||
//! the selected logits when `norm_topk_prob` is on, and the reference
|
||||
//! renormalises the *global* softmax values).
|
||||
//!
|
||||
//! ## Dispatch strategy
|
||||
//!
|
||||
//! This is the correctness-first implementation: a host-side scatter
|
||||
//! loop over the experts that actually received tokens (the pattern
|
||||
//! candle-transformers' `Qwen3SparseMoeBlock` uses). Batch-1 decode
|
||||
//! touches `num_experts_per_tok` experts per layer; prefill batches
|
||||
//! per-expert token groups. The fused grouped-GEMM path (slice 4)
|
||||
//! replaces the loop behind the same `forward` signature.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use candle_core::{DType, Module, Tensor};
|
||||
use candle_nn::Linear;
|
||||
use candle_nn::var_builder::ShardedVarBuilder;
|
||||
|
||||
use super::TextConfig;
|
||||
use super::mlp::Qwen3_5MLP;
|
||||
|
||||
pub struct Qwen3_5MoeBlock {
|
||||
/// Router: `(num_experts, hidden)`, checkpoint name `mlp.gate`.
|
||||
gate: Linear,
|
||||
/// Routed experts, checkpoint names `mlp.experts.{i}.{gate,up,down}_proj`.
|
||||
experts: Vec<Qwen3_5MLP>,
|
||||
/// Always-on expert, `mlp.shared_expert.*`. `None` when the config
|
||||
/// declares no shared expert (Qwen3-30B-A3B style).
|
||||
shared_expert: Option<Qwen3_5MLP>,
|
||||
/// Per-token sigmoid mix for the shared expert: `(1, hidden)`,
|
||||
/// checkpoint name `mlp.shared_expert_gate`.
|
||||
shared_expert_gate: Option<Linear>,
|
||||
num_experts_per_tok: usize,
|
||||
norm_topk_prob: bool,
|
||||
}
|
||||
|
||||
impl Qwen3_5MoeBlock {
|
||||
pub fn load(cfg: &TextConfig, vb: &ShardedVarBuilder) -> Result<Self> {
|
||||
anyhow::ensure!(
|
||||
cfg.num_experts > 0 && cfg.num_experts_per_tok > 0 && cfg.moe_intermediate_size > 0,
|
||||
"MoE block needs num_experts ({}), num_experts_per_tok ({}) and \
|
||||
moe_intermediate_size ({}) all > 0",
|
||||
cfg.num_experts,
|
||||
cfg.num_experts_per_tok,
|
||||
cfg.moe_intermediate_size,
|
||||
);
|
||||
anyhow::ensure!(
|
||||
cfg.num_experts_per_tok <= cfg.num_experts,
|
||||
"num_experts_per_tok ({}) exceeds num_experts ({})",
|
||||
cfg.num_experts_per_tok,
|
||||
cfg.num_experts,
|
||||
);
|
||||
|
||||
let h = cfg.hidden_size;
|
||||
|
||||
let gate_weight = vb
|
||||
.pp("gate")
|
||||
.get((cfg.num_experts, h), "weight")
|
||||
.with_context(|| format!("load '{}/gate/weight'", vb.prefix()))?;
|
||||
let gate = Linear::new(gate_weight, None);
|
||||
|
||||
let experts_vb = vb.pp("experts");
|
||||
let mut experts = Vec::with_capacity(cfg.num_experts);
|
||||
for i in 0..cfg.num_experts {
|
||||
experts.push(
|
||||
Qwen3_5MLP::load_with_dims(&experts_vb.pp(i), h, cfg.moe_intermediate_size)
|
||||
.with_context(|| format!("load expert {i}"))?,
|
||||
);
|
||||
}
|
||||
|
||||
let (shared_expert, shared_expert_gate) = if cfg.shared_expert_intermediate_size > 0 {
|
||||
let shared = Qwen3_5MLP::load_with_dims(
|
||||
&vb.pp("shared_expert"),
|
||||
h,
|
||||
cfg.shared_expert_intermediate_size,
|
||||
)
|
||||
.context("load shared_expert")?;
|
||||
let gate_w = vb
|
||||
.pp("shared_expert_gate")
|
||||
.get((1, h), "weight")
|
||||
.with_context(|| format!("load '{}/shared_expert_gate/weight'", vb.prefix()))?;
|
||||
(Some(shared), Some(Linear::new(gate_w, None)))
|
||||
} else {
|
||||
(None, None)
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
gate,
|
||||
experts,
|
||||
shared_expert,
|
||||
shared_expert_gate,
|
||||
num_experts_per_tok: cfg.num_experts_per_tok,
|
||||
norm_topk_prob: cfg.norm_topk_prob,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-expert routing assignment: `(token_rows, weights)` per expert,
|
||||
/// produced by [`route_scatter`].
|
||||
pub(crate) type ExpertAssignments = (Vec<Vec<u32>>, Vec<Vec<f32>>);
|
||||
|
||||
/// Router + host-side scatter shared by the single-GPU and TP MoE
|
||||
/// blocks (#92): softmax over ALL experts in f32 → descending-argsort
|
||||
/// top-k → renormalise iff `norm_topk_prob` → per-expert token-row and
|
||||
/// weight lists. Under TP the router weight is replicated, so every
|
||||
/// rank computes identical assignments with zero communication.
|
||||
pub(crate) fn route_scatter(
|
||||
gate: &Linear,
|
||||
xs_flat: &Tensor,
|
||||
num_experts: usize,
|
||||
num_experts_per_tok: usize,
|
||||
norm_topk_prob: bool,
|
||||
) -> candle_core::Result<ExpertAssignments> {
|
||||
let n_tokens = xs_flat.dim(0)?;
|
||||
// Router probabilities in f32 (reference uses float softmax
|
||||
// regardless of activations dtype).
|
||||
let router_logits = gate.forward(xs_flat)?;
|
||||
let probs = candle_nn::ops::softmax_last_dim(&router_logits.to_dtype(DType::F32)?)?;
|
||||
|
||||
// Top-k selection: descending argsort, take the first k. The
|
||||
// renormalisation (iff norm_topk_prob) divides by the sum of
|
||||
// the selected global-softmax values.
|
||||
let sorted = probs.arg_sort_last_dim(false)?;
|
||||
let topk_idx = sorted.narrow(1, 0, num_experts_per_tok)?.contiguous()?;
|
||||
let mut topk_w = probs.gather(&topk_idx, 1)?;
|
||||
if norm_topk_prob {
|
||||
let denom = topk_w.sum_keepdim(1)?;
|
||||
topk_w = topk_w.broadcast_div(&denom)?;
|
||||
}
|
||||
|
||||
// Host-side scatter: token row lists per expert. Cheap relative
|
||||
// to the expert GEMMs; replaced by grouped-GEMM in slice 4.
|
||||
let idx_host: Vec<Vec<u32>> = topk_idx.to_vec2()?;
|
||||
let w_host: Vec<Vec<f32>> = topk_w.to_vec2()?;
|
||||
let mut tokens_for: Vec<Vec<u32>> = vec![Vec::new(); num_experts];
|
||||
let mut weights_for: Vec<Vec<f32>> = vec![Vec::new(); num_experts];
|
||||
for t in 0..n_tokens {
|
||||
for j in 0..num_experts_per_tok {
|
||||
let e = idx_host[t][j] as usize;
|
||||
tokens_for[e].push(t as u32);
|
||||
weights_for[e].push(w_host[t][j]);
|
||||
}
|
||||
}
|
||||
Ok((tokens_for, weights_for))
|
||||
}
|
||||
|
||||
impl Module for Qwen3_5MoeBlock {
|
||||
fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||||
let (b, l, hidden) = xs.dims3()?;
|
||||
let xs_flat = xs.reshape(((), hidden))?;
|
||||
|
||||
let (tokens_for, weights_for) = route_scatter(
|
||||
&self.gate,
|
||||
&xs_flat,
|
||||
self.experts.len(),
|
||||
self.num_experts_per_tok,
|
||||
self.norm_topk_prob,
|
||||
)?;
|
||||
|
||||
let mut ys = xs_flat.zeros_like()?;
|
||||
for (e, expert) in self.experts.iter().enumerate() {
|
||||
if tokens_for[e].is_empty() {
|
||||
continue;
|
||||
}
|
||||
let rows = Tensor::new(tokens_for[e].as_slice(), xs.device())?;
|
||||
let picked = xs_flat.index_select(&rows, 0)?;
|
||||
let out = expert.forward(&picked)?;
|
||||
let w = Tensor::new(weights_for[e].as_slice(), xs.device())?
|
||||
.to_dtype(out.dtype())?
|
||||
.reshape(((), 1))?;
|
||||
ys = ys.index_add(&rows, &out.broadcast_mul(&w)?, 0)?;
|
||||
}
|
||||
|
||||
if let (Some(shared), Some(gate)) = (&self.shared_expert, &self.shared_expert_gate) {
|
||||
let mix = candle_nn::ops::sigmoid(&gate.forward(&xs_flat)?)?;
|
||||
let shared_out = shared.forward(&xs_flat)?.broadcast_mul(&mix)?;
|
||||
ys = (ys + shared_out)?;
|
||||
}
|
||||
|
||||
ys.reshape((b, l, hidden))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use candle_core::Device;
|
||||
|
||||
fn randn(shape: &[usize]) -> Tensor {
|
||||
Tensor::randn(0f32, 0.5f32, shape, &Device::Cpu).unwrap()
|
||||
}
|
||||
|
||||
fn rand_mlp(hidden: usize, inter: usize) -> Qwen3_5MLP {
|
||||
Qwen3_5MLP::from_weights(
|
||||
Linear::new(randn(&[inter, hidden]), None),
|
||||
Linear::new(randn(&[inter, hidden]), None),
|
||||
Linear::new(randn(&[hidden, inter]), None),
|
||||
)
|
||||
}
|
||||
|
||||
/// The batched scatter forward must equal a per-token dense
|
||||
/// reference: route each token independently (host softmax → top-k
|
||||
/// → renorm), run its selected experts one by one, and mix in the
|
||||
/// shared expert through the sigmoid gate. Catches indexing,
|
||||
/// weighting, and renormalisation bugs in the scatter path.
|
||||
#[test]
|
||||
fn scatter_forward_matches_per_token_reference() {
|
||||
let (hidden, inter, n_exp, top_k) = (8, 4, 6, 2);
|
||||
|
||||
let block = Qwen3_5MoeBlock {
|
||||
gate: Linear::new(randn(&[n_exp, hidden]), None),
|
||||
experts: (0..n_exp).map(|_| rand_mlp(hidden, inter)).collect(),
|
||||
shared_expert: Some(rand_mlp(hidden, inter)),
|
||||
shared_expert_gate: Some(Linear::new(randn(&[1, hidden]), None)),
|
||||
num_experts_per_tok: top_k,
|
||||
norm_topk_prob: true,
|
||||
};
|
||||
|
||||
let (b, l) = (2, 3);
|
||||
let xs = randn(&[b, l, hidden]);
|
||||
let got = block.forward(&xs).unwrap();
|
||||
assert_eq!(got.dims(), &[b, l, hidden]);
|
||||
|
||||
let xs_flat = xs.reshape(((), hidden)).unwrap();
|
||||
let logits: Vec<Vec<f32>> = block.gate.forward(&xs_flat).unwrap().to_vec2().unwrap();
|
||||
let got_flat: Vec<Vec<f32>> = got.reshape(((), hidden)).unwrap().to_vec2().unwrap();
|
||||
|
||||
for t in 0..b * l {
|
||||
// Host-side softmax over all experts, then top-k + renorm.
|
||||
let max = logits[t].iter().cloned().fold(f32::MIN, f32::max);
|
||||
let exps: Vec<f32> = logits[t].iter().map(|v| (v - max).exp()).collect();
|
||||
let sum: f32 = exps.iter().sum();
|
||||
let probs: Vec<f32> = exps.iter().map(|e| e / sum).collect();
|
||||
let mut order: Vec<usize> = (0..n_exp).collect();
|
||||
order.sort_by(|&a, &b| probs[b].partial_cmp(&probs[a]).unwrap());
|
||||
let selected = &order[..top_k];
|
||||
let denom: f32 = selected.iter().map(|&e| probs[e]).sum();
|
||||
|
||||
let row = xs_flat.narrow(0, t, 1).unwrap();
|
||||
let mut expect = vec![0f32; hidden];
|
||||
for &e in selected {
|
||||
let w = probs[e] / denom;
|
||||
let out: Vec<f32> = block.experts[e]
|
||||
.forward(&row)
|
||||
.unwrap()
|
||||
.flatten_all()
|
||||
.unwrap()
|
||||
.to_vec1()
|
||||
.unwrap();
|
||||
for (acc, o) in expect.iter_mut().zip(out) {
|
||||
*acc += w * o;
|
||||
}
|
||||
}
|
||||
let gate_v: f32 = block
|
||||
.shared_expert_gate
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.forward(&row)
|
||||
.unwrap()
|
||||
.flatten_all()
|
||||
.unwrap()
|
||||
.to_vec1::<f32>()
|
||||
.unwrap()[0];
|
||||
let mix = 1.0 / (1.0 + (-gate_v).exp());
|
||||
let shared: Vec<f32> = block
|
||||
.shared_expert
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.forward(&row)
|
||||
.unwrap()
|
||||
.flatten_all()
|
||||
.unwrap()
|
||||
.to_vec1()
|
||||
.unwrap();
|
||||
for (acc, s) in expect.iter_mut().zip(shared) {
|
||||
*acc += mix * s;
|
||||
}
|
||||
|
||||
for (i, (&g, &e)) in got_flat[t].iter().zip(expect.iter()).enumerate() {
|
||||
assert!(
|
||||
(g - e).abs() < 1e-4,
|
||||
"token {t} dim {i}: got {g}, expected {e}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Without a shared expert (Qwen3-30B-A3B shape) the block is pure
|
||||
/// routed output; without norm_topk_prob the raw global-softmax
|
||||
/// weights apply (they do NOT sum to 1 across the selected k).
|
||||
#[test]
|
||||
fn no_shared_expert_and_no_renorm() {
|
||||
let (hidden, inter, n_exp) = (4, 2, 3);
|
||||
let block = Qwen3_5MoeBlock {
|
||||
gate: Linear::new(randn(&[n_exp, hidden]), None),
|
||||
experts: (0..n_exp).map(|_| rand_mlp(hidden, inter)).collect(),
|
||||
shared_expert: None,
|
||||
shared_expert_gate: None,
|
||||
num_experts_per_tok: 1,
|
||||
norm_topk_prob: false,
|
||||
};
|
||||
let xs = randn(&[1, 1, hidden]);
|
||||
let got: Vec<f32> = block
|
||||
.forward(&xs)
|
||||
.unwrap()
|
||||
.flatten_all()
|
||||
.unwrap()
|
||||
.to_vec1()
|
||||
.unwrap();
|
||||
|
||||
// Reference: the argmax expert's output scaled by its raw
|
||||
// softmax probability.
|
||||
let logits: Vec<f32> = block
|
||||
.gate
|
||||
.forward(&xs.reshape(((), hidden)).unwrap())
|
||||
.unwrap()
|
||||
.flatten_all()
|
||||
.unwrap()
|
||||
.to_vec1()
|
||||
.unwrap();
|
||||
let max = logits.iter().cloned().fold(f32::MIN, f32::max);
|
||||
let exps: Vec<f32> = logits.iter().map(|v| (v - max).exp()).collect();
|
||||
let sum: f32 = exps.iter().sum();
|
||||
let best = (0..n_exp)
|
||||
.max_by(|&a, &b| exps[a].partial_cmp(&exps[b]).unwrap())
|
||||
.unwrap();
|
||||
let w = exps[best] / sum;
|
||||
let out: Vec<f32> = block.experts[best]
|
||||
.forward(&xs.reshape(((), hidden)).unwrap())
|
||||
.unwrap()
|
||||
.flatten_all()
|
||||
.unwrap()
|
||||
.to_vec1()
|
||||
.unwrap();
|
||||
for (i, (&g, &o)) in got.iter().zip(out.iter()).enumerate() {
|
||||
assert!(
|
||||
(g - w * o).abs() < 1e-5,
|
||||
"dim {i}: got {g}, expected {}",
|
||||
w * o
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -159,6 +159,19 @@ impl RotaryEmbedding {
|
||||
Ok((cos, sin))
|
||||
}
|
||||
|
||||
/// cos/sin gathered at arbitrary **per-row** positions — the
|
||||
/// batched-decode path (#98), where each batch row sits at its own
|
||||
/// sequence offset. Shape `(B, 1, rotary_dim/2)`: one position per
|
||||
/// row, one decode token per step. [`Self::apply_cos_sin`] detects
|
||||
/// the rank-3 shape and broadcasts per row instead of per position.
|
||||
pub fn batch_cos_sin(&self, positions: &[usize]) -> candle_core::Result<(Tensor, Tensor)> {
|
||||
let idx: Vec<u32> = positions.iter().map(|&p| p as u32).collect();
|
||||
let idx = Tensor::from_vec(idx, positions.len(), self.cos.device())?;
|
||||
let cos = self.cos.index_select(&idx, 0)?.unsqueeze(1)?;
|
||||
let sin = self.sin.index_select(&idx, 0)?.unsqueeze(1)?;
|
||||
Ok((cos, sin))
|
||||
}
|
||||
|
||||
/// cos/sin from explicit per-token 3D position ids, shape
|
||||
/// `(3, seq_len)` (axes: text, height, width). Builds each axis's
|
||||
/// frequencies and blends them at the interleave index sets, so
|
||||
@@ -185,7 +198,9 @@ impl RotaryEmbedding {
|
||||
}
|
||||
|
||||
/// Apply rotary to `q`, `k` (shape `(B, H, L, head_dim)`) using
|
||||
/// precomputed `cos`/`sin` of shape `(L, rotary_dim/2)`. Partial
|
||||
/// precomputed `cos`/`sin` of shape `(L, rotary_dim/2)` — or, for
|
||||
/// the batched-decode path (#98), `(B, L, rotary_dim/2)` with a
|
||||
/// distinct position per batch row (dispatch is on rank). Partial
|
||||
/// rotary: only the first `rotary_dim` dims rotate; the tail passes
|
||||
/// through unchanged.
|
||||
pub fn apply_cos_sin(
|
||||
@@ -197,9 +212,17 @@ impl RotaryEmbedding {
|
||||
) -> candle_core::Result<(Tensor, Tensor)> {
|
||||
let (_, _, _seq_len, head_dim_in) = q.dims4()?;
|
||||
debug_assert_eq!(head_dim_in, self.head_dim, "q head_dim mismatch");
|
||||
let per_row = cos.rank() == 3;
|
||||
let rope = |x: &Tensor| -> candle_core::Result<Tensor> {
|
||||
if per_row {
|
||||
rope_per_row(x, cos, sin)
|
||||
} else {
|
||||
candle_nn::rotary_emb::rope_slow(x, cos, sin)
|
||||
}
|
||||
};
|
||||
if self.rotary_dim == self.head_dim {
|
||||
let q_embed = candle_nn::rotary_emb::rope_slow(&q.contiguous()?, cos, sin)?;
|
||||
let k_embed = candle_nn::rotary_emb::rope_slow(&k.contiguous()?, cos, sin)?;
|
||||
let q_embed = rope(&q.contiguous()?)?;
|
||||
let k_embed = rope(&k.contiguous()?)?;
|
||||
Ok((q_embed, k_embed))
|
||||
} else {
|
||||
// Partial rotation: narrow → rotate → cat the untouched tail.
|
||||
@@ -212,8 +235,8 @@ impl RotaryEmbedding {
|
||||
.narrow(candle_core::D::Minus1, 0, self.rotary_dim)?
|
||||
.contiguous()?;
|
||||
let k_pass = k.narrow(candle_core::D::Minus1, self.rotary_dim, tail)?;
|
||||
let q_rotated = candle_nn::rotary_emb::rope_slow(&q_rot, cos, sin)?;
|
||||
let k_rotated = candle_nn::rotary_emb::rope_slow(&k_rot, cos, sin)?;
|
||||
let q_rotated = rope(&q_rot)?;
|
||||
let k_rotated = rope(&k_rot)?;
|
||||
let q_embed =
|
||||
Tensor::cat(&[&q_rotated, &q_pass.contiguous()?], candle_core::D::Minus1)?;
|
||||
let k_embed =
|
||||
@@ -223,6 +246,27 @@ impl RotaryEmbedding {
|
||||
}
|
||||
}
|
||||
|
||||
/// GLM rotate-half (same convention as candle's private
|
||||
/// `rotary_emb::rotate_half`: `cat(-x2, x1)`).
|
||||
fn rotate_half(x: &Tensor) -> candle_core::Result<Tensor> {
|
||||
let last = x.dim(candle_core::D::Minus1)?;
|
||||
let x1 = x.narrow(candle_core::D::Minus1, 0, last / 2)?;
|
||||
let x2 = x.narrow(candle_core::D::Minus1, last / 2, last - last / 2)?;
|
||||
Tensor::cat(&[&x2.neg()?, &x1], candle_core::D::Minus1)
|
||||
}
|
||||
|
||||
/// Per-row rope apply for batched decode: `x` is `(B, H, L, rot)`,
|
||||
/// `cos`/`sin` are `(B, L, rot/2)` — each batch row gets its own
|
||||
/// position's rotation (candle's `rope_slow` only broadcasts one
|
||||
/// `(L, rot/2)` table across the whole batch).
|
||||
fn rope_per_row(x: &Tensor, cos: &Tensor, sin: &Tensor) -> candle_core::Result<Tensor> {
|
||||
// (B, L, half) → duplicate pairs → (B, 1, L, rot) for broadcast
|
||||
// over the head dim.
|
||||
let cos = Tensor::cat(&[cos, cos], candle_core::D::Minus1)?.unsqueeze(1)?;
|
||||
let sin = Tensor::cat(&[sin, sin], candle_core::D::Minus1)?.unsqueeze(1)?;
|
||||
x.broadcast_mul(&cos)? + rotate_half(x)?.broadcast_mul(&sin)?
|
||||
}
|
||||
|
||||
/// Compute interleaved-M-RoPE 3D position ids for a full prompt that may
|
||||
/// contain image-placeholder runs, plus the decode `rope_delta`.
|
||||
///
|
||||
@@ -564,6 +608,49 @@ mod tests {
|
||||
assert!((last[1] - (2.0 * inv[1]).cos()).abs() < 1e-5);
|
||||
}
|
||||
|
||||
/// `batch_cos_sin` at positions [5, 9, 0] must gather exactly the
|
||||
/// rows `plain_cos_sin` would produce for each position alone.
|
||||
#[test]
|
||||
fn batch_cos_sin_gathers_per_row_positions() {
|
||||
let dev = Device::Cpu;
|
||||
let rope = RotaryEmbedding::new(DType::F32, &qwen36_cfg(), &dev).unwrap();
|
||||
let half = rope.inv_freq.dim(1).unwrap();
|
||||
let positions = [5usize, 9, 0];
|
||||
let (bc, bs) = rope.batch_cos_sin(&positions).unwrap();
|
||||
assert_eq!(bc.dims(), &[3, 1, half]);
|
||||
assert_eq!(bs.dims(), &[3, 1, half]);
|
||||
for (row, &p) in positions.iter().enumerate() {
|
||||
let (pc, ps) = rope.plain_cos_sin(p, 1).unwrap();
|
||||
let dc = (bc.i(row).unwrap() - pc).unwrap().abs().unwrap();
|
||||
let ds = (bs.i(row).unwrap() - ps).unwrap().abs().unwrap();
|
||||
assert!(dc.max_all().unwrap().to_scalar::<f32>().unwrap() < 1e-6);
|
||||
assert!(ds.max_all().unwrap().to_scalar::<f32>().unwrap() < 1e-6);
|
||||
}
|
||||
}
|
||||
|
||||
/// When every row sits at the same position, the per-row rank-3
|
||||
/// apply path must reproduce the shared rank-2 (`rope_slow`) path
|
||||
/// exactly — the invariant that makes the rank dispatch in
|
||||
/// `apply_cos_sin` safe.
|
||||
#[test]
|
||||
fn per_row_apply_matches_shared_when_uniform() {
|
||||
let dev = Device::Cpu;
|
||||
let rope = RotaryEmbedding::new(DType::F32, &qwen36_cfg(), &dev).unwrap();
|
||||
let q = Tensor::randn(0f32, 1f32, (2, 2, 1, 256), &dev).unwrap();
|
||||
let k = Tensor::randn(0f32, 1f32, (2, 2, 1, 256), &dev).unwrap();
|
||||
|
||||
let (c2, s2) = rope.plain_cos_sin(7, 1).unwrap();
|
||||
let (qa, ka) = rope.apply_cos_sin(&q, &k, &c2, &s2).unwrap();
|
||||
|
||||
let (c3, s3) = rope.batch_cos_sin(&[7, 7]).unwrap();
|
||||
let (qb, kb) = rope.apply_cos_sin(&q, &k, &c3, &s3).unwrap();
|
||||
|
||||
let dq = (qa - qb).unwrap().abs().unwrap().max_all().unwrap();
|
||||
let dk = (ka - kb).unwrap().abs().unwrap().max_all().unwrap();
|
||||
assert!(dq.to_scalar::<f32>().unwrap() < 1e-6, "q mismatch {dq:?}");
|
||||
assert!(dk.to_scalar::<f32>().unwrap() < 1e-6, "k mismatch {dk:?}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_rope_index_196_is_14x14() {
|
||||
let mut ids = vec![1u32]; // one text token
|
||||
|
||||
@@ -84,6 +84,232 @@ impl KvCacheSnapshot {
|
||||
}
|
||||
}
|
||||
|
||||
/// Batched cache state assembled from per-sequence snapshots (#98).
|
||||
/// Install with `Qwen3_5Model::restore_kv_cache(&self.snapshot)`; the
|
||||
/// forward then runs lockstep batched decode with
|
||||
/// `forward_batch_decode` using `prefix_lens[i] + step` positions and
|
||||
/// the padding mask from `Qwen3_5Model::batch_decode_mask`.
|
||||
pub struct BatchedKvState {
|
||||
/// Per-layer `(B, …)` state: attention K/V right-padded along the
|
||||
/// sequence axis to `padded_len` and `cat`ed on dim 0; GDN
|
||||
/// conv/recurrent states `cat`ed on dim 0 (position-free).
|
||||
pub snapshot: KvCacheSnapshot,
|
||||
/// The uniform KV sequence length every row was padded to — the
|
||||
/// max prefix length in the batch. Decode appends start here.
|
||||
pub padded_len: usize,
|
||||
/// Each row's true prefix length. Columns `[prefix_lens[i],
|
||||
/// padded_len)` of row `i` are zero-padding and must stay masked.
|
||||
pub prefix_lens: Vec<usize>,
|
||||
}
|
||||
|
||||
/// Assemble per-sequence snapshots into one batched cache state.
|
||||
/// `seqs` pairs each snapshot with its true token length (the caller
|
||||
/// tracks prompt token counts; attention K/V lengths are validated
|
||||
/// against it). All snapshots must come from single-sequence (`B=1`)
|
||||
/// prefills of the same model, with `rope_delta == 0` (text-only —
|
||||
/// vision requests don't batch, #98 v1).
|
||||
///
|
||||
/// Keys are stored post-RoPE, so right-padding does not disturb
|
||||
/// position correctness: a row's cached keys keep the rotation of
|
||||
/// their true positions, the garbage columns are masked, and new
|
||||
/// tokens rotate at `prefix_len + step` while landing at storage
|
||||
/// column `padded_len + step`.
|
||||
pub fn assemble_batch(seqs: &[(&KvCacheSnapshot, usize)]) -> candle_core::Result<BatchedKvState> {
|
||||
let Some((first, _)) = seqs.first() else {
|
||||
candle_core::bail!("assemble_batch: empty batch");
|
||||
};
|
||||
let n_layers = first.layers.len();
|
||||
let prefix_lens: Vec<usize> = seqs.iter().map(|&(_, len)| len).collect();
|
||||
let padded_len = *prefix_lens.iter().max().expect("non-empty");
|
||||
for (snap, len) in seqs {
|
||||
if snap.layers.len() != n_layers {
|
||||
candle_core::bail!(
|
||||
"assemble_batch: snapshot layer count mismatch ({} vs {n_layers})",
|
||||
snap.layers.len()
|
||||
);
|
||||
}
|
||||
if snap.rope_delta != 0 {
|
||||
candle_core::bail!(
|
||||
"assemble_batch: rope_delta {} != 0 — vision-positioned sequences cannot batch",
|
||||
snap.rope_delta
|
||||
);
|
||||
}
|
||||
if *len == 0 {
|
||||
candle_core::bail!("assemble_batch: zero-length sequence");
|
||||
}
|
||||
}
|
||||
|
||||
let mut layers = Vec::with_capacity(n_layers);
|
||||
for li in 0..n_layers {
|
||||
layers.push(assemble_layer(seqs, li, padded_len)?);
|
||||
}
|
||||
Ok(BatchedKvState {
|
||||
snapshot: KvCacheSnapshot {
|
||||
layers,
|
||||
rope_delta: 0,
|
||||
},
|
||||
padded_len,
|
||||
prefix_lens,
|
||||
})
|
||||
}
|
||||
|
||||
fn assemble_layer(
|
||||
seqs: &[(&KvCacheSnapshot, usize)],
|
||||
li: usize,
|
||||
padded_len: usize,
|
||||
) -> candle_core::Result<LayerKvSnapshot> {
|
||||
match &seqs[0].0.layers[li] {
|
||||
LayerKvSnapshot::Full(_) => {
|
||||
let mut ks = Vec::with_capacity(seqs.len());
|
||||
let mut vs = Vec::with_capacity(seqs.len());
|
||||
for (row, (snap, len)) in seqs.iter().enumerate() {
|
||||
let LayerKvSnapshot::Full(Some((k, v))) = &snap.layers[li] else {
|
||||
candle_core::bail!(
|
||||
"assemble_batch: row {row} layer {li} is not a populated \
|
||||
full-attention snapshot"
|
||||
);
|
||||
};
|
||||
let (b, _h, s, _d) = k.dims4()?;
|
||||
if b != 1 {
|
||||
candle_core::bail!(
|
||||
"assemble_batch: row {row} layer {li} has batch dim {b}, want 1"
|
||||
);
|
||||
}
|
||||
if s != *len {
|
||||
candle_core::bail!(
|
||||
"assemble_batch: row {row} layer {li} KV length {s} != declared \
|
||||
sequence length {len}"
|
||||
);
|
||||
}
|
||||
ks.push(pad_seq(k, padded_len)?);
|
||||
vs.push(pad_seq(v, padded_len)?);
|
||||
}
|
||||
let k = Tensor::cat(&ks, 0)?;
|
||||
let v = Tensor::cat(&vs, 0)?;
|
||||
Ok(LayerKvSnapshot::Full(Some((k, v))))
|
||||
}
|
||||
LayerKvSnapshot::Linear { .. } => {
|
||||
let mut convs = Vec::with_capacity(seqs.len());
|
||||
let mut recs = Vec::with_capacity(seqs.len());
|
||||
for (row, (snap, _)) in seqs.iter().enumerate() {
|
||||
let LayerKvSnapshot::Linear {
|
||||
conv_state: Some(conv),
|
||||
recurrent_state: Some(rec),
|
||||
} = &snap.layers[li]
|
||||
else {
|
||||
candle_core::bail!(
|
||||
"assemble_batch: row {row} layer {li} is not a populated \
|
||||
linear-attention snapshot"
|
||||
);
|
||||
};
|
||||
convs.push(conv.clone());
|
||||
recs.push(rec.clone());
|
||||
}
|
||||
Ok(LayerKvSnapshot::Linear {
|
||||
conv_state: Some(Tensor::cat(&convs, 0)?),
|
||||
recurrent_state: Some(Tensor::cat(&recs, 0)?),
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract one row of a batched cache snapshot back into a contiguous
|
||||
/// single-sequence snapshot — the defragment half of batch membership
|
||||
/// changes (#98). A join/leave rebatches by extracting every surviving
|
||||
/// row (each becomes a gap-free `B=1` snapshot of length `prefix_len +
|
||||
/// steps`) and running [`assemble_batch`] over them again, so the
|
||||
/// "each row has exactly one padding gap" invariant holds for the new
|
||||
/// batch too.
|
||||
///
|
||||
/// `row` indexes the batch dim. The row's valid attention KV is its
|
||||
/// prefix `[0, prefix_len)` plus the lockstep decode columns
|
||||
/// `[padded_len, padded_len + steps)`; the gap between them is padding
|
||||
/// and is dropped. GDN states are position-free — the row is sliced
|
||||
/// out whole and **deep-copied** (the live buffers are mutated in
|
||||
/// place by the CUDA kernels; see the module notes on copy semantics).
|
||||
pub fn extract_row(
|
||||
snap: &KvCacheSnapshot,
|
||||
row: usize,
|
||||
prefix_len: usize,
|
||||
padded_len: usize,
|
||||
steps: usize,
|
||||
) -> candle_core::Result<KvCacheSnapshot> {
|
||||
if prefix_len == 0 || prefix_len > padded_len {
|
||||
candle_core::bail!(
|
||||
"extract_row: prefix_len {prefix_len} out of range (padded_len {padded_len})"
|
||||
);
|
||||
}
|
||||
let mut layers = Vec::with_capacity(snap.layers.len());
|
||||
for (li, layer) in snap.layers.iter().enumerate() {
|
||||
layers.push(match layer {
|
||||
LayerKvSnapshot::Full(Some((k, v))) => {
|
||||
let (b, _h, s, _d) = k.dims4()?;
|
||||
if row >= b {
|
||||
candle_core::bail!("extract_row: row {row} out of range (batch {b})");
|
||||
}
|
||||
if s != padded_len + steps {
|
||||
candle_core::bail!(
|
||||
"extract_row: layer {li} KV length {s} != padded_len {padded_len} + \
|
||||
steps {steps}"
|
||||
);
|
||||
}
|
||||
LayerKvSnapshot::Full(Some((
|
||||
unpad_row(k, row, prefix_len, padded_len, steps)?,
|
||||
unpad_row(v, row, prefix_len, padded_len, steps)?,
|
||||
)))
|
||||
}
|
||||
LayerKvSnapshot::Full(None) => {
|
||||
candle_core::bail!("extract_row: layer {li} has an empty attention cache")
|
||||
}
|
||||
LayerKvSnapshot::Linear {
|
||||
conv_state: Some(conv),
|
||||
recurrent_state: Some(rec),
|
||||
} => LayerKvSnapshot::Linear {
|
||||
conv_state: Some(conv.narrow(0, row, 1)?.copy()?),
|
||||
recurrent_state: Some(rec.narrow(0, row, 1)?.copy()?),
|
||||
},
|
||||
LayerKvSnapshot::Linear { .. } => {
|
||||
candle_core::bail!("extract_row: layer {li} has unpopulated GDN state")
|
||||
}
|
||||
});
|
||||
}
|
||||
Ok(KvCacheSnapshot {
|
||||
layers,
|
||||
rope_delta: 0,
|
||||
})
|
||||
}
|
||||
|
||||
/// Slice `row` out of a batched `(B, H, padded_len + steps, D)` K or V
|
||||
/// tensor and drop its padding gap, yielding an owned contiguous
|
||||
/// `(1, H, prefix_len + steps, D)` tensor.
|
||||
fn unpad_row(
|
||||
t: &Tensor,
|
||||
row: usize,
|
||||
prefix_len: usize,
|
||||
padded_len: usize,
|
||||
steps: usize,
|
||||
) -> candle_core::Result<Tensor> {
|
||||
let r = t.narrow(0, row, 1)?;
|
||||
let prefix = r.narrow(2, 0, prefix_len)?;
|
||||
if steps == 0 {
|
||||
return prefix.contiguous()?.copy();
|
||||
}
|
||||
let decoded = r.narrow(2, padded_len, steps)?;
|
||||
Tensor::cat(&[&prefix.contiguous()?, &decoded.contiguous()?], 2)
|
||||
}
|
||||
|
||||
/// Right-pad a `(1, H, S, D)` K or V tensor with zeros along the
|
||||
/// sequence axis to `padded_len`. Zero columns are inert: the padding
|
||||
/// mask keeps every query from attending to them.
|
||||
fn pad_seq(t: &Tensor, padded_len: usize) -> candle_core::Result<Tensor> {
|
||||
let (b, h, s, d) = t.dims4()?;
|
||||
if s == padded_len {
|
||||
return Ok(t.clone());
|
||||
}
|
||||
let pad = Tensor::zeros((b, h, padded_len - s, d), t.dtype(), t.device())?;
|
||||
Tensor::cat(&[t, &pad], 2)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::super::{Qwen3_5Model, RopeParameters, TextConfig};
|
||||
@@ -119,6 +345,13 @@ mod tests {
|
||||
linear_key_head_dim: 4,
|
||||
linear_value_head_dim: 4,
|
||||
linear_conv_kernel_dim: 4,
|
||||
num_experts: 0,
|
||||
num_experts_per_tok: 0,
|
||||
moe_intermediate_size: 0,
|
||||
shared_expert_intermediate_size: 0,
|
||||
decoder_sparse_step: 1,
|
||||
mlp_only_layers: Vec::new(),
|
||||
norm_topk_prob: false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -201,7 +434,7 @@ mod tests {
|
||||
)
|
||||
.expect("build ShardedVarBuilder")
|
||||
};
|
||||
Qwen3_5Model::load(cfg, &vb).expect("load tiny qwen3_5 model")
|
||||
Qwen3_5Model::load(cfg, &vb, "model.language_model").expect("load tiny qwen3_5 model")
|
||||
}
|
||||
|
||||
fn forward_tokens(model: &mut Qwen3_5Model, tokens: &[u32], offset: usize) -> Vec<f32> {
|
||||
@@ -271,6 +504,228 @@ mod tests {
|
||||
assert!(diff < 1e-6, "second restore diverged: {diff}");
|
||||
}
|
||||
|
||||
/// The gold test for #98 slice 1: ragged sequences decoded in one
|
||||
/// lockstep batch (assembled from per-sequence snapshots, per-row
|
||||
/// positions, padding mask) must match the same sequences decoded
|
||||
/// sequentially, hidden-state for hidden-state at every step.
|
||||
#[test]
|
||||
fn batched_decode_matches_sequential() {
|
||||
use candle_core::IndexOp;
|
||||
let cfg = tiny_config();
|
||||
let mut model = tiny_model(&cfg);
|
||||
|
||||
let prompts: [&[u32]; 3] = [&[1, 2, 3], &[4, 5], &[7, 7, 2, 5, 6]];
|
||||
let steps: [&[u32]; 3] = [&[11, 12, 13, 14], &[9, 8, 7, 6], &[21, 22, 23, 24]];
|
||||
let n_steps = 4;
|
||||
|
||||
// Sequential reference: each sequence decoded alone.
|
||||
let mut expected: Vec<Vec<Vec<f32>>> = Vec::new(); // [row][step]
|
||||
for (prompt, toks) in prompts.iter().zip(steps.iter()) {
|
||||
model.clear_kv_cache();
|
||||
forward_tokens(&mut model, prompt, 0);
|
||||
let mut per_step = Vec::new();
|
||||
for (t, tok) in toks.iter().enumerate() {
|
||||
per_step.push(forward_tokens(&mut model, &[*tok], prompt.len() + t));
|
||||
}
|
||||
expected.push(per_step);
|
||||
}
|
||||
|
||||
// Batched: prefill each sequence alone, snapshot, assemble.
|
||||
let mut snaps = Vec::new();
|
||||
for prompt in prompts.iter() {
|
||||
model.clear_kv_cache();
|
||||
forward_tokens(&mut model, prompt, 0);
|
||||
snaps.push(model.snapshot_kv_cache().expect("snapshot"));
|
||||
}
|
||||
let seqs: Vec<(&super::KvCacheSnapshot, usize)> = snaps
|
||||
.iter()
|
||||
.zip(prompts.iter())
|
||||
.map(|(s, p)| (s, p.len()))
|
||||
.collect();
|
||||
let batch = super::assemble_batch(&seqs).expect("assemble");
|
||||
assert_eq!(batch.padded_len, 5);
|
||||
assert_eq!(batch.prefix_lens, vec![3, 2, 5]);
|
||||
model
|
||||
.restore_kv_cache(&batch.snapshot)
|
||||
.expect("install batched state");
|
||||
|
||||
for t in 0..n_steps {
|
||||
let toks: Vec<u32> = steps.iter().map(|s| s[t]).collect();
|
||||
let input = Tensor::from_vec(toks, (3, 1), &Device::Cpu).unwrap();
|
||||
let positions: Vec<usize> = prompts.iter().map(|p| p.len() + t).collect();
|
||||
let total_len = batch.padded_len + t + 1;
|
||||
let mask = model
|
||||
.batch_decode_mask(&batch.prefix_lens, batch.padded_len, total_len)
|
||||
.expect("mask");
|
||||
assert!(mask.is_some(), "ragged batch must be masked");
|
||||
let h = model
|
||||
.forward_batch_decode(&input, &positions, mask.as_ref())
|
||||
.expect("batched step");
|
||||
assert_eq!(h.dims()[0], 3);
|
||||
for row in 0..3 {
|
||||
let got: Vec<f32> = h.i((row, 0, ..)).unwrap().to_vec1().unwrap();
|
||||
let diff = max_abs_diff(&expected[row][t], &got);
|
||||
assert!(diff < 1e-4, "row {row} step {t} diverged: {diff}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Uniform-length batch: no padding → `batch_decode_mask` returns
|
||||
/// `None`, and unmasked lockstep decode still matches sequential.
|
||||
#[test]
|
||||
fn batched_decode_uniform_lengths_needs_no_mask() {
|
||||
use candle_core::IndexOp;
|
||||
let cfg = tiny_config();
|
||||
let mut model = tiny_model(&cfg);
|
||||
|
||||
let prompts: [&[u32]; 2] = [&[1, 2, 3], &[6, 5, 4]];
|
||||
let toks = [13u32, 17];
|
||||
|
||||
let mut expected = Vec::new();
|
||||
for (prompt, tok) in prompts.iter().zip(toks.iter()) {
|
||||
model.clear_kv_cache();
|
||||
forward_tokens(&mut model, prompt, 0);
|
||||
expected.push(forward_tokens(&mut model, &[*tok], prompt.len()));
|
||||
}
|
||||
|
||||
let mut snaps = Vec::new();
|
||||
for prompt in prompts.iter() {
|
||||
model.clear_kv_cache();
|
||||
forward_tokens(&mut model, prompt, 0);
|
||||
snaps.push(model.snapshot_kv_cache().expect("snapshot"));
|
||||
}
|
||||
let seqs: Vec<(&super::KvCacheSnapshot, usize)> = snaps
|
||||
.iter()
|
||||
.zip(prompts.iter())
|
||||
.map(|(s, p)| (s, p.len()))
|
||||
.collect();
|
||||
let batch = super::assemble_batch(&seqs).expect("assemble");
|
||||
let mask = model
|
||||
.batch_decode_mask(&batch.prefix_lens, batch.padded_len, batch.padded_len + 1)
|
||||
.expect("mask");
|
||||
assert!(mask.is_none(), "uniform lengths must not build a mask");
|
||||
model.restore_kv_cache(&batch.snapshot).expect("install");
|
||||
|
||||
let input = Tensor::from_vec(toks.to_vec(), (2, 1), &Device::Cpu).unwrap();
|
||||
let h = model
|
||||
.forward_batch_decode(&input, &[3, 3], None)
|
||||
.expect("step");
|
||||
for row in 0..2 {
|
||||
let got: Vec<f32> = h.i((row, 0, ..)).unwrap().to_vec1().unwrap();
|
||||
let diff = max_abs_diff(&expected[row], &got);
|
||||
assert!(diff < 1e-4, "row {row} diverged: {diff}");
|
||||
}
|
||||
}
|
||||
|
||||
/// Round-trip for batch membership changes (#98): decode two steps
|
||||
/// in a lockstep batch, extract each row back to a contiguous
|
||||
/// single-sequence snapshot, restore it alone, and continue
|
||||
/// decoding at B=1 — the continuation must match a pure-sequential
|
||||
/// run of the same tokens. This is the primitive a join/leave
|
||||
/// rebatch is built from.
|
||||
#[test]
|
||||
fn extract_row_continues_like_sequential() {
|
||||
use candle_core::IndexOp;
|
||||
let cfg = tiny_config();
|
||||
let mut model = tiny_model(&cfg);
|
||||
|
||||
let prompts: [&[u32]; 2] = [&[1, 2, 3], &[4, 5]];
|
||||
let toks: [&[u32]; 2] = [&[11, 12, 13, 14], &[9, 8, 7, 6]];
|
||||
let batched_steps = 2; // decoded in the batch
|
||||
let solo_steps = 2; // decoded after extraction, B=1
|
||||
|
||||
// Pure-sequential reference over all four steps.
|
||||
let mut expected: Vec<Vec<Vec<f32>>> = Vec::new();
|
||||
for (prompt, t) in prompts.iter().zip(toks.iter()) {
|
||||
model.clear_kv_cache();
|
||||
forward_tokens(&mut model, prompt, 0);
|
||||
let mut per_step = Vec::new();
|
||||
for (i, tok) in t.iter().enumerate() {
|
||||
per_step.push(forward_tokens(&mut model, &[*tok], prompt.len() + i));
|
||||
}
|
||||
expected.push(per_step);
|
||||
}
|
||||
|
||||
// Prefill + assemble + two lockstep batched steps.
|
||||
let mut snaps = Vec::new();
|
||||
for prompt in prompts.iter() {
|
||||
model.clear_kv_cache();
|
||||
forward_tokens(&mut model, prompt, 0);
|
||||
snaps.push(model.snapshot_kv_cache().expect("snapshot"));
|
||||
}
|
||||
let seqs: Vec<(&super::KvCacheSnapshot, usize)> = snaps
|
||||
.iter()
|
||||
.zip(prompts.iter())
|
||||
.map(|(s, p)| (s, p.len()))
|
||||
.collect();
|
||||
let batch = super::assemble_batch(&seqs).expect("assemble");
|
||||
model.restore_kv_cache(&batch.snapshot).expect("install");
|
||||
for t in 0..batched_steps {
|
||||
let step_toks: Vec<u32> = toks.iter().map(|s| s[t]).collect();
|
||||
let input = Tensor::from_vec(step_toks, (2, 1), &Device::Cpu).unwrap();
|
||||
let positions: Vec<usize> = prompts.iter().map(|p| p.len() + t).collect();
|
||||
let mask = model
|
||||
.batch_decode_mask(
|
||||
&batch.prefix_lens,
|
||||
batch.padded_len,
|
||||
batch.padded_len + t + 1,
|
||||
)
|
||||
.expect("mask");
|
||||
let h = model
|
||||
.forward_batch_decode(&input, &positions, mask.as_ref())
|
||||
.expect("batched step");
|
||||
for row in 0..2 {
|
||||
let got: Vec<f32> = h.i((row, 0, ..)).unwrap().to_vec1().unwrap();
|
||||
let diff = max_abs_diff(&expected[row][t], &got);
|
||||
assert!(diff < 1e-4, "batched row {row} step {t}: {diff}");
|
||||
}
|
||||
}
|
||||
|
||||
// Extract each row from the live batched state, restore it
|
||||
// alone, and continue at B=1.
|
||||
let live = model.snapshot_kv_cache().expect("snapshot live batch");
|
||||
for row in 0..2 {
|
||||
let solo = super::extract_row(
|
||||
&live,
|
||||
row,
|
||||
prompts[row].len(),
|
||||
batch.padded_len,
|
||||
batched_steps,
|
||||
)
|
||||
.expect("extract row");
|
||||
model.restore_kv_cache(&solo).expect("restore solo");
|
||||
for i in 0..solo_steps {
|
||||
let t = batched_steps + i;
|
||||
let got = forward_tokens(&mut model, &[toks[row][t]], prompts[row].len() + t);
|
||||
let diff = max_abs_diff(&expected[row][t], &got);
|
||||
assert!(diff < 1e-4, "solo row {row} step {t}: {diff}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Mask geometry: `-inf` exactly on `[prefix_len, padded_len)` per
|
||||
/// row, zero elsewhere (including the decode columns past
|
||||
/// `padded_len`).
|
||||
#[test]
|
||||
fn batch_decode_mask_covers_only_padding_gap() {
|
||||
let model = tiny_model(&tiny_config());
|
||||
let m = model
|
||||
.batch_decode_mask(&[3, 5], 5, 7)
|
||||
.unwrap()
|
||||
.expect("ragged → mask");
|
||||
assert_eq!(m.dims(), &[2, 1, 1, 7]);
|
||||
let flat: Vec<f32> = m.flatten_all().unwrap().to_vec1().unwrap();
|
||||
let (row0, row1) = flat.split_at(7);
|
||||
for (j, &v) in row0.iter().enumerate() {
|
||||
if (3..5).contains(&j) {
|
||||
assert_eq!(v, f32::NEG_INFINITY, "row0 col {j} must be masked");
|
||||
} else {
|
||||
assert_eq!(v, 0.0, "row0 col {j} must be open");
|
||||
}
|
||||
}
|
||||
assert!(row1.iter().all(|&v| v == 0.0), "unpadded row must be open");
|
||||
}
|
||||
|
||||
/// 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.
|
||||
|
||||
@@ -23,11 +23,11 @@ use candle_transformers::models::qwen3_moe as qwen3_moe_dense;
|
||||
use cortex_core::harness::{Harness, HarnessHealth, ModelInfo, ModelSpec};
|
||||
use cortex_core::openai::{
|
||||
ChatCompletionChoice, ChatCompletionChunk, ChatCompletionRequest, ChatCompletionResponse,
|
||||
ChatMessage, MessageContent, Usage,
|
||||
ChatMessage, CompletionTokensDetails, MessageContent, Usage,
|
||||
};
|
||||
|
||||
use crate::wire::{
|
||||
FinishReason, InferenceEvent, ReasoningTokenPair, ToolCallTokenPair,
|
||||
FinishReason, FinishTiming, InferenceEvent, ReasoningTokenPair, ToolCallTokenPair,
|
||||
detect_reasoning_token_pair, detect_tool_call_token_pair, openai_chat as wire_chat,
|
||||
};
|
||||
use std::collections::HashMap;
|
||||
@@ -320,7 +320,7 @@ pub struct LoadedModel {
|
||||
/// error so an operator knows to unload+reload to recover. See
|
||||
/// the 2026-05-26 beast incident where a 14k-token prefill OOM
|
||||
/// silently turned every subsequent request into a stuck wait.
|
||||
pub poisoned: AtomicBool,
|
||||
pub poisoned: Arc<AtomicBool>,
|
||||
/// Handle to the per-device CUDA worker thread for this model's
|
||||
/// device. `None` for CPU loads (no context to own). VRAM queries
|
||||
/// and — for CUDA loads — forward / kv-cache / drop ops route
|
||||
@@ -344,7 +344,7 @@ pub struct LoadedModel {
|
||||
/// shape-mismatch failure mid-prefill. Mirrors TpLoadedModel.pool
|
||||
/// for the TP path (which already had this invariant by accident
|
||||
/// because the pool lock covered the same window).
|
||||
pub inference_lock: tokio::sync::Mutex<()>,
|
||||
pub inference_lock: Arc<tokio::sync::Mutex<()>>,
|
||||
/// Bounded admission scheduler (#53). Gated *before* `inference_lock`
|
||||
/// so a busy model refuses overflow fast instead of growing an
|
||||
/// unbounded, untimed queue of lock waiters.
|
||||
@@ -401,7 +401,7 @@ pub struct LoadedModel {
|
||||
/// the pre-#11 behaviour. Dropped with the model, so unload and
|
||||
/// auto-recovery invalidate every entry for free (the worker-side
|
||||
/// snapshots go with `Job::DropArch`).
|
||||
pub prefix_cache: Option<ModelPrefixCache>,
|
||||
pub prefix_cache: Option<Arc<ModelPrefixCache>>,
|
||||
/// Context-limit physics (#67), captured at load. `None` for arches
|
||||
/// whose KV layout we don't yet introspect (GGUF/CPU/non-qwen3_5) —
|
||||
/// those fall back to the static prompt cap with no advertised limit.
|
||||
@@ -412,7 +412,7 @@ pub struct LoadedModel {
|
||||
/// at the end of each streaming request's prefill phase. Feeds the
|
||||
/// throughput ceiling in the derived limit; falls back to the
|
||||
/// configured bootstrap estimate before the first sample.
|
||||
pub prefill_rate: super::context_limit::PrefillRateEma,
|
||||
pub prefill_rate: Arc<super::context_limit::PrefillRateEma>,
|
||||
/// Last derived input-token cap (#67), refreshed each time
|
||||
/// `derived_limit` runs (i.e. on every `/models` poll). The
|
||||
/// request-path enforcement reads this — `0` means "not derived yet"
|
||||
@@ -425,6 +425,14 @@ pub struct LoadedModel {
|
||||
/// cortex's health poller into marking the node unhealthy. Refreshed off
|
||||
/// the request path: seeded at load, then by a background task.
|
||||
pub last_free_mb: AtomicU64,
|
||||
/// Lockstep batched decode engine (#98). `Some` when the operator
|
||||
/// raised `[admission] max_in_flight` above 1 on a snapshot-capable
|
||||
/// worker-path model (and `NEURON_BATCHING` isn't 0). Text chat
|
||||
/// streams route through it instead of taking `inference_lock` per
|
||||
/// request; the engine holds the lock while it has active slots, so
|
||||
/// vision and non-streaming requests still serialize safely against
|
||||
/// the batch.
|
||||
pub engine: Option<super::engine::EngineHandle>,
|
||||
}
|
||||
|
||||
impl LoadedModel {
|
||||
@@ -679,6 +687,41 @@ impl ModelArch {
|
||||
}
|
||||
}
|
||||
|
||||
/// One lockstep batched decode step (#98): `(B, 1)` input, per-row
|
||||
/// positions, optional padding mask. Returns `(B, 1, vocab)` — the
|
||||
/// caller extracts one logits row per batch row (no
|
||||
/// `squeeze_to_vocab`, which would collapse the batch dim). Only
|
||||
/// the qwen3_5 arch batches; the engine only forms batches where
|
||||
/// [`Self::supports_kv_snapshot`] holds, so other archs erroring
|
||||
/// here is defence in depth.
|
||||
pub fn forward_batch_decode(
|
||||
&mut self,
|
||||
input: &Tensor,
|
||||
positions: &[usize],
|
||||
attn_mask: Option<&Tensor>,
|
||||
) -> Result<Tensor> {
|
||||
match self {
|
||||
ModelArch::Qwen3_5Dense(m) => Ok(m.forward_batch_decode(input, positions, attn_mask)?),
|
||||
_ => anyhow::bail!("forward_batch_decode: architecture has no batched-decode support"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Padding mask for a batched decode step — see
|
||||
/// `Qwen3_5Model::batch_decode_mask`.
|
||||
pub fn batch_decode_mask(
|
||||
&self,
|
||||
prefix_lens: &[usize],
|
||||
padded_len: usize,
|
||||
total_len: usize,
|
||||
) -> Result<Option<Tensor>> {
|
||||
match self {
|
||||
ModelArch::Qwen3_5Dense(m) => {
|
||||
Ok(m.batch_decode_mask(prefix_lens, padded_len, total_len)?)
|
||||
}
|
||||
_ => anyhow::bail!("batch_decode_mask: architecture has no batched-decode support"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Forward step that splices vision-tower output at
|
||||
/// `<|image_pad|>` token positions. Stage B2.
|
||||
///
|
||||
@@ -784,6 +827,48 @@ impl ModelArch {
|
||||
}
|
||||
}
|
||||
|
||||
/// Split a non-streaming completion's generated tokens into the
|
||||
/// visible answer and the leading reasoning span.
|
||||
///
|
||||
/// Reasoning models (Qwen3 `<think>`, DeepSeek-R1, …) emit their
|
||||
/// chain-of-thought *before* the answer, and the chat template injects
|
||||
/// the **opening** marker into the prompt — so the generated tokens look
|
||||
/// like `…reasoning… </think> …answer…` with no opening marker present
|
||||
/// in the output. The streaming path drops reasoning as
|
||||
/// [`InferenceEvent::ReasoningDelta`]; the non-streaming path has to do
|
||||
/// the equivalent post-hoc or the chain-of-thought leaks into the
|
||||
/// assistant `content` (which broke agent-zero v2.0, whose parser
|
||||
/// expected the bare JSON answer, not a `<think>` preamble).
|
||||
///
|
||||
/// Returns `(content_ids, reasoning_token_count)`. Strategy: if the model
|
||||
/// declares a reasoning marker pair and its **close** token appears in
|
||||
/// `generated_ids`, everything up to and including the last close token is
|
||||
/// reasoning and only the tail is the answer. When no close token was
|
||||
/// generated, `prompt_opened` decides (#112): the chat template may have
|
||||
/// force-opened the think block inside the generation prompt
|
||||
/// (Qwen3-Next-80B-A3B-Thinking ends its prompt with
|
||||
/// `<|im_start|>assistant\n<think>\n`), in which case a generation
|
||||
/// truncated mid-reasoning is ALL reasoning and the visible answer is
|
||||
/// empty. Otherwise (non-reasoning model, thinking disabled) the tokens
|
||||
/// are returned unchanged. Splitting on the token id — not a decoded
|
||||
/// `</think>` string — keeps this robust against tokenizer byte-fallback
|
||||
/// and special-token handling.
|
||||
fn split_off_reasoning<'a>(
|
||||
generated_ids: &'a [u32],
|
||||
reasoning: Option<&ReasoningTokenPair>,
|
||||
prompt_opened: bool,
|
||||
) -> (&'a [u32], u64) {
|
||||
if let Some(pair) = reasoning {
|
||||
if let Some(idx) = generated_ids.iter().rposition(|&t| t == pair.close_id) {
|
||||
return (&generated_ids[idx + 1..], (idx + 1) as u64);
|
||||
}
|
||||
if prompt_opened {
|
||||
return (&[], generated_ids.len() as u64);
|
||||
}
|
||||
}
|
||||
(generated_ids, 0)
|
||||
}
|
||||
|
||||
/// Squeeze any leading singleton dims off the logits tensor so the
|
||||
/// caller gets a rank-1 `[vocab_size]` slice ready for sampling. Bails
|
||||
/// on a non-singleton leading dim (would mean a batched forward, which
|
||||
@@ -861,7 +946,8 @@ const REPEAT_LAST_N: usize = 64;
|
||||
/// value. New entries land alongside a new `ModelArch` variant + a
|
||||
/// dispatch branch in `load_arch_dense` (plus, for TP, a parallel
|
||||
/// pattern in `tp_qwen3.rs`).
|
||||
const DENSE_SUPPORTED_MODEL_TYPES: &[&str] = &["llama", "qwen3", "qwen3_5", "qwen3_moe"];
|
||||
const DENSE_SUPPORTED_MODEL_TYPES: &[&str] =
|
||||
&["llama", "qwen3", "qwen3_5", "qwen3_moe", "qwen3_next"];
|
||||
|
||||
/// Pre-flight check the operator's `config.json` against the set of
|
||||
/// architectures the dense path actually knows how to build. Surfaces
|
||||
@@ -916,7 +1002,7 @@ pub(crate) fn check_dense_config_supported(config_json: &str, model_id: &str) ->
|
||||
/// families than the TP path because each TP-aware module is a real
|
||||
/// chunk of work (`tp_qwen3.rs` is the only one shipped today).
|
||||
#[cfg(feature = "cuda")]
|
||||
const TP_SUPPORTED_MODEL_TYPES: &[&str] = &["qwen3", "qwen3_5"];
|
||||
const TP_SUPPORTED_MODEL_TYPES: &[&str] = &["qwen3", "qwen3_5", "qwen3_next"];
|
||||
|
||||
/// TP-side counterpart to `check_dense_config_supported`. Gates the
|
||||
/// `load_tp` path on a narrower architecture set: even though the
|
||||
@@ -987,7 +1073,7 @@ fn resolve_hf_cache(explicit: Option<PathBuf>) -> Option<PathBuf> {
|
||||
/// paid at most once per poisoned model.
|
||||
#[derive(Debug)]
|
||||
#[allow(dead_code)]
|
||||
struct LogitsHealth {
|
||||
pub(crate) struct LogitsHealth {
|
||||
len: usize,
|
||||
nan: usize,
|
||||
pos_inf: usize,
|
||||
@@ -1028,7 +1114,7 @@ fn logits_health(t: &Tensor) -> LogitsHealth {
|
||||
/// the async caller has the values in hand. Avoids the round-trip of
|
||||
/// rebuilding a Tensor just to call to_vec1 again.
|
||||
#[allow(dead_code)]
|
||||
fn logits_health_slice(values: &[f32]) -> LogitsHealth {
|
||||
pub(crate) fn logits_health_slice(values: &[f32]) -> LogitsHealth {
|
||||
let mut nan = 0usize;
|
||||
let mut pos_inf = 0usize;
|
||||
let mut neg_inf = 0usize;
|
||||
@@ -1088,7 +1174,7 @@ fn logits_health_slice(values: &[f32]) -> LogitsHealth {
|
||||
/// the TP streaming task). Matching against the full chain lets the
|
||||
/// classification survive `.context("…")` and `format!("…: {e}")`
|
||||
/// wrappers in the call sites.
|
||||
fn is_device_fault(chain_text: &str) -> bool {
|
||||
pub(crate) fn is_device_fault(chain_text: &str) -> bool {
|
||||
let chain = chain_text.to_lowercase();
|
||||
// Non-device patterns: shape errors are pre-kernel and don't touch
|
||||
// GPU state; NaN-logits failures happen on the CPU side after the
|
||||
@@ -1462,7 +1548,7 @@ async fn acquire_pool_lock<'a>(
|
||||
/// Apply the repetition penalty (if any) to the prediction logits and
|
||||
/// then sample. Centralises the prefill / generation-loop call sites
|
||||
/// so they share identical sampling behaviour.
|
||||
fn sample_with_penalty(
|
||||
pub(crate) fn sample_with_penalty(
|
||||
logits: &Tensor,
|
||||
history: &[u32],
|
||||
logits_processor: &mut LogitsProcessor,
|
||||
@@ -1521,8 +1607,7 @@ fn chunked_prefill_local(
|
||||
/// chunk's last position. Tensors never escape the worker.
|
||||
/// `start_offset` skips a restored cached prefix, as in
|
||||
/// [`chunked_prefill_local`].
|
||||
#[cfg(feature = "cuda")]
|
||||
async fn chunked_prefill_via_worker(
|
||||
pub(crate) async fn chunked_prefill_via_worker(
|
||||
worker: &super::device_worker::DeviceWorkerHandle,
|
||||
handle: super::device_worker::ArchHandle,
|
||||
prompt_tokens: &[u32],
|
||||
@@ -1974,9 +2059,14 @@ impl CandleHarness {
|
||||
);
|
||||
|
||||
// bf16 is the canonical distribution dtype for Qwen3 /
|
||||
// Llama 3 / Qwen3 MoE. CUDA on Ada+ has hardware bf16;
|
||||
// Ampere has it too. CPU emulates.
|
||||
let dtype = DType::BF16;
|
||||
// Llama 3 / Qwen3 MoE; CUDA on Ampere+ has hardware bf16.
|
||||
// candle's CPU backend has no bf16 matmul, so the CPU
|
||||
// fallback upcasts to f32 at load.
|
||||
let dtype = if device_for_load.is_cuda() {
|
||||
DType::BF16
|
||||
} else {
|
||||
DType::F32
|
||||
};
|
||||
// SAFETY: VarBuilder::from_mmaped_safetensors mmaps the files;
|
||||
// mutation by another process while we hold the mapping is
|
||||
// UB. We trust the HF cache is immutable-by-design.
|
||||
@@ -2021,14 +2111,16 @@ impl CandleHarness {
|
||||
device: device_for_load,
|
||||
})))
|
||||
}
|
||||
"qwen3_5" => {
|
||||
"qwen3_5" | "qwen3_next" => {
|
||||
// Qwen3-Next needs a ShardedVarBuilder because its
|
||||
// load functions use the sharded backend (so they
|
||||
// can be reused unchanged by the future TP variant).
|
||||
// With world_size=1 the backend falls through to
|
||||
// the unsharded path, so there is no per-load cost.
|
||||
let cfg: super::arch::qwen3_5::Config = serde_json::from_str(&cfg_text)
|
||||
.context("parse Qwen3-Next (qwen3_5) config.json")?;
|
||||
// `from_config_json` normalises the flat qwen3_next
|
||||
// layout (#92) into the nested qwen3_5 shape.
|
||||
let cfg = super::arch::qwen3_5::Config::from_config_json(&cfg_text)
|
||||
.context("parse Qwen3-Next (qwen3_5/qwen3_next) config.json")?;
|
||||
let sharded_vb = unsafe {
|
||||
candle_nn::var_builder::ShardedSafeTensors::var_builder(
|
||||
&safetensors_paths,
|
||||
@@ -2269,6 +2361,12 @@ impl CandleHarness {
|
||||
};
|
||||
|
||||
let prompt_len = prompt_tokens.len();
|
||||
// Whether the chat template left the think block open in the
|
||||
// generation prompt (#112) — decides the truncated-mid-think
|
||||
// case in `split_off_reasoning`. Computed before the
|
||||
// inference closure takes ownership of `prompt_tokens`.
|
||||
let prompt_opened_reasoning =
|
||||
prompt_opens_reasoning(&prompt_tokens, loaded.reasoning_tokens.as_ref());
|
||||
let temperature = request.temperature.unwrap_or(0.7);
|
||||
let top_p = request.top_p;
|
||||
let max_new = request.max_tokens.unwrap_or(8192) as usize;
|
||||
@@ -2340,7 +2438,7 @@ impl CandleHarness {
|
||||
worker,
|
||||
handle,
|
||||
&prompt_tokens,
|
||||
loaded.prefix_cache.as_ref(),
|
||||
loaded.prefix_cache.as_deref(),
|
||||
loaded.tokenizer.token_to_id("<|im_start|>"),
|
||||
max_new,
|
||||
temperature,
|
||||
@@ -2392,7 +2490,7 @@ impl CandleHarness {
|
||||
&mut guard,
|
||||
&device,
|
||||
&prompt_tokens,
|
||||
loaded_for_cache.prefix_cache.as_ref(),
|
||||
loaded_for_cache.prefix_cache.as_deref(),
|
||||
im_start_id,
|
||||
max_new,
|
||||
temperature,
|
||||
@@ -2454,21 +2552,40 @@ impl CandleHarness {
|
||||
)));
|
||||
};
|
||||
|
||||
// Strip the leading `<think>` span so the chain-of-thought
|
||||
// doesn't leak into `content` (the streaming path drops it
|
||||
// as ReasoningDelta; this is the non-streaming equivalent).
|
||||
let (content_ids, reasoning_tokens) =
|
||||
split_off_reasoning(
|
||||
&generated_ids,
|
||||
loaded.reasoning_tokens.as_ref(),
|
||||
prompt_opened_reasoning,
|
||||
);
|
||||
let completion_text = loaded
|
||||
.tokenizer
|
||||
.decode(&generated_ids, true)
|
||||
.decode(content_ids, true)
|
||||
.map_err(|e| InferenceError::Other(anyhow::anyhow!("detokenize: {e}")))?;
|
||||
// The first answer token after `</think>` is usually a
|
||||
// newline pair; trim it so `content` starts at the answer.
|
||||
let completion_text = if reasoning_tokens > 0 {
|
||||
completion_text.trim_start().to_string()
|
||||
} else {
|
||||
completion_text
|
||||
};
|
||||
|
||||
let usage = Usage {
|
||||
prompt_tokens: prompt_len as u64,
|
||||
completion_tokens: generated_ids.len() as u64,
|
||||
total_tokens: (prompt_len + generated_ids.len()) as u64,
|
||||
// Reasoning accounting is streaming-only: the
|
||||
// non-streaming path doesn't track `in_reasoning`
|
||||
// (would require post-hoc <think> span parsing).
|
||||
// Deferred — see #64.
|
||||
completion_tokens_details: None,
|
||||
// `reasoning_tokens` is an additive sub-count of
|
||||
// `completion_tokens` (which still counts every
|
||||
// generated token, reasoning included).
|
||||
completion_tokens_details: (reasoning_tokens > 0)
|
||||
.then_some(CompletionTokensDetails { reasoning_tokens }),
|
||||
prompt_tokens_details: None,
|
||||
// Non-streaming path: prefill/decode split is only
|
||||
// surfaced on the streaming Finish event today (#85).
|
||||
helexa_timing: None,
|
||||
};
|
||||
|
||||
tracing::info!(
|
||||
@@ -2783,7 +2900,29 @@ impl CandleHarness {
|
||||
.map_err(InferenceError::from)?;
|
||||
|
||||
let tool_schemas = build_tool_schemas(&request);
|
||||
if let (Some(worker), Some(handle)) = (loaded.worker.clone(), loaded.arch_handle) {
|
||||
// Batched decode engine (#98): text streams multiplex through
|
||||
// the per-model engine instead of serializing on
|
||||
// inference_lock. Vision requests keep the direct path (they
|
||||
// can't batch — M-RoPE positions) and serialize against the
|
||||
// engine via the lock it holds while active.
|
||||
if vision_route.is_none() && loaded.engine.is_some() {
|
||||
let engine = loaded.engine.clone().expect("checked is_some");
|
||||
engine
|
||||
.submit(super::engine::EngineRequest {
|
||||
prompt_tokens,
|
||||
max_new,
|
||||
temperature,
|
||||
top_p,
|
||||
seed,
|
||||
eos_id,
|
||||
tool_schemas,
|
||||
tx,
|
||||
admit,
|
||||
span: span_for_task,
|
||||
})
|
||||
.await
|
||||
.map_err(InferenceError::Other)?;
|
||||
} else if let (Some(worker), Some(handle)) = (loaded.worker.clone(), loaded.arch_handle) {
|
||||
#[cfg(feature = "cuda")]
|
||||
{
|
||||
let prompt_tokens = prompt_tokens.clone();
|
||||
@@ -2800,7 +2939,7 @@ impl CandleHarness {
|
||||
tokenizer,
|
||||
prompt_tokens,
|
||||
vision_route,
|
||||
loaded_for_task.prefix_cache.as_ref(),
|
||||
loaded_for_task.prefix_cache.as_deref(),
|
||||
&loaded_for_task.prefill_rate,
|
||||
max_new,
|
||||
temperature,
|
||||
@@ -2865,7 +3004,7 @@ impl CandleHarness {
|
||||
&device,
|
||||
&tokenizer,
|
||||
&prompt_tokens,
|
||||
loaded_for_task.prefix_cache.as_ref(),
|
||||
loaded_for_task.prefix_cache.as_deref(),
|
||||
max_new,
|
||||
temperature,
|
||||
top_p,
|
||||
@@ -3309,6 +3448,41 @@ impl Harness for CandleHarness {
|
||||
);
|
||||
}
|
||||
|
||||
let poisoned = Arc::new(AtomicBool::new(false));
|
||||
let inference_lock = Arc::new(tokio::sync::Mutex::new(()));
|
||||
let prefix_cache = self.new_prefix_cache(snapshot_capable).map(Arc::new);
|
||||
let prefill_rate = Arc::new(super::context_limit::PrefillRateEma::new());
|
||||
// Batched decode engine (#98): spawned when the operator raised
|
||||
// max_in_flight above 1 on a snapshot-capable worker-path model.
|
||||
let engine = match (&worker, arch_handle) {
|
||||
(Some(w), Some(h))
|
||||
if snapshot_capable
|
||||
&& self.admission_cfg.max_in_flight > 1
|
||||
&& super::engine::batching_enabled() =>
|
||||
{
|
||||
tracing::info!(
|
||||
model = %spec.model_id,
|
||||
max_slots = self.admission_cfg.max_in_flight,
|
||||
"batched decode engine enabled (#98)"
|
||||
);
|
||||
Some(super::engine::EngineHandle::spawn(
|
||||
super::engine::EngineConfig {
|
||||
model_id: spec.model_id.clone(),
|
||||
worker: Arc::clone(w),
|
||||
handle: h,
|
||||
tokenizer: tokenizer.clone(),
|
||||
prefix_cache: prefix_cache.clone(),
|
||||
prefill_rate: Arc::clone(&prefill_rate),
|
||||
reasoning_tokens: reasoning_tokens.clone(),
|
||||
tool_call_tokens: tool_call_tokens.clone(),
|
||||
poisoned: Arc::clone(&poisoned),
|
||||
inference_lock: Arc::clone(&inference_lock),
|
||||
max_slots: self.admission_cfg.max_in_flight,
|
||||
},
|
||||
))
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
let loaded = Arc::new(LoadedModel {
|
||||
model_id: spec.model_id.clone(),
|
||||
arch: arch_local,
|
||||
@@ -3316,10 +3490,10 @@ impl Harness for CandleHarness {
|
||||
device,
|
||||
quant: spec.quant.clone(),
|
||||
devices,
|
||||
poisoned: AtomicBool::new(false),
|
||||
poisoned,
|
||||
worker,
|
||||
arch_handle,
|
||||
inference_lock: tokio::sync::Mutex::new(()),
|
||||
inference_lock,
|
||||
admission: super::admission::AdmissionController::new(&self.admission_cfg),
|
||||
reasoning_tokens,
|
||||
tool_call_tokens,
|
||||
@@ -3328,11 +3502,12 @@ impl Harness for CandleHarness {
|
||||
image_token_id: vision_meta.image_token_id,
|
||||
image_grid_factor: vision_meta.image_grid_factor,
|
||||
spec: spec.clone(),
|
||||
prefix_cache: self.new_prefix_cache(snapshot_capable),
|
||||
prefix_cache,
|
||||
context_profile,
|
||||
prefill_rate: super::context_limit::PrefillRateEma::new(),
|
||||
prefill_rate,
|
||||
derived_input_cap: AtomicUsize::new(0),
|
||||
last_free_mb: AtomicU64::new(0),
|
||||
engine,
|
||||
});
|
||||
if loaded.prefix_cache.is_some() {
|
||||
tracing::info!(
|
||||
@@ -3957,6 +4132,11 @@ impl CandleHarness {
|
||||
// call — promotes the terminal finish_reason to ToolCalls
|
||||
// so Anthropic clients see stop_reason: tool_use.
|
||||
let mut emitted_tool_call = false;
|
||||
// Prefill/decode split timers (#85). Declared outside 'work
|
||||
// so the terminal Finish — built after the block exits — can
|
||||
// read them; populated at the prefill→decode boundary inside.
|
||||
let mut prefill_ms_measured: u32 = 0;
|
||||
let mut decode_start: Option<std::time::Instant> = None;
|
||||
|
||||
'work: {
|
||||
// Prefix-cache decision (#11): vision requests
|
||||
@@ -4084,14 +4264,16 @@ impl CandleHarness {
|
||||
break 'work;
|
||||
}
|
||||
};
|
||||
let prefill_elapsed = prefill_start.elapsed();
|
||||
prefill_ms_measured = prefill_elapsed.as_millis() as u32;
|
||||
tp_for_task
|
||||
.prefill_rate
|
||||
.record(prompt_len, prefill_start.elapsed());
|
||||
.record(prompt_len, prefill_elapsed);
|
||||
let (post_prefill_vram_free_mb, _) = tp_for_task.query_vram().await;
|
||||
tracing::info!(
|
||||
model = %model_id,
|
||||
prompt_len,
|
||||
prefill_ms = prefill_start.elapsed().as_millis(),
|
||||
prefill_ms = prefill_elapsed.as_millis(),
|
||||
vram_free_mb = post_prefill_vram_free_mb,
|
||||
"TP chat_completion (stream): prefill complete"
|
||||
);
|
||||
@@ -4116,6 +4298,8 @@ impl CandleHarness {
|
||||
break 'work;
|
||||
}
|
||||
};
|
||||
// Decode-phase timer for the Finish prefill/decode split (#85).
|
||||
decode_start = Some(std::time::Instant::now());
|
||||
|
||||
if Some(next_token) == eos_id {
|
||||
finish_reason = FinishReason::Stop;
|
||||
@@ -4393,6 +4577,13 @@ impl CandleHarness {
|
||||
prompt_tokens: prompt_len as u32,
|
||||
completion_tokens: all_tokens.len() as u32,
|
||||
reasoning_tokens: reasoning_token_count,
|
||||
timing: Some(FinishTiming {
|
||||
prefill_ms: prefill_ms_measured,
|
||||
decode_ms: decode_start
|
||||
.map(|d| d.elapsed().as_millis() as u32)
|
||||
.unwrap_or(0),
|
||||
prefill_tokens: prompt_len as u32,
|
||||
}),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
@@ -4722,19 +4913,34 @@ async fn chat_completion_tp_inner(
|
||||
}
|
||||
drop(pool);
|
||||
|
||||
// Strip the leading `<think>` span (see `split_off_reasoning` and the
|
||||
// single-GPU path) so the chain-of-thought doesn't leak into `content`.
|
||||
let (content_ids, reasoning_tokens) = split_off_reasoning(
|
||||
&generated,
|
||||
tp.reasoning_tokens.as_ref(),
|
||||
prompt_opens_reasoning(&prompt_tokens, tp.reasoning_tokens.as_ref()),
|
||||
);
|
||||
let completion_text = tp
|
||||
.tokenizer
|
||||
.decode(&generated, true)
|
||||
.decode(content_ids, true)
|
||||
.map_err(|e| InferenceError::Other(anyhow::anyhow!("detokenize: {e}")))?;
|
||||
let completion_text = if reasoning_tokens > 0 {
|
||||
completion_text.trim_start().to_string()
|
||||
} else {
|
||||
completion_text
|
||||
};
|
||||
|
||||
let usage = Usage {
|
||||
prompt_tokens: prompt_len as u64,
|
||||
completion_tokens: generated.len() as u64,
|
||||
total_tokens: (prompt_len + generated.len()) as u64,
|
||||
// Reasoning accounting is streaming-only (non-streaming TP path
|
||||
// doesn't track `in_reasoning`). Deferred — see #64.
|
||||
completion_tokens_details: None,
|
||||
// `reasoning_tokens` is an additive sub-count of `completion_tokens`.
|
||||
completion_tokens_details: (reasoning_tokens > 0)
|
||||
.then_some(CompletionTokensDetails { reasoning_tokens }),
|
||||
prompt_tokens_details: None,
|
||||
// Non-streaming path: prefill/decode split is only surfaced on
|
||||
// the streaming Finish event today (#85).
|
||||
helexa_timing: None,
|
||||
};
|
||||
|
||||
tracing::info!(
|
||||
@@ -4776,8 +4982,11 @@ async fn chat_completion_tp_inner(
|
||||
/// stays out of this function — the wire projector in
|
||||
/// [`crate::wire::openai_chat`] stamps it onto every chunk
|
||||
/// downstream.
|
||||
#[cfg(feature = "cuda")]
|
||||
async fn emit_delta(delta: &str, tx: &mpsc::Sender<InferenceEvent>, in_reasoning: bool) -> bool {
|
||||
pub(crate) async fn emit_delta(
|
||||
delta: &str,
|
||||
tx: &mpsc::Sender<InferenceEvent>,
|
||||
in_reasoning: bool,
|
||||
) -> bool {
|
||||
if delta.is_empty() {
|
||||
return true;
|
||||
}
|
||||
@@ -4813,7 +5022,7 @@ fn emit_delta_blocking(delta: &str, tx: &mpsc::Sender<InferenceEvent>, in_reason
|
||||
///
|
||||
/// `pair = None` short-circuits to `false` (no reasoning markers
|
||||
/// configured for this model → pass-through).
|
||||
fn handle_reasoning_marker(
|
||||
pub(crate) fn handle_reasoning_marker(
|
||||
next_token: u32,
|
||||
pair: Option<&ReasoningTokenPair>,
|
||||
in_reasoning: &mut bool,
|
||||
@@ -4838,7 +5047,10 @@ fn handle_reasoning_marker(
|
||||
/// visible text. Replaying the prompt's reasoning markers and starting
|
||||
/// the loop in whatever state the prompt ends in fixes that without
|
||||
/// disabling thinking. `None` pair (non-reasoning model) → false.
|
||||
fn prompt_opens_reasoning(prompt_tokens: &[u32], pair: Option<&ReasoningTokenPair>) -> bool {
|
||||
pub(crate) fn prompt_opens_reasoning(
|
||||
prompt_tokens: &[u32],
|
||||
pair: Option<&ReasoningTokenPair>,
|
||||
) -> bool {
|
||||
let Some(pair) = pair else { return false };
|
||||
let mut open = false;
|
||||
for &t in prompt_tokens {
|
||||
@@ -4853,7 +5065,7 @@ fn prompt_opens_reasoning(prompt_tokens: &[u32], pair: Option<&ReasoningTokenPai
|
||||
|
||||
/// Outcome of checking a sampled token against the model's
|
||||
/// tool-call markers.
|
||||
enum ToolCallMarker {
|
||||
pub(crate) enum ToolCallMarker {
|
||||
/// Not a tool-call marker — caller proceeds with the normal
|
||||
/// detokenize-and-emit path.
|
||||
None,
|
||||
@@ -4870,7 +5082,7 @@ enum ToolCallMarker {
|
||||
Exit { buffer: String },
|
||||
}
|
||||
|
||||
fn handle_tool_call_marker(
|
||||
pub(crate) fn handle_tool_call_marker(
|
||||
next_token: u32,
|
||||
pair: Option<&ToolCallTokenPair>,
|
||||
in_tool_call: &mut bool,
|
||||
@@ -4908,7 +5120,8 @@ fn handle_tool_call_marker(
|
||||
/// the Qwen-XML tool-call parser can coerce each `<parameter>` string
|
||||
/// to its declared JSON type. An empty map (no tools, or untyped
|
||||
/// params) makes the parser fall back to value-sniffing.
|
||||
type ToolSchemas = std::collections::HashMap<String, std::collections::HashMap<String, String>>;
|
||||
pub(crate) type ToolSchemas =
|
||||
std::collections::HashMap<String, std::collections::HashMap<String, String>>;
|
||||
|
||||
/// Extract [`ToolSchemas`] from a request's `tools` (OpenAI shape:
|
||||
/// `{type:"function", function:{name, parameters:{properties:{p:{type}}}}}`).
|
||||
@@ -4952,7 +5165,7 @@ fn build_tool_schemas(request: &ChatCompletionRequest) -> ToolSchemas {
|
||||
///
|
||||
/// Returns `None` only when neither form yields a usable name, so the
|
||||
/// caller can re-emit the raw block as text instead of swallowing it.
|
||||
fn parse_tool_call_body(
|
||||
pub(crate) fn parse_tool_call_body(
|
||||
body: &str,
|
||||
index: usize,
|
||||
schemas: &ToolSchemas,
|
||||
@@ -5559,7 +5772,10 @@ async fn run_inference_with_images_via_worker(
|
||||
///
|
||||
/// Returns `None` (run the request without storing a snapshot) when
|
||||
/// the marker id is unknown or the prompt has no usable boundary.
|
||||
fn stable_snapshot_cut(prompt_tokens: &[u32], im_start_id: Option<u32>) -> Option<usize> {
|
||||
pub(crate) fn stable_snapshot_cut(
|
||||
prompt_tokens: &[u32],
|
||||
im_start_id: Option<u32>,
|
||||
) -> Option<usize> {
|
||||
let id = im_start_id?;
|
||||
let cut = prompt_tokens.iter().rposition(|&t| t == id)? + 1;
|
||||
(cut < prompt_tokens.len()).then_some(cut)
|
||||
@@ -5583,8 +5799,7 @@ fn lock_prefix_cache(
|
||||
/// prompt tokens already in the cache after this call — prefill
|
||||
/// resumes at that offset. A failed restore drops the entry and falls
|
||||
/// back to clear + full prefill.
|
||||
#[cfg(feature = "cuda")]
|
||||
async fn restore_or_clear_via_worker(
|
||||
pub(crate) async fn restore_or_clear_via_worker(
|
||||
worker: &super::device_worker::DeviceWorkerHandle,
|
||||
handle: super::device_worker::ArchHandle,
|
||||
prefix_cache: Option<&ModelPrefixCache>,
|
||||
@@ -5633,8 +5848,7 @@ async fn restore_or_clear_via_worker(
|
||||
/// `prompt_tokens`) and register it. Eviction decided by the
|
||||
/// registry; evicted worker snapshots are dropped here. Best-effort —
|
||||
/// a failed snapshot only costs the next request its prefill saving.
|
||||
#[cfg(feature = "cuda")]
|
||||
async fn store_prefix_snapshot_via_worker(
|
||||
pub(crate) async fn store_prefix_snapshot_via_worker(
|
||||
worker: &super::device_worker::DeviceWorkerHandle,
|
||||
handle: super::device_worker::ArchHandle,
|
||||
prefix_cache: Option<&ModelPrefixCache>,
|
||||
@@ -6064,7 +6278,8 @@ async fn stream_inference_via_worker(
|
||||
}
|
||||
}
|
||||
};
|
||||
prefill_rate.record(prefill_prompt_len, prefill_start.elapsed());
|
||||
let prefill_elapsed = prefill_start.elapsed();
|
||||
prefill_rate.record(prefill_prompt_len, prefill_elapsed);
|
||||
let logits = Tensor::new(logits_vec.as_slice(), &Device::Cpu)?;
|
||||
let mut next_token = match sample_with_penalty(&logits, &all_tokens, &mut logits_processor) {
|
||||
Ok(t) => t,
|
||||
@@ -6077,6 +6292,8 @@ async fn stream_inference_via_worker(
|
||||
return Err(e);
|
||||
}
|
||||
};
|
||||
// Decode-phase timer for the Finish prefill/decode split (#85).
|
||||
let decode_start = std::time::Instant::now();
|
||||
|
||||
// Per-token routing. `tokenizers::DecodeStream` carries five
|
||||
// generic parameters (`M, N, PT, PP, D`) which makes naming
|
||||
@@ -6221,6 +6438,11 @@ async fn stream_inference_via_worker(
|
||||
prompt_tokens: prompt_tokens.len() as u32,
|
||||
completion_tokens: all_tokens.len() as u32,
|
||||
reasoning_tokens: reasoning_token_count,
|
||||
timing: Some(FinishTiming {
|
||||
prefill_ms: prefill_elapsed.as_millis() as u32,
|
||||
decode_ms: decode_start.elapsed().as_millis() as u32,
|
||||
prefill_tokens: prefill_prompt_len as u32,
|
||||
}),
|
||||
})
|
||||
.await;
|
||||
|
||||
@@ -6355,6 +6577,10 @@ fn run_inference_streaming(
|
||||
// See `inference_tp_stream`: promotes finish_reason to ToolCalls.
|
||||
let mut emitted_tool_call = false;
|
||||
|
||||
// Time prefill and decode separately so the Finish event can carry
|
||||
// a server-measured prefill/decode split (#85) instead of leaving
|
||||
// the client to infer both from SSE chunk arrival.
|
||||
let prefill_start = std::time::Instant::now();
|
||||
let reused = restore_or_clear_local(arch, prefix_cache, prompt_tokens)?;
|
||||
// Two-stage prefill around the retokenization-stable snapshot
|
||||
// boundary — see `run_inference_via_worker`.
|
||||
@@ -6373,6 +6599,8 @@ fn run_inference_streaming(
|
||||
None => chunked_prefill_local(arch, device, prompt_tokens, reused)?,
|
||||
};
|
||||
let mut next_token = sample_with_penalty(&logits, &all_tokens, &mut logits_processor)?;
|
||||
let prefill_elapsed = prefill_start.elapsed();
|
||||
let decode_start = std::time::Instant::now();
|
||||
|
||||
// Per-token routing block, used at both the prefill-sample
|
||||
// tail and the decode loop. Macros are ugly but Rust's
|
||||
@@ -6481,6 +6709,11 @@ fn run_inference_streaming(
|
||||
prompt_tokens: prompt_tokens.len() as u32,
|
||||
completion_tokens: all_tokens.len() as u32,
|
||||
reasoning_tokens: reasoning_token_count,
|
||||
timing: Some(FinishTiming {
|
||||
prefill_ms: prefill_elapsed.as_millis() as u32,
|
||||
decode_ms: decode_start.elapsed().as_millis() as u32,
|
||||
prefill_tokens: prompt_tokens.len() as u32,
|
||||
}),
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
@@ -6505,6 +6738,71 @@ mod tests {
|
||||
|
||||
const IM_START: u32 = 999;
|
||||
|
||||
fn think_pair() -> ReasoningTokenPair {
|
||||
ReasoningTokenPair {
|
||||
open_id: 100,
|
||||
close_id: 200,
|
||||
open_text: "<think>".into(),
|
||||
close_text: "</think>".into(),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_off_reasoning_strips_up_to_close_marker() {
|
||||
// [reasoning_a, reasoning_b, </think>, answer_x, answer_y]
|
||||
let ids = [10, 11, 200, 42, 43];
|
||||
let (content, reasoning) = split_off_reasoning(&ids, Some(&think_pair()), false);
|
||||
assert_eq!(content, &[42, 43]);
|
||||
assert_eq!(reasoning, 3); // two reasoning tokens + the close marker
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_off_reasoning_no_close_marker_returns_all() {
|
||||
// Thinking disabled / span not opened by the prompt: return as-is.
|
||||
let ids = [42, 43, 44];
|
||||
let (content, reasoning) = split_off_reasoning(&ids, Some(&think_pair()), false);
|
||||
assert_eq!(content, &ids);
|
||||
assert_eq!(reasoning, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_off_reasoning_no_marker_pair_is_noop() {
|
||||
let ids = [1, 2, 3];
|
||||
let (content, reasoning) = split_off_reasoning(&ids, None, true);
|
||||
assert_eq!(content, &ids);
|
||||
assert_eq!(reasoning, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_off_reasoning_close_at_end_yields_empty_content() {
|
||||
// All reasoning, answer truncated to nothing after the marker.
|
||||
let ids = [10, 11, 200];
|
||||
let (content, reasoning) = split_off_reasoning(&ids, Some(&think_pair()), false);
|
||||
assert!(content.is_empty());
|
||||
assert_eq!(reasoning, 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_off_reasoning_prompt_opened_truncation_is_all_reasoning() {
|
||||
// Template force-opened the think block (#112:
|
||||
// Qwen3-Next-80B-A3B-Thinking) and generation hit max_tokens
|
||||
// before emitting </think> — everything is chain-of-thought.
|
||||
let ids = [10, 11, 12];
|
||||
let (content, reasoning) = split_off_reasoning(&ids, Some(&think_pair()), true);
|
||||
assert!(content.is_empty());
|
||||
assert_eq!(reasoning, 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_off_reasoning_splits_on_last_close_marker() {
|
||||
// Defensive: if the model emits its own <think></think> pair plus
|
||||
// the prompt-injected one, split on the LAST close marker.
|
||||
let ids = [200, 10, 200, 42];
|
||||
let (content, reasoning) = split_off_reasoning(&ids, Some(&think_pair()), true);
|
||||
assert_eq!(content, &[42]);
|
||||
assert_eq!(reasoning, 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stable_snapshot_cut_lands_after_last_im_start() {
|
||||
// ChatML shape: [im_start, "system", ..., im_start, "user",
|
||||
@@ -6580,6 +6878,20 @@ mod tests {
|
||||
.expect("qwen3_5 should be in the supported set as of Stage 8c scaffold");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_dense_config_accepts_qwen3_next() {
|
||||
// The MoE sibling family (Qwen3-Next-80B-A3B, #92) routes into
|
||||
// the same qwen3_5 arch module via Config::from_config_json.
|
||||
let cfg = r#"{
|
||||
"model_type": "qwen3_next",
|
||||
"architectures": ["Qwen3NextForCausalLM"],
|
||||
"hidden_size": 2048,
|
||||
"num_experts": 512
|
||||
}"#;
|
||||
check_dense_config_supported(cfg, "Qwen/Qwen3-Next-80B-A3B-Instruct")
|
||||
.expect("qwen3_next should be in the supported set (#92)");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_dense_config_rejects_missing_model_type() {
|
||||
let cfg = r#"{ "vocab_size": 1234 }"#;
|
||||
|
||||
@@ -248,6 +248,46 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
|
||||
let result = forward_logits(&mut state, handle, &tokens, offset);
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
Job::AssembleKvBatch {
|
||||
handle,
|
||||
seqs,
|
||||
reply,
|
||||
} => {
|
||||
let result = assemble_kv_batch(&mut state, handle, &seqs);
|
||||
// The replaced live cache state just freed its tensors.
|
||||
if result.is_ok() {
|
||||
trim_device_pool(&state);
|
||||
}
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
Job::ForwardLogitsBatch {
|
||||
handle,
|
||||
tokens,
|
||||
prefix_lens,
|
||||
padded_len,
|
||||
step,
|
||||
reply,
|
||||
} => {
|
||||
let result = forward_logits_batch(
|
||||
&mut state,
|
||||
handle,
|
||||
&tokens,
|
||||
&prefix_lens,
|
||||
padded_len,
|
||||
step,
|
||||
);
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
Job::ExtractKvRows {
|
||||
handle,
|
||||
rows,
|
||||
padded_len,
|
||||
steps,
|
||||
reply,
|
||||
} => {
|
||||
let result = extract_kv_rows(&mut state, handle, &rows, padded_len, steps);
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
Job::EncodeImage {
|
||||
handle,
|
||||
pixels,
|
||||
@@ -760,9 +800,14 @@ fn load_dense_inner(
|
||||
);
|
||||
|
||||
// bf16 is the canonical distribution dtype for Qwen3 / Llama 3 /
|
||||
// Qwen3 MoE. CUDA on Ada+ has hardware bf16; Ampere has it too.
|
||||
// CPU emulates.
|
||||
let dtype = DType::BF16;
|
||||
// Qwen3 MoE; CUDA on Ampere+ has hardware bf16. candle's CPU
|
||||
// backend has no bf16 matmul, so the CPU fallback (and the CPU
|
||||
// test worker) upcasts to f32 at load.
|
||||
let dtype = if device.is_cuda() {
|
||||
DType::BF16
|
||||
} else {
|
||||
DType::F32
|
||||
};
|
||||
// SAFETY: VarBuilder::from_mmaped_safetensors mmaps the files;
|
||||
// mutation by another process while we hold the mapping is UB.
|
||||
// We trust the HF cache is immutable-by-design.
|
||||
@@ -804,9 +849,11 @@ fn load_dense_inner(
|
||||
),
|
||||
)))
|
||||
}
|
||||
"qwen3_5" => {
|
||||
let cfg: crate::harness::arch::qwen3_5::Config = serde_json::from_str(&cfg_text)
|
||||
.context("parse Qwen3-Next (qwen3_5) config.json")?;
|
||||
"qwen3_5" | "qwen3_next" => {
|
||||
// `from_config_json` normalises the flat qwen3_next layout
|
||||
// (#92) into the nested qwen3_5 shape.
|
||||
let cfg = crate::harness::arch::qwen3_5::Config::from_config_json(&cfg_text)
|
||||
.context("parse Qwen3-Next (qwen3_5/qwen3_next) config.json")?;
|
||||
let sharded_vb = unsafe {
|
||||
candle_nn::var_builder::ShardedSafeTensors::var_builder(
|
||||
safetensors_paths,
|
||||
@@ -873,8 +920,8 @@ fn tp_load_shard_inner(
|
||||
&cfg, &vb, 0, world_size, comm,
|
||||
)?)
|
||||
}
|
||||
"qwen3_5" => {
|
||||
let cfg: crate::harness::tp::tp_qwen3_5::Config = serde_json::from_str(config_json)
|
||||
"qwen3_5" | "qwen3_next" => {
|
||||
let cfg = crate::harness::tp::tp_qwen3_5::Config::from_config_json(config_json)
|
||||
.context("parse Qwen3-Next Config JSON for leader load")?;
|
||||
let quant_dtype = crate::harness::tp::worker::parse_quant_string(quant)?;
|
||||
TpLeaderModel::Qwen3_5(crate::harness::tp::tp_qwen3_5::TpQwen3_5ForCausalLM::load(
|
||||
@@ -888,7 +935,8 @@ fn tp_load_shard_inner(
|
||||
)?)
|
||||
}
|
||||
other => anyhow::bail!(
|
||||
"TP dispatch: unsupported model_type '{other}' on leader (supported: qwen3, qwen3_5)"
|
||||
"TP dispatch: unsupported model_type '{other}' on leader \
|
||||
(supported: qwen3, qwen3_5, qwen3_next)"
|
||||
),
|
||||
};
|
||||
|
||||
@@ -1034,6 +1082,119 @@ fn forward_logits(
|
||||
Ok(values)
|
||||
}
|
||||
|
||||
/// Assemble stored per-sequence snapshots into a batched cache state
|
||||
/// and install it as the model's live state (#98). Split-borrows the
|
||||
/// snapshot map (immutable) and the model slab (mutable) — disjoint
|
||||
/// fields of the worker state.
|
||||
fn assemble_kv_batch(
|
||||
state: &mut DeviceWorkerState,
|
||||
handle: ArchHandle,
|
||||
seqs: &[(KvSnapshotId, usize)],
|
||||
) -> anyhow::Result<usize> {
|
||||
let DeviceWorkerState {
|
||||
models,
|
||||
kv_snapshots,
|
||||
..
|
||||
} = state;
|
||||
let mut pairs = Vec::with_capacity(seqs.len());
|
||||
for (id, len) in seqs {
|
||||
let snap = kv_snapshots.get(&(handle, id.0)).ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"AssembleKvBatch: no snapshot {} for handle {}",
|
||||
id.0,
|
||||
handle.0
|
||||
)
|
||||
})?;
|
||||
pairs.push((snap, *len));
|
||||
}
|
||||
let batch = crate::harness::arch::qwen3_5::snapshot::assemble_batch(&pairs)?;
|
||||
let arch = models
|
||||
.get_mut(&handle)
|
||||
.ok_or_else(|| anyhow::anyhow!("AssembleKvBatch: no model for handle {}", handle.0))?;
|
||||
arch.restore_kv_cache(&batch.snapshot)?;
|
||||
Ok(batch.padded_len)
|
||||
}
|
||||
|
||||
/// Extract live batched-state rows into stored per-sequence snapshots
|
||||
/// (#98) — see `Job::ExtractKvRows`. Captures the live state once
|
||||
/// (shallow attention KV, deep-copied GDN) and slices each requested
|
||||
/// row out gap-free.
|
||||
fn extract_kv_rows(
|
||||
state: &mut DeviceWorkerState,
|
||||
handle: ArchHandle,
|
||||
rows: &[(usize, usize)],
|
||||
padded_len: usize,
|
||||
steps: usize,
|
||||
) -> anyhow::Result<Vec<(KvSnapshotId, u64)>> {
|
||||
let live = state
|
||||
.models
|
||||
.get(&handle)
|
||||
.ok_or_else(|| anyhow::anyhow!("ExtractKvRows: no model for handle {}", handle.0))?
|
||||
.snapshot_kv_cache()?;
|
||||
let mut out = Vec::with_capacity(rows.len());
|
||||
for &(row, prefix_len) in rows {
|
||||
let snap = crate::harness::arch::qwen3_5::snapshot::extract_row(
|
||||
&live, row, prefix_len, padded_len, steps,
|
||||
)?;
|
||||
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);
|
||||
out.push((id, bytes));
|
||||
}
|
||||
tracing::debug!(
|
||||
handle = handle.0,
|
||||
rows = rows.len(),
|
||||
stored = state.kv_snapshots.len(),
|
||||
"device worker: batch rows extracted"
|
||||
);
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// One lockstep batched decode step (#98). Builds the `(B, 1)` input
|
||||
/// on the worker's device, derives per-row positions and the padding
|
||||
/// mask, and copies each row's logits back to CPU — same
|
||||
/// "tensors never escape the worker" contract as `forward_logits`.
|
||||
fn forward_logits_batch(
|
||||
state: &mut DeviceWorkerState,
|
||||
handle: ArchHandle,
|
||||
tokens: &[u32],
|
||||
prefix_lens: &[usize],
|
||||
padded_len: usize,
|
||||
step: usize,
|
||||
) -> anyhow::Result<Vec<Vec<f32>>> {
|
||||
use candle_core::{DType, Tensor};
|
||||
|
||||
let b = tokens.len();
|
||||
anyhow::ensure!(b > 0, "ForwardLogitsBatch: empty batch");
|
||||
anyhow::ensure!(
|
||||
prefix_lens.len() == b,
|
||||
"ForwardLogitsBatch: {} prefix_lens for batch of {b}",
|
||||
prefix_lens.len()
|
||||
);
|
||||
|
||||
let input = Tensor::from_vec(tokens.to_vec(), (b, 1), &state.device)?;
|
||||
let arch = state
|
||||
.models
|
||||
.get_mut(&handle)
|
||||
.ok_or_else(|| anyhow::anyhow!("ForwardLogitsBatch: no model for handle {}", handle.0))?;
|
||||
|
||||
let positions: Vec<usize> = prefix_lens.iter().map(|&len| len + step).collect();
|
||||
let total_len = padded_len + step + 1;
|
||||
let mask = arch.batch_decode_mask(prefix_lens, padded_len, total_len)?;
|
||||
let logits = arch.forward_batch_decode(&input, &positions, mask.as_ref())?;
|
||||
|
||||
// (B, 1, vocab) → per-row CPU Vec<f32>.
|
||||
let logits = logits.to_dtype(DType::F32)?;
|
||||
let (rows, _, vocab) = logits.dims3()?;
|
||||
anyhow::ensure!(
|
||||
rows == b,
|
||||
"ForwardLogitsBatch: model returned {rows} logits rows for batch of {b}"
|
||||
);
|
||||
let flat: Vec<f32> = logits.flatten_all()?.to_vec1()?;
|
||||
Ok(flat.chunks(vocab).map(<[f32]>::to_vec).collect())
|
||||
}
|
||||
|
||||
/// Run the LM forward with vision-tower image splicing. Stage B3.
|
||||
///
|
||||
/// Encodes each image through the vision tower (`VisionTower::forward`,
|
||||
@@ -1187,6 +1348,15 @@ fn drain_poisoned(job: Job, device_index: u32) {
|
||||
Job::ForwardLogits { reply, .. } => {
|
||||
let _ = reply.send(Err(err()));
|
||||
}
|
||||
Job::AssembleKvBatch { reply, .. } => {
|
||||
let _ = reply.send(Err(err()));
|
||||
}
|
||||
Job::ForwardLogitsBatch { reply, .. } => {
|
||||
let _ = reply.send(Err(err()));
|
||||
}
|
||||
Job::ExtractKvRows { reply, .. } => {
|
||||
let _ = reply.send(Err(err()));
|
||||
}
|
||||
Job::EncodeImage { reply, .. } => {
|
||||
let _ = reply.send(Err(err()));
|
||||
}
|
||||
|
||||
@@ -149,6 +149,50 @@ pub enum Job {
|
||||
offset: usize,
|
||||
reply: oneshot::Sender<Result<Vec<f32>>>,
|
||||
},
|
||||
/// Assemble stored per-sequence snapshots into one batched cache
|
||||
/// state and install it as the model's live state (#98). `seqs`
|
||||
/// pairs each snapshot id with its true token length; attention
|
||||
/// K/V is right-padded to the batch max and `cat`ed on dim 0 (see
|
||||
/// `arch::qwen3_5::snapshot::assemble_batch`). Replies with the
|
||||
/// padded uniform KV length. The source snapshots remain stored —
|
||||
/// the caller drops them via `DropKvSnapshot` when the sequences
|
||||
/// leave the batch.
|
||||
AssembleKvBatch {
|
||||
handle: ArchHandle,
|
||||
seqs: Vec<(KvSnapshotId, usize)>,
|
||||
reply: oneshot::Sender<Result<usize>>,
|
||||
},
|
||||
/// Extract rows of the model's **live** batched cache state back
|
||||
/// into contiguous single-sequence snapshots stored in the
|
||||
/// worker's slab (#98) — the first half of a rebatch (join or
|
||||
/// leave). `rows` pairs each batch-row index with its prefix
|
||||
/// length; `padded_len`/`steps` describe the live batch geometry.
|
||||
/// Replies one `(snapshot id, bytes)` per requested row, in
|
||||
/// order. Compose with `AssembleKvBatch` to form the new batch,
|
||||
/// then `DropKvSnapshot` the intermediates.
|
||||
ExtractKvRows {
|
||||
handle: ArchHandle,
|
||||
/// `(batch row index, prefix_len)` per surviving sequence.
|
||||
rows: Vec<(usize, usize)>,
|
||||
padded_len: usize,
|
||||
steps: usize,
|
||||
reply: oneshot::Sender<Result<Vec<(KvSnapshotId, u64)>>>,
|
||||
},
|
||||
/// One lockstep batched decode step (#98): `tokens[i]` is batch
|
||||
/// row i's next token, sitting at sequence position
|
||||
/// `prefix_lens[i] + step`. The handler derives per-row positions
|
||||
/// and the padding mask from `prefix_lens`/`padded_len` (the
|
||||
/// values `AssembleKvBatch` was built from) and replies one CPU
|
||||
/// `[vocab]` logits row per batch row, ready for per-slot
|
||||
/// sampling on the async side.
|
||||
ForwardLogitsBatch {
|
||||
handle: ArchHandle,
|
||||
tokens: Vec<u32>,
|
||||
prefix_lens: Vec<usize>,
|
||||
padded_len: usize,
|
||||
step: usize,
|
||||
reply: oneshot::Sender<Result<Vec<Vec<f32>>>>,
|
||||
},
|
||||
/// Run the LM forward with vision splicing in one round-trip.
|
||||
/// Stage B3 of the vision plan.
|
||||
///
|
||||
|
||||
@@ -420,6 +420,111 @@ impl DeviceWorkerHandle {
|
||||
}
|
||||
}
|
||||
|
||||
/// Assemble stored per-sequence snapshots into one batched cache
|
||||
/// state and install it as the model's live state (#98). Returns
|
||||
/// the padded uniform KV length the batch was assembled to. The
|
||||
/// source snapshots remain stored.
|
||||
pub async fn assemble_kv_batch(
|
||||
&self,
|
||||
handle: ArchHandle,
|
||||
seqs: Vec<(jobs::KvSnapshotId, usize)>,
|
||||
) -> Result<usize, 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::AssembleKvBatch {
|
||||
handle,
|
||||
seqs,
|
||||
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,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract live batched-state rows into stored per-sequence
|
||||
/// snapshots (#98) — the first half of a rebatch. Returns one
|
||||
/// `(snapshot id, bytes)` per requested row, in order.
|
||||
pub async fn extract_kv_rows(
|
||||
&self,
|
||||
handle: ArchHandle,
|
||||
rows: Vec<(usize, usize)>,
|
||||
padded_len: usize,
|
||||
steps: usize,
|
||||
) -> Result<Vec<(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::ExtractKvRows {
|
||||
handle,
|
||||
rows,
|
||||
padded_len,
|
||||
steps,
|
||||
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,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// One lockstep batched decode step (#98): row i's next token at
|
||||
/// position `prefix_lens[i] + step`. Returns one CPU `[vocab]`
|
||||
/// logits row per batch row, ready for per-slot sampling — same
|
||||
/// no-device-tensor contract as [`Self::forward_logits`].
|
||||
pub async fn forward_logits_batch(
|
||||
&self,
|
||||
handle: ArchHandle,
|
||||
tokens: Vec<u32>,
|
||||
prefix_lens: Vec<usize>,
|
||||
padded_len: usize,
|
||||
step: usize,
|
||||
) -> Result<Vec<Vec<f32>>, 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::ForwardLogitsBatch {
|
||||
handle,
|
||||
tokens,
|
||||
prefix_lens,
|
||||
padded_len,
|
||||
step,
|
||||
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,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Forward with image-aware splicing in one round-trip. Stage B3.
|
||||
///
|
||||
/// Encodes each image on the worker thread (device-resident), then
|
||||
@@ -925,6 +1030,189 @@ mod tests {
|
||||
handle.shutdown().expect("shutdown ok");
|
||||
}
|
||||
|
||||
/// #98 slice 2 end-to-end: load the tiny qwen3_next fixture through
|
||||
/// the worker, prefill three ragged sequences (snapshotting each),
|
||||
/// assemble them into one batched state via `AssembleKvBatch`, and
|
||||
/// check every `ForwardLogitsBatch` row against the sequential
|
||||
/// `ForwardLogits` reference at each decode step. Exercises the
|
||||
/// whole job plumbing on the real dispatch thread (CPU device in
|
||||
/// CI), not just the arch-level primitives.
|
||||
#[tokio::test]
|
||||
async fn batched_decode_jobs_match_sequential_forward() {
|
||||
let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
|
||||
.join("tests/fixtures/numerical/qwen3_next-tiny");
|
||||
if !fixture.join("model.safetensors").exists() {
|
||||
eprintln!("SKIP batched_decode_jobs_match_sequential_forward: fixture not generated");
|
||||
return;
|
||||
}
|
||||
|
||||
let worker = DeviceWorkerHandle::spawn(0).expect("spawn ok");
|
||||
let arch = worker
|
||||
.load_dense(
|
||||
fixture.join("config.json"),
|
||||
vec![fixture.join("model.safetensors")],
|
||||
"qwen3_next-tiny".into(),
|
||||
)
|
||||
.await
|
||||
.expect("load tiny fixture");
|
||||
|
||||
let prompts: [&[u32]; 3] = [&[1, 2, 3], &[4, 5], &[7, 3, 2, 5, 6]];
|
||||
let steps: [&[u32]; 3] = [&[11, 12, 13], &[9, 8, 7], &[21, 22, 23]];
|
||||
let n_steps = 3;
|
||||
|
||||
// Sequential reference + per-sequence snapshots. Snapshot is
|
||||
// taken at the prefill boundary; the reference decode that
|
||||
// follows mutates only the live state (GDN snapshots are deep
|
||||
// copies, attention KV never mutates in place).
|
||||
let mut snaps = Vec::new();
|
||||
let mut expected: Vec<Vec<Vec<f32>>> = Vec::new(); // [row][step]
|
||||
for (prompt, toks) in prompts.iter().zip(steps.iter()) {
|
||||
worker.clear_kv_cache(arch).await.expect("clear");
|
||||
worker
|
||||
.forward_logits(arch, prompt.to_vec(), 0)
|
||||
.await
|
||||
.expect("prefill");
|
||||
let (snap, bytes) = worker.snapshot_kv(arch).await.expect("snapshot");
|
||||
assert!(bytes > 0);
|
||||
snaps.push(snap);
|
||||
let mut per_step = Vec::new();
|
||||
for (t, tok) in toks.iter().enumerate() {
|
||||
per_step.push(
|
||||
worker
|
||||
.forward_logits(arch, vec![*tok], prompt.len() + t)
|
||||
.await
|
||||
.expect("sequential decode step"),
|
||||
);
|
||||
}
|
||||
expected.push(per_step);
|
||||
}
|
||||
|
||||
// Assemble and decode lockstep.
|
||||
let seqs: Vec<(jobs::KvSnapshotId, usize)> = snaps
|
||||
.iter()
|
||||
.zip(prompts.iter())
|
||||
.map(|(s, p)| (*s, p.len()))
|
||||
.collect();
|
||||
let prefix_lens: Vec<usize> = prompts.iter().map(|p| p.len()).collect();
|
||||
let padded_len = worker
|
||||
.assemble_kv_batch(arch, seqs)
|
||||
.await
|
||||
.expect("assemble batch");
|
||||
assert_eq!(padded_len, 5);
|
||||
|
||||
for t in 0..n_steps {
|
||||
let toks: Vec<u32> = steps.iter().map(|s| s[t]).collect();
|
||||
let rows = worker
|
||||
.forward_logits_batch(arch, toks, prefix_lens.clone(), padded_len, t)
|
||||
.await
|
||||
.expect("batched decode step");
|
||||
assert_eq!(rows.len(), 3);
|
||||
for (row, got) in rows.iter().enumerate() {
|
||||
let want = &expected[row][t];
|
||||
assert_eq!(got.len(), want.len(), "row {row} vocab width");
|
||||
let diff = want
|
||||
.iter()
|
||||
.zip(got)
|
||||
.map(|(a, b)| (a - b).abs())
|
||||
.fold(0f32, f32::max);
|
||||
assert!(diff < 1e-4, "row {row} step {t} diverged: {diff}");
|
||||
}
|
||||
}
|
||||
|
||||
// Rebatch mid-decode (the leave path): extract rows 0 and 2
|
||||
// from the live batch, re-assemble without row 1, and check the
|
||||
// shrunken batch still tracks the sequential reference. The
|
||||
// extra decode tokens per remaining row are appended after the
|
||||
// original per-row streams for the reference run.
|
||||
let survivors = [0usize, 2];
|
||||
let extra: [&[u32]; 2] = [&[15, 16], &[25, 26]];
|
||||
let mut expected_extra: Vec<Vec<Vec<f32>>> = Vec::new();
|
||||
for (i, &row) in survivors.iter().enumerate() {
|
||||
worker.clear_kv_cache(arch).await.expect("clear");
|
||||
worker
|
||||
.forward_logits(arch, prompts[row].to_vec(), 0)
|
||||
.await
|
||||
.expect("re-prefill");
|
||||
for (t, tok) in steps[row].iter().enumerate() {
|
||||
worker
|
||||
.forward_logits(arch, vec![*tok], prompts[row].len() + t)
|
||||
.await
|
||||
.expect("re-decode");
|
||||
}
|
||||
let mut per_step = Vec::new();
|
||||
for (j, tok) in extra[i].iter().enumerate() {
|
||||
per_step.push(
|
||||
worker
|
||||
.forward_logits(arch, vec![*tok], prompts[row].len() + n_steps + j)
|
||||
.await
|
||||
.expect("reference extra step"),
|
||||
);
|
||||
}
|
||||
expected_extra.push(per_step);
|
||||
}
|
||||
|
||||
// Rebuild the 3-row batch state (it was clobbered by the
|
||||
// reference runs above), replay the lockstep steps, then
|
||||
// extract + re-assemble the survivors.
|
||||
let seqs: Vec<(jobs::KvSnapshotId, usize)> = snaps
|
||||
.iter()
|
||||
.zip(prompts.iter())
|
||||
.map(|(s, p)| (*s, p.len()))
|
||||
.collect();
|
||||
worker
|
||||
.assemble_kv_batch(arch, seqs)
|
||||
.await
|
||||
.expect("re-assemble");
|
||||
for t in 0..n_steps {
|
||||
let toks: Vec<u32> = steps.iter().map(|s| s[t]).collect();
|
||||
worker
|
||||
.forward_logits_batch(arch, toks, prefix_lens.clone(), padded_len, t)
|
||||
.await
|
||||
.expect("replay batched step");
|
||||
}
|
||||
let extracted = worker
|
||||
.extract_kv_rows(
|
||||
arch,
|
||||
survivors.iter().map(|&r| (r, prompts[r].len())).collect(),
|
||||
padded_len,
|
||||
n_steps,
|
||||
)
|
||||
.await
|
||||
.expect("extract survivors");
|
||||
let new_lens: Vec<usize> = survivors
|
||||
.iter()
|
||||
.map(|&r| prompts[r].len() + n_steps)
|
||||
.collect();
|
||||
let new_seqs: Vec<(jobs::KvSnapshotId, usize)> = extracted
|
||||
.iter()
|
||||
.zip(new_lens.iter())
|
||||
.map(|((id, _), &len)| (*id, len))
|
||||
.collect();
|
||||
let new_padded = worker
|
||||
.assemble_kv_batch(arch, new_seqs)
|
||||
.await
|
||||
.expect("assemble survivors");
|
||||
assert_eq!(new_padded, 5 + n_steps);
|
||||
for j in 0..extra[0].len() {
|
||||
let toks: Vec<u32> = extra.iter().map(|s| s[j]).collect();
|
||||
let rows = worker
|
||||
.forward_logits_batch(arch, toks, new_lens.clone(), new_padded, j)
|
||||
.await
|
||||
.expect("post-rebatch step");
|
||||
for (i, got) in rows.iter().enumerate() {
|
||||
let want = &expected_extra[i][j];
|
||||
let diff = want
|
||||
.iter()
|
||||
.zip(got)
|
||||
.map(|(a, b)| (a - b).abs())
|
||||
.fold(0f32, f32::max);
|
||||
assert!(diff < 1e-4, "survivor {i} extra step {j} diverged: {diff}");
|
||||
}
|
||||
}
|
||||
|
||||
worker.shutdown().expect("shutdown ok");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn shutdown_drains_pending_jobs() {
|
||||
let handle = DeviceWorkerHandle::spawn(0).expect("spawn ok");
|
||||
|
||||
860
crates/neuron/src/harness/engine.rs
Normal file
860
crates/neuron/src/harness/engine.rs
Normal file
@@ -0,0 +1,860 @@
|
||||
//! Lockstep batched decode engine (#98) — single-GPU worker path.
|
||||
//!
|
||||
//! One engine task per loaded model replaces the per-request
|
||||
//! `inference_lock` serialization for **text** chat streams when the
|
||||
//! operator raises `[admission] max_in_flight` above 1 (and the arch
|
||||
//! supports cache snapshots — qwen3_5 only today). The engine owns the
|
||||
//! model's forward exclusively and multiplexes up to `max_slots`
|
||||
//! concurrent streams through one `(B, 1)` forward per decode step:
|
||||
//!
|
||||
//! - **Join** (new request): prefill runs alone at B=1 through the
|
||||
//! existing chunked-prefill + prefix-cache paths, then the fresh
|
||||
//! state is snapshotted and the batch re-assembled
|
||||
//! (`ExtractKvRows` survivors → `AssembleKvBatch` everyone). Decode
|
||||
//! for running slots stalls for the duration of the newcomer's
|
||||
//! prefill — the accepted v1 cost, bounded by chunked prefill.
|
||||
//! - **Step**: one `ForwardLogitsBatch` job; per-slot CPU sampling
|
||||
//! (each slot has its own `LogitsProcessor` + repeat-penalty
|
||||
//! history); sampled tokens go to per-slot **router tasks** that own
|
||||
//! the incremental detokenizer and the reasoning/tool-call state
|
||||
//! machine and emit `InferenceEvent`s on the request's channel.
|
||||
//! - **Leave** (EOS / length / consumer hangup): the slot's Finish is
|
||||
//! emitted and the batch compacts at the next rebatch point (which
|
||||
//! runs immediately after any step that finished a slot).
|
||||
//!
|
||||
//! Routers are separate tasks (not inline state) because
|
||||
//! `tokenizers::DecodeStream` borrows the tokenizer and carries five
|
||||
//! generic parameters — owning both inside one async block sidesteps
|
||||
//! the self-referential-struct problem the same way the
|
||||
//! `route_token!` macro does at its call sites, and decouples slow
|
||||
//! consumers from the lockstep loop.
|
||||
//!
|
||||
//! A worker error mid-step is fatal for the whole engine: every
|
||||
//! active slot's stream ends, the model is marked poisoned when the
|
||||
//! error classifies as a device fault, and the engine exits (later
|
||||
//! submits fail fast).
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
use anyhow::Result;
|
||||
use candle_core::{Device, Tensor};
|
||||
use candle_transformers::generation::LogitsProcessor;
|
||||
use tokenizers::Tokenizer;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use super::admission::AdmissionPermit;
|
||||
use super::candle::{
|
||||
ModelPrefixCache, ToolCallMarker, ToolSchemas, chunked_prefill_via_worker, emit_delta,
|
||||
handle_reasoning_marker, handle_tool_call_marker, is_device_fault, logits_health_slice,
|
||||
parse_tool_call_body, prompt_opens_reasoning, restore_or_clear_via_worker, sample_with_penalty,
|
||||
stable_snapshot_cut, store_prefix_snapshot_via_worker,
|
||||
};
|
||||
use super::context_limit::PrefillRateEma;
|
||||
use super::device_worker::{ArchHandle, DeviceWorkerHandle};
|
||||
use crate::wire::event::{
|
||||
FinishReason, FinishTiming, InferenceEvent, ReasoningTokenPair, ToolCallTokenPair,
|
||||
};
|
||||
|
||||
/// Runtime kill switch: `NEURON_BATCHING=0` (or `false`) keeps the
|
||||
/// per-request `inference_lock` path even when `max_in_flight > 1`.
|
||||
/// Read once.
|
||||
pub fn batching_enabled() -> bool {
|
||||
static ENABLED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
|
||||
*ENABLED.get_or_init(|| {
|
||||
let on = !std::env::var("NEURON_BATCHING").is_ok_and(|v| v == "0" || v == "false");
|
||||
tracing::info!(enabled = on, "batched decode engine (#98)");
|
||||
on
|
||||
})
|
||||
}
|
||||
|
||||
/// Everything the engine needs that is per-model (not per-request).
|
||||
/// Deliberately does NOT hold `Arc<LoadedModel>` — the engine task
|
||||
/// must not keep the model alive (the task exits when the model drops
|
||||
/// its `EngineHandle` and the channel closes).
|
||||
pub struct EngineConfig {
|
||||
pub model_id: String,
|
||||
pub worker: Arc<DeviceWorkerHandle>,
|
||||
pub handle: ArchHandle,
|
||||
pub tokenizer: Tokenizer,
|
||||
pub prefix_cache: Option<Arc<ModelPrefixCache>>,
|
||||
pub prefill_rate: Arc<PrefillRateEma>,
|
||||
pub reasoning_tokens: Option<ReasoningTokenPair>,
|
||||
pub tool_call_tokens: Option<ToolCallTokenPair>,
|
||||
/// Shared with `LoadedModel.poisoned` so a device fault inside the
|
||||
/// engine fast-rejects subsequent requests at the harness boundary.
|
||||
pub poisoned: Arc<AtomicBool>,
|
||||
/// Shared with `LoadedModel.inference_lock`. Held for the whole
|
||||
/// active phase (first join → last slot finished) so the
|
||||
/// non-engine forward paths (vision, non-streaming chat) can never
|
||||
/// clobber the live batched cache state mid-decode. Released while
|
||||
/// idle.
|
||||
pub inference_lock: Arc<tokio::sync::Mutex<()>>,
|
||||
pub max_slots: usize,
|
||||
}
|
||||
|
||||
/// One queued request. Admission has already been passed — the permit
|
||||
/// rides along and is released when the slot finishes.
|
||||
pub struct EngineRequest {
|
||||
pub prompt_tokens: Vec<u32>,
|
||||
pub max_new: usize,
|
||||
pub temperature: f64,
|
||||
pub top_p: Option<f64>,
|
||||
pub seed: u64,
|
||||
pub eos_id: Option<u32>,
|
||||
pub tool_schemas: ToolSchemas,
|
||||
pub tx: mpsc::Sender<InferenceEvent>,
|
||||
pub admit: AdmissionPermit,
|
||||
pub span: tracing::Span,
|
||||
}
|
||||
|
||||
/// Cheap handle held by `LoadedModel`. Submitting fails once the
|
||||
/// engine task has exited (fatal worker error) — callers surface that
|
||||
/// as an inference error; the model is typically poisoned by then.
|
||||
#[derive(Clone)]
|
||||
pub struct EngineHandle {
|
||||
tx: mpsc::Sender<EngineRequest>,
|
||||
}
|
||||
|
||||
impl EngineHandle {
|
||||
pub fn spawn(cfg: EngineConfig) -> Self {
|
||||
// Depth beyond max_slots only buffers between admission and the
|
||||
// engine momentarily; admission's queue is the real bound.
|
||||
let (tx, rx) = mpsc::channel::<EngineRequest>(cfg.max_slots.max(1) * 2);
|
||||
tokio::spawn(run_engine(cfg, rx));
|
||||
Self { tx }
|
||||
}
|
||||
|
||||
pub async fn submit(&self, req: EngineRequest) -> Result<()> {
|
||||
self.tx
|
||||
.send(req)
|
||||
.await
|
||||
.map_err(|_| anyhow::anyhow!("batch engine is not running (model poisoned?)"))
|
||||
}
|
||||
}
|
||||
|
||||
/// Messages from the engine loop to a slot's router task.
|
||||
enum RouterMsg {
|
||||
Token(u32),
|
||||
Finish {
|
||||
reason: FinishReason,
|
||||
prompt_tokens: u32,
|
||||
completion_tokens: u32,
|
||||
timing: FinishTiming,
|
||||
},
|
||||
}
|
||||
|
||||
struct Slot {
|
||||
/// Contiguous valid tokens this row held at the last rebatch
|
||||
/// (prompt + tokens decoded before that point).
|
||||
prefix_len: usize,
|
||||
prompt_len: usize,
|
||||
/// Completion tokens so far — the repeat-penalty history. EOS is
|
||||
/// never pushed (mirrors the B=1 paths).
|
||||
generated: Vec<u32>,
|
||||
next_token: u32,
|
||||
max_new: usize,
|
||||
eos_id: Option<u32>,
|
||||
lp: LogitsProcessor,
|
||||
router: mpsc::Sender<RouterMsg>,
|
||||
/// Set by the router when the consumer hangs up; the engine stops
|
||||
/// feeding the slot and compacts it out.
|
||||
hangup: Arc<AtomicBool>,
|
||||
finished: Option<FinishReason>,
|
||||
prefill_ms: u32,
|
||||
prefill_tokens: u32,
|
||||
decode_start: std::time::Instant,
|
||||
_admit: AdmissionPermit,
|
||||
}
|
||||
|
||||
impl Slot {
|
||||
fn finish(&mut self, reason: FinishReason) {
|
||||
self.finished = Some(reason);
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_engine(cfg: EngineConfig, mut rx: mpsc::Receiver<EngineRequest>) {
|
||||
let mut slots: Vec<Slot> = Vec::new();
|
||||
// Uniform padded KV length of the current batch and steps decoded
|
||||
// since the last rebatch — the geometry `ForwardLogitsBatch` and
|
||||
// `ExtractKvRows` key off.
|
||||
let mut padded_len = 0usize;
|
||||
let mut step = 0usize;
|
||||
// Held while any slot is active — see `EngineConfig.inference_lock`.
|
||||
let mut lock_guard: Option<tokio::sync::OwnedMutexGuard<()>> = None;
|
||||
|
||||
tracing::info!(
|
||||
model = %cfg.model_id,
|
||||
max_slots = cfg.max_slots,
|
||||
"batch engine started"
|
||||
);
|
||||
|
||||
'main: loop {
|
||||
// Gather joins: block when idle, drain opportunistically when
|
||||
// busy. Slots that finished or hung up leave at the rebatch
|
||||
// below.
|
||||
let mut joins: Vec<EngineRequest> = Vec::new();
|
||||
if slots.is_empty() {
|
||||
match rx.recv().await {
|
||||
Some(r) => joins.push(r),
|
||||
None => break 'main, // model unloaded
|
||||
}
|
||||
}
|
||||
while slots.len() + joins.len() < cfg.max_slots {
|
||||
match rx.try_recv() {
|
||||
Ok(r) => joins.push(r),
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
|
||||
// Take the model's inference lock before touching cache state;
|
||||
// release it whenever the batch drains so vision/non-streaming
|
||||
// requests get their turn.
|
||||
if !joins.is_empty() && lock_guard.is_none() {
|
||||
lock_guard = Some(Arc::clone(&cfg.inference_lock).lock_owned().await);
|
||||
}
|
||||
|
||||
let needs_compaction = slots
|
||||
.iter()
|
||||
.any(|s| s.finished.is_some() || s.hangup.load(Ordering::Acquire));
|
||||
if (!joins.is_empty() || needs_compaction)
|
||||
&& let Err(e) = rebatch(&cfg, &mut slots, joins, &mut padded_len, &mut step).await
|
||||
{
|
||||
fail_engine(&cfg, &mut slots, &mut rx, &e);
|
||||
break 'main;
|
||||
}
|
||||
if slots.is_empty() {
|
||||
lock_guard = None; // every join finished during prefill
|
||||
continue;
|
||||
}
|
||||
|
||||
// One lockstep decode step.
|
||||
let tokens: Vec<u32> = slots.iter().map(|s| s.next_token).collect();
|
||||
let prefix_lens: Vec<usize> = slots.iter().map(|s| s.prefix_len).collect();
|
||||
let rows = match cfg
|
||||
.worker
|
||||
.forward_logits_batch(cfg.handle, tokens, prefix_lens, padded_len, step)
|
||||
.await
|
||||
{
|
||||
Ok(rows) => rows,
|
||||
Err(e) => {
|
||||
let e = anyhow::anyhow!("batched decode step {step}: {e}");
|
||||
fail_engine(&cfg, &mut slots, &mut rx, &e);
|
||||
break 'main;
|
||||
}
|
||||
};
|
||||
step += 1;
|
||||
|
||||
let mut fatal: Option<anyhow::Error> = None;
|
||||
for (slot, logits_vec) in slots.iter_mut().zip(rows) {
|
||||
if slot.finished.is_some() || slot.hangup.load(Ordering::Acquire) {
|
||||
// Compacted out at the next rebatch; discard its row.
|
||||
continue;
|
||||
}
|
||||
let logits = match Tensor::new(logits_vec.as_slice(), &Device::Cpu) {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
fatal = Some(e.into());
|
||||
break;
|
||||
}
|
||||
};
|
||||
let nt = match sample_with_penalty(&logits, &slot.generated, &mut slot.lp) {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
let health = logits_health_slice(&logits_vec);
|
||||
tracing::warn!(
|
||||
?health,
|
||||
error = %e,
|
||||
"batch engine: sample failed; logits unhealthy"
|
||||
);
|
||||
// Unhealthy logits are a device-level problem —
|
||||
// fail the whole engine, mirroring the B=1 path's
|
||||
// poison classification.
|
||||
fatal = Some(e);
|
||||
break;
|
||||
}
|
||||
};
|
||||
if Some(nt) == slot.eos_id {
|
||||
finish_slot(slot, FinishReason::Stop).await;
|
||||
continue;
|
||||
}
|
||||
slot.generated.push(nt);
|
||||
slot.next_token = nt;
|
||||
if slot.router.send(RouterMsg::Token(nt)).await.is_err() {
|
||||
// Router exited (consumer hung up mid-drain).
|
||||
slot.hangup.store(true, Ordering::Release);
|
||||
slot.finish(FinishReason::Stop);
|
||||
continue;
|
||||
}
|
||||
if slot.generated.len() >= slot.max_new {
|
||||
finish_slot(slot, FinishReason::Length).await;
|
||||
}
|
||||
}
|
||||
if let Some(e) = fatal {
|
||||
fail_engine(&cfg, &mut slots, &mut rx, &e);
|
||||
break 'main;
|
||||
}
|
||||
}
|
||||
|
||||
tracing::info!(model = %cfg.model_id, "batch engine stopped");
|
||||
}
|
||||
|
||||
/// Emit the slot's Finish through its router and mark it for
|
||||
/// compaction.
|
||||
async fn finish_slot(slot: &mut Slot, reason: FinishReason) {
|
||||
slot.finish(reason);
|
||||
let _ = slot
|
||||
.router
|
||||
.send(RouterMsg::Finish {
|
||||
reason,
|
||||
prompt_tokens: slot.prompt_len as u32,
|
||||
completion_tokens: slot.generated.len() as u32,
|
||||
timing: FinishTiming {
|
||||
prefill_ms: slot.prefill_ms,
|
||||
decode_ms: slot.decode_start.elapsed().as_millis() as u32,
|
||||
prefill_tokens: slot.prefill_tokens,
|
||||
},
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
/// Fatal-path teardown: classify + record the poison flag, end every
|
||||
/// active stream (routers exit when their channel drops without a
|
||||
/// Finish), and drain queued requests so their clients aren't left
|
||||
/// hanging on a dead channel.
|
||||
fn fail_engine(
|
||||
cfg: &EngineConfig,
|
||||
slots: &mut Vec<Slot>,
|
||||
rx: &mut mpsc::Receiver<EngineRequest>,
|
||||
error: &anyhow::Error,
|
||||
) {
|
||||
let chain = format!("{error:#}");
|
||||
if is_device_fault(&chain) {
|
||||
cfg.poisoned.store(true, Ordering::Release);
|
||||
tracing::error!(
|
||||
model = %cfg.model_id,
|
||||
error = %chain,
|
||||
"batch engine: device fault, model marked poisoned"
|
||||
);
|
||||
} else {
|
||||
tracing::error!(
|
||||
model = %cfg.model_id,
|
||||
error = %chain,
|
||||
"batch engine: fatal error (non-device fault)"
|
||||
);
|
||||
}
|
||||
slots.clear();
|
||||
rx.close();
|
||||
while let Ok(req) = rx.try_recv() {
|
||||
drop(req); // dropping tx ends the client stream
|
||||
}
|
||||
}
|
||||
|
||||
/// Rebatch point: drop finished/hung slots, extract survivors from the
|
||||
/// live batched state, prefill every join at B=1, and assemble the new
|
||||
/// batch. On return `step == 0` and `padded_len` describes the new
|
||||
/// geometry.
|
||||
async fn rebatch(
|
||||
cfg: &EngineConfig,
|
||||
slots: &mut Vec<Slot>,
|
||||
joins: Vec<EngineRequest>,
|
||||
padded_len: &mut usize,
|
||||
step: &mut usize,
|
||||
) -> Result<()> {
|
||||
// 1. Extract survivors BEFORE any prefill clobbers the live state.
|
||||
let mut kept: Vec<Slot> = Vec::new();
|
||||
let mut extracted: Vec<(super::device_worker::jobs::KvSnapshotId, usize)> = Vec::new();
|
||||
let leavers_or_joiners = joins.len()
|
||||
+ slots
|
||||
.iter()
|
||||
.filter(|s| s.finished.is_some() || s.hangup.load(Ordering::Acquire))
|
||||
.count();
|
||||
let survivors: Vec<usize> = slots
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, s)| s.finished.is_none() && !s.hangup.load(Ordering::Acquire))
|
||||
.map(|(i, _)| i)
|
||||
.collect();
|
||||
if !survivors.is_empty() && leavers_or_joiners > 0 {
|
||||
let rows: Vec<(usize, usize)> = survivors
|
||||
.iter()
|
||||
.map(|&i| (i, slots[i].prefix_len))
|
||||
.collect();
|
||||
let ids = cfg
|
||||
.worker
|
||||
.extract_kv_rows(cfg.handle, rows, *padded_len, *step)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("extract_kv_rows: {e}"))?;
|
||||
for (&i, (id, _bytes)) in survivors.iter().zip(ids) {
|
||||
let new_len = slots[i].prefix_len + *step;
|
||||
extracted.push((id, new_len));
|
||||
}
|
||||
}
|
||||
// Drain slots preserving survivor order; finished/hung slots drop
|
||||
// here (permits release, routers wind down).
|
||||
for (order, i) in survivors.iter().enumerate() {
|
||||
debug_assert!(*i >= order);
|
||||
let mut s = slots.remove(*i - order);
|
||||
s.prefix_len += *step;
|
||||
kept.push(s);
|
||||
}
|
||||
slots.clear();
|
||||
|
||||
// 2. Prefill each join at B=1 (prefix cache + chunked prefill
|
||||
// exactly as the per-request path).
|
||||
let mut assemble: Vec<(super::device_worker::jobs::KvSnapshotId, usize)> = extracted.clone();
|
||||
for req in joins {
|
||||
let req_span = req.span.clone();
|
||||
// `None` = finished during prefill (EOS / hangup / max_new 0).
|
||||
if let Some((slot, snap_id)) = prefill_join(cfg, req).instrument_in(req_span).await? {
|
||||
assemble.push((snap_id, slot.prompt_len));
|
||||
kept.push(slot);
|
||||
}
|
||||
}
|
||||
|
||||
// 3. Assemble the new batch (or go idle).
|
||||
if kept.is_empty() {
|
||||
// Nothing active. Temp snapshots for extraction are dropped.
|
||||
for (id, _) in &assemble {
|
||||
let _ = cfg.worker.drop_kv_snapshot(cfg.handle, *id).await;
|
||||
}
|
||||
*padded_len = 0;
|
||||
*step = 0;
|
||||
return Ok(());
|
||||
}
|
||||
let seqs: Vec<(super::device_worker::jobs::KvSnapshotId, usize)> = assemble.clone();
|
||||
let new_padded = cfg
|
||||
.worker
|
||||
.assemble_kv_batch(cfg.handle, seqs)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("assemble_kv_batch: {e}"))?;
|
||||
for (id, _) in &assemble {
|
||||
let _ = cfg.worker.drop_kv_snapshot(cfg.handle, *id).await;
|
||||
}
|
||||
*padded_len = new_padded;
|
||||
*step = 0;
|
||||
*slots = kept;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Prefill one joining request at B=1 and snapshot its state for
|
||||
/// assembly. Returns `None` when the request finished during prefill
|
||||
/// (EOS as first token, `max_new == 0`, or the consumer already hung
|
||||
/// up) — its Finish has been emitted and no slot joins the batch.
|
||||
async fn prefill_join(
|
||||
cfg: &EngineConfig,
|
||||
req: EngineRequest,
|
||||
) -> Result<Option<(Slot, super::device_worker::jobs::KvSnapshotId)>> {
|
||||
use candle_transformers::generation::Sampling;
|
||||
|
||||
let EngineRequest {
|
||||
prompt_tokens,
|
||||
max_new,
|
||||
temperature,
|
||||
top_p,
|
||||
seed,
|
||||
eos_id,
|
||||
tool_schemas,
|
||||
tx,
|
||||
admit,
|
||||
span,
|
||||
} = req;
|
||||
|
||||
let mut lp = {
|
||||
let sampling = if temperature <= 0.0 {
|
||||
Sampling::ArgMax
|
||||
} else {
|
||||
match top_p {
|
||||
Some(p) => Sampling::TopP { p, temperature },
|
||||
None => Sampling::All { temperature },
|
||||
}
|
||||
};
|
||||
LogitsProcessor::from_sampling(seed, sampling)
|
||||
};
|
||||
|
||||
let prefix_cache = cfg.prefix_cache.as_deref();
|
||||
let prompt_len = prompt_tokens.len();
|
||||
let prefill_start = std::time::Instant::now();
|
||||
let reused =
|
||||
restore_or_clear_via_worker(&cfg.worker, cfg.handle, prefix_cache, &prompt_tokens).await?;
|
||||
let cut = if prefix_cache.is_some() {
|
||||
stable_snapshot_cut(&prompt_tokens, cfg.tokenizer.token_to_id("<|im_start|>"))
|
||||
.filter(|&c| c > reused)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let logits_vec = match cut {
|
||||
Some(c) => {
|
||||
chunked_prefill_via_worker(&cfg.worker, cfg.handle, &prompt_tokens[..c], reused)
|
||||
.await?;
|
||||
store_prefix_snapshot_via_worker(
|
||||
&cfg.worker,
|
||||
cfg.handle,
|
||||
prefix_cache,
|
||||
prompt_tokens[..c].to_vec(),
|
||||
)
|
||||
.await;
|
||||
chunked_prefill_via_worker(&cfg.worker, cfg.handle, &prompt_tokens, c).await?
|
||||
}
|
||||
None => chunked_prefill_via_worker(&cfg.worker, cfg.handle, &prompt_tokens, reused).await?,
|
||||
};
|
||||
let prefill_elapsed = prefill_start.elapsed();
|
||||
cfg.prefill_rate.record(prompt_len, prefill_elapsed);
|
||||
|
||||
// First token from the prefill logits.
|
||||
let generated: Vec<u32> = Vec::new();
|
||||
let logits = Tensor::new(logits_vec.as_slice(), &Device::Cpu)?;
|
||||
let first = match sample_with_penalty(&logits, &generated, &mut lp) {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
let health = logits_health_slice(&logits_vec);
|
||||
tracing::warn!(
|
||||
?health,
|
||||
"batch engine: prefill sample failed; logits unhealthy"
|
||||
);
|
||||
return Err(e);
|
||||
}
|
||||
};
|
||||
|
||||
// Router task for this slot.
|
||||
let hangup = Arc::new(AtomicBool::new(false));
|
||||
let (router_tx, router_rx) = mpsc::channel::<RouterMsg>(1024);
|
||||
let starts_in_reasoning = prompt_opens_reasoning(&prompt_tokens, cfg.reasoning_tokens.as_ref());
|
||||
tokio::spawn(
|
||||
run_router(
|
||||
cfg.tokenizer.clone(),
|
||||
cfg.reasoning_tokens.clone(),
|
||||
cfg.tool_call_tokens.clone(),
|
||||
tool_schemas,
|
||||
starts_in_reasoning,
|
||||
tx,
|
||||
Arc::clone(&hangup),
|
||||
router_rx,
|
||||
)
|
||||
.instrument_in(span.clone()),
|
||||
);
|
||||
|
||||
let mut slot = Slot {
|
||||
prefix_len: prompt_len,
|
||||
prompt_len,
|
||||
generated,
|
||||
next_token: first,
|
||||
max_new,
|
||||
eos_id,
|
||||
lp,
|
||||
router: router_tx,
|
||||
hangup,
|
||||
finished: None,
|
||||
prefill_ms: prefill_elapsed.as_millis() as u32,
|
||||
prefill_tokens: prompt_len as u32,
|
||||
decode_start: std::time::Instant::now(),
|
||||
_admit: admit,
|
||||
};
|
||||
|
||||
// First-token bookkeeping mirrors the B=1 path: EOS finishes
|
||||
// without routing; max_new bounds include the first token.
|
||||
if Some(first) == slot.eos_id || slot.max_new == 0 {
|
||||
let reason = if slot.max_new == 0 {
|
||||
FinishReason::Length
|
||||
} else {
|
||||
FinishReason::Stop
|
||||
};
|
||||
finish_slot(&mut slot, reason).await;
|
||||
return Ok(None);
|
||||
}
|
||||
slot.generated.push(first);
|
||||
if slot.router.send(RouterMsg::Token(first)).await.is_err() {
|
||||
return Ok(None); // consumer already gone
|
||||
}
|
||||
if slot.generated.len() >= slot.max_new {
|
||||
finish_slot(&mut slot, FinishReason::Length).await;
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// Snapshot the freshly prefilled state for assembly.
|
||||
let (snap_id, _bytes) = cfg
|
||||
.worker
|
||||
.snapshot_kv(cfg.handle)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("snapshot after prefill: {e}"))?;
|
||||
Ok(Some((slot, snap_id)))
|
||||
}
|
||||
|
||||
/// Per-slot router: owns the incremental detokenizer and the
|
||||
/// reasoning/tool-call state machine (the same logic as the
|
||||
/// `route_token!` macro in the B=1 stream path) and emits
|
||||
/// `InferenceEvent`s on the request's channel. Sets `hangup` and
|
||||
/// drains silently once the consumer goes away.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn run_router(
|
||||
tokenizer: Tokenizer,
|
||||
reasoning_tokens: Option<ReasoningTokenPair>,
|
||||
tool_call_tokens: Option<ToolCallTokenPair>,
|
||||
tool_schemas: ToolSchemas,
|
||||
starts_in_reasoning: bool,
|
||||
tx: mpsc::Sender<InferenceEvent>,
|
||||
hangup: Arc<AtomicBool>,
|
||||
mut rx: mpsc::Receiver<RouterMsg>,
|
||||
) {
|
||||
let mut decode_stream = tokenizer.decode_stream(true);
|
||||
let mut in_reasoning = starts_in_reasoning;
|
||||
let mut reasoning_token_count: u32 = 0;
|
||||
let mut in_tool_call = false;
|
||||
let mut tool_call_buf = String::new();
|
||||
let mut tool_call_idx: usize = 0;
|
||||
let mut emitted_tool_call = false;
|
||||
let mut consumer_alive = true;
|
||||
|
||||
while let Some(msg) = rx.recv().await {
|
||||
match msg {
|
||||
RouterMsg::Token(nt) => {
|
||||
if !consumer_alive {
|
||||
continue; // drain
|
||||
}
|
||||
'route: {
|
||||
match handle_tool_call_marker(
|
||||
nt,
|
||||
tool_call_tokens.as_ref(),
|
||||
&mut in_tool_call,
|
||||
&mut tool_call_buf,
|
||||
) {
|
||||
ToolCallMarker::Enter => break 'route,
|
||||
ToolCallMarker::Exit { buffer } => {
|
||||
let idx = tool_call_idx;
|
||||
tool_call_idx += 1;
|
||||
match parse_tool_call_body(&buffer, idx, &tool_schemas) {
|
||||
Some((id, name, arguments)) => {
|
||||
emitted_tool_call = true;
|
||||
if tx
|
||||
.send(InferenceEvent::ToolCall {
|
||||
index: idx,
|
||||
id,
|
||||
name,
|
||||
arguments,
|
||||
})
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
consumer_alive = false;
|
||||
}
|
||||
}
|
||||
None => {
|
||||
let open = tool_call_tokens
|
||||
.as_ref()
|
||||
.map(|p| p.open_text.as_str())
|
||||
.unwrap_or("<tool_call>");
|
||||
let close = tool_call_tokens
|
||||
.as_ref()
|
||||
.map(|p| p.close_text.as_str())
|
||||
.unwrap_or("</tool_call>");
|
||||
let raw = format!("{open}{buffer}{close}");
|
||||
if !emit_delta(&raw, &tx, in_reasoning).await {
|
||||
consumer_alive = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
break 'route;
|
||||
}
|
||||
ToolCallMarker::None => {}
|
||||
}
|
||||
if in_tool_call {
|
||||
match decode_stream.step(nt) {
|
||||
Ok(Some(s)) => tool_call_buf.push_str(&s),
|
||||
Ok(None) => {}
|
||||
Err(e) => tracing::warn!(
|
||||
error = %e,
|
||||
"decode_stream step failed (in tool_call)"
|
||||
),
|
||||
}
|
||||
break 'route;
|
||||
}
|
||||
if handle_reasoning_marker(nt, reasoning_tokens.as_ref(), &mut in_reasoning) {
|
||||
break 'route;
|
||||
}
|
||||
if in_reasoning {
|
||||
reasoning_token_count += 1;
|
||||
}
|
||||
match decode_stream.step(nt) {
|
||||
Ok(Some(delta)) => {
|
||||
if !emit_delta(&delta, &tx, in_reasoning).await {
|
||||
consumer_alive = false;
|
||||
}
|
||||
}
|
||||
Ok(None) => {}
|
||||
Err(e) => tracing::warn!(error = %e, "decode_stream step failed"),
|
||||
}
|
||||
}
|
||||
if !consumer_alive {
|
||||
hangup.store(true, Ordering::Release);
|
||||
}
|
||||
}
|
||||
RouterMsg::Finish {
|
||||
mut reason,
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
timing,
|
||||
} => {
|
||||
if emitted_tool_call && reason == FinishReason::Stop {
|
||||
reason = FinishReason::ToolCalls;
|
||||
}
|
||||
let _ = tx
|
||||
.send(InferenceEvent::Finish {
|
||||
reason,
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
reasoning_tokens: reasoning_token_count,
|
||||
timing: Some(timing),
|
||||
})
|
||||
.await;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `tracing::Instrument` without importing the trait at every use
|
||||
/// site.
|
||||
trait InstrumentExt: Sized + std::future::Future {
|
||||
fn instrument_in(self, span: tracing::Span) -> tracing::instrument::Instrumented<Self> {
|
||||
tracing::Instrument::instrument(self, span)
|
||||
}
|
||||
}
|
||||
impl<F: std::future::Future> InstrumentExt for F {}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::config::AdmissionConfig;
|
||||
use crate::harness::admission::AdmissionController;
|
||||
|
||||
/// A WordLevel tokenizer whose vocab covers the whole fixture
|
||||
/// vocab (`w0`..`w511`), so every decoded token maps to a unique
|
||||
/// word and the emitted text uniquely encodes the token sequence.
|
||||
fn tiny_tokenizer(vocab_size: usize) -> Tokenizer {
|
||||
// WordLevel's builder vocab type (AHashMap) is private, so go
|
||||
// through the vocab-file loader instead.
|
||||
let vocab: std::collections::HashMap<String, u32> = (0..vocab_size as u32)
|
||||
.map(|i| (format!("w{i}"), i))
|
||||
.collect();
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let path = dir.path().join("vocab.json");
|
||||
std::fs::write(&path, serde_json::to_string(&vocab).expect("vocab json"))
|
||||
.expect("write vocab");
|
||||
let model = tokenizers::models::wordlevel::WordLevel::from_file(
|
||||
path.to_str().expect("utf8 path"),
|
||||
"w0".into(),
|
||||
)
|
||||
.expect("build WordLevel");
|
||||
Tokenizer::new(tokenizers::ModelWrapper::WordLevel(model))
|
||||
}
|
||||
|
||||
async fn collect_run(
|
||||
engine: &EngineHandle,
|
||||
admission: &AdmissionController,
|
||||
prompt: Vec<u32>,
|
||||
max_new: usize,
|
||||
) -> (String, u32, FinishReason) {
|
||||
let admit = admission.enter(None).await.expect("admitted");
|
||||
let (tx, mut rx) = mpsc::channel::<InferenceEvent>(32);
|
||||
engine
|
||||
.submit(EngineRequest {
|
||||
prompt_tokens: prompt,
|
||||
max_new,
|
||||
temperature: 0.0, // greedy — deterministic
|
||||
top_p: None,
|
||||
seed: 0,
|
||||
eos_id: None,
|
||||
tool_schemas: ToolSchemas::new(),
|
||||
tx,
|
||||
admit,
|
||||
span: tracing::Span::none(),
|
||||
})
|
||||
.await
|
||||
.expect("submit");
|
||||
let mut text = String::new();
|
||||
loop {
|
||||
match rx.recv().await {
|
||||
Some(InferenceEvent::TextDelta(d)) | Some(InferenceEvent::ReasoningDelta(d)) => {
|
||||
text.push_str(&d)
|
||||
}
|
||||
Some(InferenceEvent::Finish {
|
||||
reason,
|
||||
completion_tokens,
|
||||
..
|
||||
}) => return (text, completion_tokens, reason),
|
||||
Some(_) => {}
|
||||
None => panic!("stream ended without Finish"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The engine's gold test: three greedy requests submitted
|
||||
/// concurrently (batched lockstep decode, ragged prompts, joins
|
||||
/// mid-flight) must produce byte-identical output to the same
|
||||
/// requests run one-at-a-time through the same engine.
|
||||
#[tokio::test]
|
||||
async fn concurrent_engine_output_matches_sequential() {
|
||||
let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
|
||||
.join("tests/fixtures/numerical/qwen3_next-tiny");
|
||||
if !fixture.join("model.safetensors").exists() {
|
||||
eprintln!("SKIP concurrent_engine_output_matches_sequential: fixture not generated");
|
||||
return;
|
||||
}
|
||||
let worker = DeviceWorkerHandle::spawn(0).expect("spawn worker");
|
||||
let handle = worker
|
||||
.load_dense(
|
||||
fixture.join("config.json"),
|
||||
vec![fixture.join("model.safetensors")],
|
||||
"qwen3_next-tiny".into(),
|
||||
)
|
||||
.await
|
||||
.expect("load fixture");
|
||||
|
||||
let admission_cfg = AdmissionConfig {
|
||||
max_in_flight: 3,
|
||||
..Default::default()
|
||||
};
|
||||
let admission = AdmissionController::new(&admission_cfg);
|
||||
let engine = EngineHandle::spawn(EngineConfig {
|
||||
model_id: "qwen3_next-tiny".into(),
|
||||
worker: Arc::clone(&worker),
|
||||
handle,
|
||||
tokenizer: tiny_tokenizer(512),
|
||||
prefix_cache: None,
|
||||
prefill_rate: Arc::new(PrefillRateEma::new()),
|
||||
reasoning_tokens: None,
|
||||
tool_call_tokens: None,
|
||||
poisoned: Arc::new(AtomicBool::new(false)),
|
||||
inference_lock: Arc::new(tokio::sync::Mutex::new(())),
|
||||
max_slots: 3,
|
||||
});
|
||||
|
||||
let prompts: [&[u32]; 3] = [&[1, 2, 3], &[4, 5], &[7, 3, 2, 5, 6]];
|
||||
let max_new = 6;
|
||||
|
||||
// Sequential reference: one at a time through the same engine.
|
||||
let mut expected = Vec::new();
|
||||
for p in prompts {
|
||||
expected.push(collect_run(&engine, &admission, p.to_vec(), max_new).await);
|
||||
}
|
||||
|
||||
// Concurrent: all three at once — they batch.
|
||||
let futs: Vec<_> = prompts
|
||||
.iter()
|
||||
.map(|p| collect_run(&engine, &admission, p.to_vec(), max_new))
|
||||
.collect();
|
||||
let got = futures::future::join_all(futs).await;
|
||||
|
||||
for (i, (want, got)) in expected.iter().zip(got.iter()).enumerate() {
|
||||
assert_eq!(want.2, got.2, "request {i} finish reason");
|
||||
assert_eq!(want.1, got.1, "request {i} completion tokens");
|
||||
assert_eq!(want.0, got.0, "request {i} text");
|
||||
}
|
||||
assert!(
|
||||
expected
|
||||
.iter()
|
||||
.all(|(t, n, _)| !t.is_empty() && *n as usize == max_new),
|
||||
"greedy runs should hit the length cap with non-empty text: {expected:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -6,6 +6,7 @@ pub mod candle;
|
||||
pub mod chat_template;
|
||||
pub mod context_limit;
|
||||
pub mod device_worker;
|
||||
pub mod engine;
|
||||
pub mod prefix_cache;
|
||||
pub mod preflight;
|
||||
pub mod preprocess;
|
||||
|
||||
@@ -156,6 +156,15 @@ pub struct ColumnParallelLinear {
|
||||
}
|
||||
|
||||
impl ColumnParallelLinear {
|
||||
/// Wrap an already-materialised per-rank weight slice (used by the
|
||||
/// fused-checkpoint loaders that de-interleave a rank's regions
|
||||
/// before construction, #92).
|
||||
pub fn from_weight(weight: Tensor, quant: Option<GgmlDType>) -> Result<Self> {
|
||||
Ok(Self {
|
||||
inner: MaybeQuantLinear::from_weight(weight, quant)?,
|
||||
})
|
||||
}
|
||||
|
||||
/// Load this rank's column-parallel slice from a
|
||||
/// `ShardedVarBuilder`. The provided `vb` must already be `pp`-ed
|
||||
/// to the layer's path (e.g. `vb.pp("model.layers.0.self_attn.q_proj")`).
|
||||
|
||||
@@ -36,7 +36,6 @@ use candle_core::safetensors::MmapedSafetensors;
|
||||
use candle_core::{DType, Device, IndexOp, Module, Tensor};
|
||||
use candle_nn::var_builder::ShardedVarBuilder;
|
||||
use candle_nn::{Embedding, kv_cache::ConcatKvCache};
|
||||
use candle_transformers::utils::repeat_kv;
|
||||
use std::sync::Arc;
|
||||
|
||||
#[cfg(feature = "cuda")]
|
||||
@@ -174,21 +173,88 @@ impl TpQwen3_5GatedDeltaNet {
|
||||
// source on consumer GPUs near their VRAM ceiling.
|
||||
let dtype = vb.dtype();
|
||||
let device = vb.device().clone();
|
||||
let in_proj_qkv_name = format!("{}.in_proj_qkv.weight", vb.prefix());
|
||||
let in_proj_qkv_weight = super::fused_load::load_fused_qkv_2d(
|
||||
mmap,
|
||||
&in_proj_qkv_name,
|
||||
hidden_size,
|
||||
key_dim,
|
||||
value_dim,
|
||||
rank,
|
||||
world_size,
|
||||
dtype,
|
||||
&device,
|
||||
)?;
|
||||
let in_proj_qkv =
|
||||
super::tp_linear::MaybeQuantLinear::from_weight(in_proj_qkv_weight, quant)
|
||||
.with_context(|| format!("wrap fused in_proj_qkv for '{}'", vb.prefix()))?;
|
||||
|
||||
// Two checkpoint layouts (#92, mirroring the single-GPU loader):
|
||||
// - Qwen3.6: separate `in_proj_qkv` ([Q|K|V] contiguous) +
|
||||
// `in_proj_z`/`in_proj_b`/`in_proj_a` — per-region mmap
|
||||
// slicing for qkv, uniform column-parallel for the rest.
|
||||
// - Qwen3-Next 80B-A3B: fused `in_proj_qkvz` + `in_proj_ba`,
|
||||
// interleaved per key-head group. K- and V-heads shard
|
||||
// uniformly together, so each rank owns a CONTIGUOUS span of
|
||||
// whole groups — a uniform dim-0 shard of the fused tensor is
|
||||
// exactly the rank's groups, and the per-rank de-interleave
|
||||
// (`split_fused_qkvz`/`split_fused_ba` with per-rank head
|
||||
// counts) restores the contiguous [Q|K|V] + Z / B + A layout
|
||||
// the forward path expects.
|
||||
let (in_proj_qkv, in_proj_z, in_proj_b, in_proj_a) =
|
||||
if vb.contains_tensor("in_proj_qkvz.weight") {
|
||||
let qkvz_slice = vb
|
||||
.pp("in_proj_qkvz")
|
||||
.get_with_hints((), "weight", super::tp_linear::shard(0, rank, world_size))
|
||||
.with_context(|| format!("load '{}/in_proj_qkvz' rank slice", vb.prefix()))?;
|
||||
let ba_slice = vb
|
||||
.pp("in_proj_ba")
|
||||
.get_with_hints((), "weight", super::tp_linear::shard(0, rank, world_size))
|
||||
.with_context(|| format!("load '{}/in_proj_ba' rank slice", vb.prefix()))?;
|
||||
let (qkv_w, z_w) = crate::harness::arch::qwen3_5::linear_attn::split_fused_qkvz(
|
||||
&qkvz_slice,
|
||||
per_rank_num_k_heads,
|
||||
per_rank_num_v_heads,
|
||||
head_k_dim,
|
||||
head_v_dim,
|
||||
)
|
||||
.with_context(|| format!("de-interleave '{}/in_proj_qkvz'", vb.prefix()))?;
|
||||
let (b_w, a_w) = crate::harness::arch::qwen3_5::linear_attn::split_fused_ba(
|
||||
&ba_slice,
|
||||
per_rank_num_k_heads,
|
||||
per_rank_num_v_heads,
|
||||
)
|
||||
.with_context(|| format!("de-interleave '{}/in_proj_ba'", vb.prefix()))?;
|
||||
(
|
||||
super::tp_linear::MaybeQuantLinear::from_weight(qkv_w, quant)
|
||||
.with_context(|| format!("wrap fused in_proj_qkv '{}'", vb.prefix()))?,
|
||||
ColumnParallelLinear::from_weight(z_w, quant)?,
|
||||
ColumnParallelLinear::from_weight(b_w, quant)?,
|
||||
ColumnParallelLinear::from_weight(a_w, quant)?,
|
||||
)
|
||||
} else {
|
||||
let in_proj_qkv_name = format!("{}.in_proj_qkv.weight", vb.prefix());
|
||||
let in_proj_qkv_weight = super::fused_load::load_fused_qkv_2d(
|
||||
mmap,
|
||||
&in_proj_qkv_name,
|
||||
hidden_size,
|
||||
key_dim,
|
||||
value_dim,
|
||||
rank,
|
||||
world_size,
|
||||
dtype,
|
||||
&device,
|
||||
)?;
|
||||
(
|
||||
super::tp_linear::MaybeQuantLinear::from_weight(in_proj_qkv_weight, quant)
|
||||
.with_context(|| format!("wrap fused in_proj_qkv for '{}'", vb.prefix()))?,
|
||||
// in_proj_z: hidden → value_dim, sharded along value_dim
|
||||
// (V-head); in_proj_b / in_proj_a: hidden → num_v_heads.
|
||||
ColumnParallelLinear::load_with_quant(
|
||||
&vb.pp("in_proj_z"),
|
||||
rank,
|
||||
world_size,
|
||||
quant,
|
||||
)?,
|
||||
ColumnParallelLinear::load_with_quant(
|
||||
&vb.pp("in_proj_b"),
|
||||
rank,
|
||||
world_size,
|
||||
quant,
|
||||
)?,
|
||||
ColumnParallelLinear::load_with_quant(
|
||||
&vb.pp("in_proj_a"),
|
||||
rank,
|
||||
world_size,
|
||||
quant,
|
||||
)?,
|
||||
)
|
||||
};
|
||||
|
||||
let conv1d_name = format!("{}.conv1d.weight", vb.prefix());
|
||||
let conv1d_weight = super::fused_load::load_fused_qkv_3d(
|
||||
@@ -204,16 +270,6 @@ impl TpQwen3_5GatedDeltaNet {
|
||||
&device,
|
||||
)?;
|
||||
|
||||
// ----- Uniformly-sharded projections (along output dim 0). -----
|
||||
// in_proj_z: hidden → value_dim, sharded along value_dim (V-head).
|
||||
let in_proj_z =
|
||||
ColumnParallelLinear::load_with_quant(&vb.pp("in_proj_z"), rank, world_size, quant)?;
|
||||
// in_proj_b, in_proj_a: hidden → num_v_heads, sharded along output.
|
||||
let in_proj_b =
|
||||
ColumnParallelLinear::load_with_quant(&vb.pp("in_proj_b"), rank, world_size, quant)?;
|
||||
let in_proj_a =
|
||||
ColumnParallelLinear::load_with_quant(&vb.pp("in_proj_a"), rank, world_size, quant)?;
|
||||
|
||||
// ----- Per-V-head 1D params (sharded uniformly). -----
|
||||
let a_log = vb
|
||||
.get_with_hints((), "A_log", super::tp_linear::shard(0, rank, world_size))
|
||||
@@ -596,16 +652,19 @@ impl TpQwen3_5Attention {
|
||||
|
||||
let (q, k) = self.rotary.apply_cos_sin(&q, &k, cos, sin)?;
|
||||
let (k, v) = self.kv_cache.append(&k, &v)?;
|
||||
let k = repeat_kv(k, self.num_kv_groups)?.contiguous()?;
|
||||
let v = repeat_kv(v, self.num_kv_groups)?.contiguous()?;
|
||||
|
||||
// Attention core — FlashAttention when available, eager
|
||||
// GQA-repeat + masked softmax otherwise (#95). Per-rank heads,
|
||||
// same kernel semantics as the single-GPU path.
|
||||
let scale = 1.0_f64 / (self.head_dim as f64).sqrt();
|
||||
let mut scores = (q.matmul(&k.transpose(2, 3)?)? * scale)?;
|
||||
if let Some(m) = attn_mask {
|
||||
scores = scores.broadcast_add(m)?;
|
||||
}
|
||||
let probs = candle_nn::ops::softmax_last_dim(&scores)?;
|
||||
let ctx = probs.matmul(&v)?;
|
||||
let ctx = crate::harness::arch::qwen3_5::full_attn::attention_context(
|
||||
&q,
|
||||
&k,
|
||||
&v,
|
||||
attn_mask,
|
||||
self.num_kv_groups,
|
||||
scale,
|
||||
)?;
|
||||
|
||||
let ctx = ctx
|
||||
.transpose(1, 2)?
|
||||
@@ -735,6 +794,502 @@ impl Module for TpQwen3_5MLP {
|
||||
}
|
||||
}
|
||||
|
||||
// ─── MoE FFN (qwen3_next 80B-A3B, #92) ──────────────────────────────
|
||||
|
||||
/// One routed expert's per-rank slice: `gate_proj`/`up_proj`
|
||||
/// column-sharded along `moe_intermediate_size`, `down_proj`
|
||||
/// input-sharded — `forward_partial` yields a PARTIAL hidden output.
|
||||
///
|
||||
/// Deliberately NOT built from [`RowParallelLinear`]: that embeds an
|
||||
/// AllReduce per call, which with top-10 routing would mean ten
|
||||
/// collectives per layer. Partial sums from every selected expert and
|
||||
/// the shared expert add linearly, so the whole MoE block needs
|
||||
/// exactly ONE AllReduce at the end — preserving the existing
|
||||
/// one-reduce-per-FFN pattern.
|
||||
struct TpExpert {
|
||||
gate_proj: super::tp_linear::MaybeQuantLinear,
|
||||
up_proj: super::tp_linear::MaybeQuantLinear,
|
||||
down_proj: super::tp_linear::MaybeQuantLinear,
|
||||
}
|
||||
|
||||
impl TpExpert {
|
||||
fn load(
|
||||
vb: &ShardedVarBuilder,
|
||||
rank: u32,
|
||||
world_size: u32,
|
||||
quant: Option<GgmlDType>,
|
||||
) -> Result<Self> {
|
||||
let col = |name: &str| -> Result<super::tp_linear::MaybeQuantLinear> {
|
||||
let w = vb
|
||||
.pp(name)
|
||||
.get_with_hints((), "weight", super::tp_linear::shard(0, rank, world_size))
|
||||
.with_context(|| format!("load expert '{}/{name}'", vb.prefix()))?;
|
||||
super::tp_linear::MaybeQuantLinear::from_weight(w, quant)
|
||||
};
|
||||
let gate_proj = col("gate_proj")?;
|
||||
let up_proj = col("up_proj")?;
|
||||
let down_w = vb
|
||||
.pp("down_proj")
|
||||
.get_with_hints((), "weight", super::tp_linear::shard(1, rank, world_size))
|
||||
.with_context(|| format!("load expert '{}/down_proj'", vb.prefix()))?;
|
||||
let down_proj = super::tp_linear::MaybeQuantLinear::from_weight(down_w, quant)?;
|
||||
Ok(Self {
|
||||
gate_proj,
|
||||
up_proj,
|
||||
down_proj,
|
||||
})
|
||||
}
|
||||
|
||||
/// SwiGLU over this rank's intermediate slice; the output is a
|
||||
/// partial sum awaiting the block-end AllReduce.
|
||||
fn forward_partial(&self, x: &Tensor) -> candle_core::Result<Tensor> {
|
||||
let lhs = candle_nn::ops::silu(&self.gate_proj.forward(x)?)?;
|
||||
let rhs = self.up_proj.forward(x)?;
|
||||
self.down_proj.forward(&(lhs * rhs)?)
|
||||
}
|
||||
}
|
||||
|
||||
/// How the routed experts are stored and dispatched (#92 slice 4).
|
||||
///
|
||||
/// `Scatter` is the correctness-first path: per-expert
|
||||
/// `MaybeQuantLinear`s driven by a host-side token scatter. Correct
|
||||
/// everywhere, but at batch-1 decode it costs a GPU→CPU routing sync
|
||||
/// plus ~top-k tiny GEMV launches per layer per token — measured at
|
||||
/// **4.3 tok/s** on the 80B (vs ~27 for the dense 27B).
|
||||
///
|
||||
/// `Fused` holds each projection as ONE stacked per-rank `QTensor`
|
||||
/// (`[num_experts, out/ws, in]`) driven by candle-nn's grouped-GEMM
|
||||
/// GGUF kernels (`moe_gemm_gguf`): routing, index sort, and all expert
|
||||
/// GEMMs stay on-device — three kernel launches per layer regardless
|
||||
/// of k. Chosen at load on CUDA devices when ISQ is active and the
|
||||
/// block dims satisfy the GGUF constraints; `NEURON_MOE_FUSED=0`
|
||||
/// forces Scatter as an escape hatch / A-B lever.
|
||||
enum TpExpertStore {
|
||||
Scatter(Vec<TpExpert>),
|
||||
#[cfg(feature = "cuda")]
|
||||
Fused {
|
||||
gate_experts: Arc<candle_core::quantized::QTensor>,
|
||||
up_experts: Arc<candle_core::quantized::QTensor>,
|
||||
down_experts: Arc<candle_core::quantized::QTensor>,
|
||||
num_experts: usize,
|
||||
},
|
||||
}
|
||||
|
||||
/// TP counterpart of `arch::qwen3_5::moe::Qwen3_5MoeBlock`. The router
|
||||
/// and the shared-expert sigmoid gate are replicated (tiny; every rank
|
||||
/// computes identical routing with zero communication); expert FFNs
|
||||
/// shard per-expert along the intermediate dim; one AllReduce at block
|
||||
/// end recovers the full activation.
|
||||
pub(crate) struct TpQwen3_5MoeBlock {
|
||||
gate: candle_nn::Linear,
|
||||
experts: TpExpertStore,
|
||||
shared_expert: Option<TpExpert>,
|
||||
shared_expert_gate: Option<candle_nn::Linear>,
|
||||
num_experts_per_tok: usize,
|
||||
norm_topk_prob: bool,
|
||||
#[cfg(feature = "cuda")]
|
||||
all_reduce: super::all_reduce::AllReduce,
|
||||
needs_reduce: bool,
|
||||
}
|
||||
|
||||
impl TpQwen3_5MoeBlock {
|
||||
fn check_and_gate(
|
||||
cfg: &TextConfig,
|
||||
vb: &ShardedVarBuilder,
|
||||
world_size: u32,
|
||||
) -> Result<candle_nn::Linear> {
|
||||
let ws = world_size as usize;
|
||||
if !cfg.moe_intermediate_size.is_multiple_of(ws) {
|
||||
bail!(
|
||||
"moe_intermediate_size {} not divisible by world_size {ws}",
|
||||
cfg.moe_intermediate_size
|
||||
);
|
||||
}
|
||||
if cfg.shared_expert_intermediate_size > 0
|
||||
&& !cfg.shared_expert_intermediate_size.is_multiple_of(ws)
|
||||
{
|
||||
bail!(
|
||||
"shared_expert_intermediate_size {} not divisible by world_size {ws}",
|
||||
cfg.shared_expert_intermediate_size
|
||||
);
|
||||
}
|
||||
let gate_w = load_replicated(&vb.pp("gate"), (cfg.num_experts, cfg.hidden_size), "weight")?;
|
||||
Ok(candle_nn::Linear::new(gate_w, None))
|
||||
}
|
||||
|
||||
fn load_experts_and_shared(
|
||||
cfg: &TextConfig,
|
||||
vb: &ShardedVarBuilder,
|
||||
rank: u32,
|
||||
world_size: u32,
|
||||
quant: Option<GgmlDType>,
|
||||
) -> Result<(TpExpertStore, Option<TpExpert>, Option<candle_nn::Linear>)> {
|
||||
let experts = Self::load_expert_store(cfg, vb, rank, world_size, quant)?;
|
||||
let (shared_expert, shared_expert_gate) = if cfg.shared_expert_intermediate_size > 0 {
|
||||
let shared = TpExpert::load(&vb.pp("shared_expert"), rank, world_size, quant)
|
||||
.context("load TP shared_expert")?;
|
||||
let gate_w =
|
||||
load_replicated(&vb.pp("shared_expert_gate"), (1, cfg.hidden_size), "weight")?;
|
||||
(Some(shared), Some(candle_nn::Linear::new(gate_w, None)))
|
||||
} else {
|
||||
(None, None)
|
||||
};
|
||||
Ok((experts, shared_expert, shared_expert_gate))
|
||||
}
|
||||
|
||||
fn load_expert_store(
|
||||
cfg: &TextConfig,
|
||||
vb: &ShardedVarBuilder,
|
||||
rank: u32,
|
||||
world_size: u32,
|
||||
quant: Option<GgmlDType>,
|
||||
) -> Result<TpExpertStore> {
|
||||
#[cfg(feature = "cuda")]
|
||||
if Self::fused_eligible(cfg, vb, world_size, quant) {
|
||||
return Self::load_fused_experts(cfg, vb, rank, world_size, quant);
|
||||
}
|
||||
let experts_vb = vb.pp("experts");
|
||||
let mut experts = Vec::with_capacity(cfg.num_experts);
|
||||
for i in 0..cfg.num_experts {
|
||||
experts.push(
|
||||
TpExpert::load(&experts_vb.pp(i), rank, world_size, quant)
|
||||
.with_context(|| format!("load TP expert {i}"))?,
|
||||
);
|
||||
}
|
||||
Ok(TpExpertStore::Scatter(experts))
|
||||
}
|
||||
|
||||
/// Whether the fused grouped-GEMM path can serve this block: CUDA
|
||||
/// device, ISQ active with a kernel-supported GGML dtype, GGUF
|
||||
/// block alignment on both GEMM K dims (hidden for gate/up, the
|
||||
/// per-rank intermediate slice for down), and not vetoed by the
|
||||
/// `NEURON_MOE_FUSED=0` escape hatch.
|
||||
#[cfg(feature = "cuda")]
|
||||
fn fused_eligible(
|
||||
cfg: &TextConfig,
|
||||
vb: &ShardedVarBuilder,
|
||||
world_size: u32,
|
||||
quant: Option<GgmlDType>,
|
||||
) -> bool {
|
||||
if std::env::var("NEURON_MOE_FUSED").is_ok_and(|v| v == "0" || v == "false") {
|
||||
tracing::info!("NEURON_MOE_FUSED=0 — MoE block using scatter dispatch");
|
||||
return false;
|
||||
}
|
||||
if !vb.device().is_cuda() {
|
||||
return false;
|
||||
}
|
||||
let Some(q) = quant else { return false };
|
||||
let kernel_supported = matches!(
|
||||
q,
|
||||
GgmlDType::Q8_0
|
||||
| GgmlDType::Q4K
|
||||
| GgmlDType::Q2K
|
||||
| GgmlDType::Q3K
|
||||
| GgmlDType::Q5K
|
||||
| GgmlDType::Q6K
|
||||
);
|
||||
let per_rank_inter = cfg.moe_intermediate_size / world_size as usize;
|
||||
let aligned = cfg.hidden_size.is_multiple_of(q.block_size())
|
||||
&& per_rank_inter.is_multiple_of(q.block_size());
|
||||
if !kernel_supported || !aligned {
|
||||
tracing::warn!(
|
||||
quant = ?q,
|
||||
hidden = cfg.hidden_size,
|
||||
per_rank_inter,
|
||||
"MoE fused path ineligible — falling back to scatter dispatch"
|
||||
);
|
||||
return false;
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
/// Build the stacked per-rank expert QTensors for the fused path:
|
||||
/// read each expert's rank slice, stack into `[E, out/ws, in]`
|
||||
/// (gate/up) and `[E, hidden, inter/ws]` (down), ISQ the stack in
|
||||
/// one parallel pass per projection. Transient cost: one bf16
|
||||
/// stack per projection (~0.5 GB at 80B dims) alive on-device
|
||||
/// until its QTensor replaces it.
|
||||
#[cfg(feature = "cuda")]
|
||||
fn load_fused_experts(
|
||||
cfg: &TextConfig,
|
||||
vb: &ShardedVarBuilder,
|
||||
rank: u32,
|
||||
world_size: u32,
|
||||
quant: Option<GgmlDType>,
|
||||
) -> Result<TpExpertStore> {
|
||||
let q = quant.expect("fused_eligible ensured quant");
|
||||
let experts_vb = vb.pp("experts");
|
||||
|
||||
let stack_proj = |name: &str, shard_dim: usize| -> Result<Tensor> {
|
||||
let mut slices = Vec::with_capacity(cfg.num_experts);
|
||||
for i in 0..cfg.num_experts {
|
||||
let w = experts_vb
|
||||
.pp(i)
|
||||
.pp(name)
|
||||
.get_with_hints(
|
||||
(),
|
||||
"weight",
|
||||
super::tp_linear::shard(shard_dim, rank, world_size),
|
||||
)
|
||||
.with_context(|| format!("load expert {i} '{name}' rank slice"))?;
|
||||
slices.push(w);
|
||||
}
|
||||
Tensor::stack(&slices, 0).with_context(|| format!("stack {name} experts"))
|
||||
};
|
||||
|
||||
let quantize =
|
||||
|stack: Tensor, name: &str| -> Result<Arc<candle_core::quantized::QTensor>> {
|
||||
let qt = super::isq::quantize_parallel(&stack, q)
|
||||
.with_context(|| format!("ISQ {name} expert stack to {q:?}"))?;
|
||||
Ok(Arc::new(qt))
|
||||
};
|
||||
|
||||
let gate_experts = quantize(stack_proj("gate_proj", 0)?, "gate_proj")?;
|
||||
let up_experts = quantize(stack_proj("up_proj", 0)?, "up_proj")?;
|
||||
let down_experts = quantize(stack_proj("down_proj", 1)?, "down_proj")?;
|
||||
|
||||
Ok(TpExpertStore::Fused {
|
||||
gate_experts,
|
||||
up_experts,
|
||||
down_experts,
|
||||
num_experts: cfg.num_experts,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(feature = "cuda")]
|
||||
pub fn load(
|
||||
cfg: &TextConfig,
|
||||
vb: &ShardedVarBuilder,
|
||||
rank: u32,
|
||||
world_size: u32,
|
||||
comm: Arc<Comm>,
|
||||
quant: Option<GgmlDType>,
|
||||
) -> Result<Self> {
|
||||
let gate = Self::check_and_gate(cfg, vb, world_size)?;
|
||||
let (experts, shared_expert, shared_expert_gate) =
|
||||
Self::load_experts_and_shared(cfg, vb, rank, world_size, quant)?;
|
||||
Ok(Self {
|
||||
gate,
|
||||
experts,
|
||||
shared_expert,
|
||||
shared_expert_gate,
|
||||
num_experts_per_tok: cfg.num_experts_per_tok,
|
||||
norm_topk_prob: cfg.norm_topk_prob,
|
||||
all_reduce: super::all_reduce::AllReduce::new(comm),
|
||||
needs_reduce: world_size > 1,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "cuda"))]
|
||||
pub fn load(
|
||||
cfg: &TextConfig,
|
||||
vb: &ShardedVarBuilder,
|
||||
rank: u32,
|
||||
world_size: u32,
|
||||
quant: Option<GgmlDType>,
|
||||
) -> Result<Self> {
|
||||
let gate = Self::check_and_gate(cfg, vb, world_size)?;
|
||||
let (experts, shared_expert, shared_expert_gate) =
|
||||
Self::load_experts_and_shared(cfg, vb, rank, world_size, quant)?;
|
||||
Ok(Self {
|
||||
gate,
|
||||
experts,
|
||||
shared_expert,
|
||||
shared_expert_gate,
|
||||
num_experts_per_tok: cfg.num_experts_per_tok,
|
||||
norm_topk_prob: cfg.norm_topk_prob,
|
||||
needs_reduce: world_size > 1,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl TpQwen3_5MoeBlock {
|
||||
/// Correctness-first dispatch: shared `route_scatter` host routing
|
||||
/// with per-expert GEMMs. Returns the rank's PARTIAL routed
|
||||
/// output, `(tokens, hidden)`, in the input dtype.
|
||||
fn forward_scatter(
|
||||
&self,
|
||||
experts: &[TpExpert],
|
||||
xs_flat: &Tensor,
|
||||
) -> candle_core::Result<Tensor> {
|
||||
let (tokens_for, weights_for) = crate::harness::arch::qwen3_5::moe::route_scatter(
|
||||
&self.gate,
|
||||
xs_flat,
|
||||
experts.len(),
|
||||
self.num_experts_per_tok,
|
||||
self.norm_topk_prob,
|
||||
)?;
|
||||
|
||||
let mut ys = xs_flat.zeros_like()?;
|
||||
for (e, expert) in experts.iter().enumerate() {
|
||||
if tokens_for[e].is_empty() {
|
||||
continue;
|
||||
}
|
||||
let rows = Tensor::new(tokens_for[e].as_slice(), xs_flat.device())?;
|
||||
let picked = xs_flat.index_select(&rows, 0)?;
|
||||
let out = expert.forward_partial(&picked)?;
|
||||
let w = Tensor::new(weights_for[e].as_slice(), xs_flat.device())?
|
||||
.to_dtype(out.dtype())?
|
||||
.reshape(((), 1))?;
|
||||
ys = ys.index_add(&rows, &out.broadcast_mul(&w)?, 0)?;
|
||||
}
|
||||
Ok(ys)
|
||||
}
|
||||
|
||||
/// Fused grouped-GEMM dispatch (#92 slice 4): routing, index sort,
|
||||
/// and all expert GEMMs stay on-device — the port of
|
||||
/// candle-transformers' `FusedMoeGGUF::forward` onto per-rank
|
||||
/// expert stacks. Returns the rank's PARTIAL routed output,
|
||||
/// `(tokens, hidden)`, in the input dtype.
|
||||
///
|
||||
/// Kernel contract (candle-nn `moe_gemm_gguf`): decode wants an
|
||||
/// F32 input and always emits F32; the `dtype` argument only
|
||||
/// selects the f16/bf16 conversion used by the prefill kernel.
|
||||
/// gate/up run with `topk_weights: None` → `tokens×topk` output
|
||||
/// rows; the down GEMM folds the routing weights in-kernel and the
|
||||
/// final `(tokens, topk, hidden)` view sums over `topk`.
|
||||
#[cfg(feature = "cuda")]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn forward_fused(
|
||||
&self,
|
||||
gate_experts: &candle_core::quantized::QTensor,
|
||||
up_experts: &candle_core::quantized::QTensor,
|
||||
down_experts: &candle_core::quantized::QTensor,
|
||||
num_experts: usize,
|
||||
xs_flat: &Tensor,
|
||||
is_prefill: bool,
|
||||
) -> candle_core::Result<Tensor> {
|
||||
use candle_core::D;
|
||||
|
||||
let original_dtype = xs_flat.dtype();
|
||||
let n_tokens = xs_flat.dim(0)?;
|
||||
let xs_f32 = if original_dtype == DType::F32 {
|
||||
xs_flat.clone()
|
||||
} else {
|
||||
xs_flat.to_dtype(DType::F32)?
|
||||
};
|
||||
|
||||
// Replicated routing, all on-device (contrast route_scatter's
|
||||
// host round-trip): softmax over all experts → top-k → renorm.
|
||||
// The router runs in the ORIGINAL activation dtype — its
|
||||
// replicated weight is bf16, and a bf16×f32 matmul is a dtype
|
||||
// error (the first live fused request poisoned on exactly
|
||||
// this). Only the softmax and everything downstream are f32,
|
||||
// matching route_scatter.
|
||||
let router_logits = self.gate.forward(xs_flat)?;
|
||||
let probs = candle_nn::ops::softmax_last_dim(&router_logits.to_dtype(DType::F32)?)?;
|
||||
let topk_ids = probs
|
||||
.arg_sort_last_dim(false)?
|
||||
.narrow(D::Minus1, 0, self.num_experts_per_tok)?
|
||||
.contiguous()?;
|
||||
let mut topk_weights = probs.gather(&topk_ids, D::Minus1)?;
|
||||
if self.norm_topk_prob {
|
||||
topk_weights = topk_weights.broadcast_div(&topk_weights.sum_keepdim(D::Minus1)?)?;
|
||||
}
|
||||
let (expert_ids, sorted_token_ids) = topk_ids.flatten_all()?.sort_last_dim(true)?;
|
||||
let _ = num_experts;
|
||||
|
||||
let gate = candle_nn::moe::moe_gemm_gguf(
|
||||
&xs_f32,
|
||||
gate_experts,
|
||||
&None,
|
||||
&sorted_token_ids,
|
||||
&expert_ids,
|
||||
self.num_experts_per_tok,
|
||||
is_prefill,
|
||||
DType::BF16,
|
||||
)?;
|
||||
let up = candle_nn::moe::moe_gemm_gguf(
|
||||
&xs_f32,
|
||||
up_experts,
|
||||
&None,
|
||||
&sorted_token_ids,
|
||||
&expert_ids,
|
||||
self.num_experts_per_tok,
|
||||
is_prefill,
|
||||
DType::BF16,
|
||||
)?;
|
||||
let down_inputs = (up * candle_nn::ops::silu(&gate)?)?;
|
||||
let ys = candle_nn::moe::moe_gemm_gguf(
|
||||
&down_inputs,
|
||||
down_experts,
|
||||
&Some(topk_weights),
|
||||
&sorted_token_ids,
|
||||
&expert_ids,
|
||||
self.num_experts_per_tok,
|
||||
is_prefill,
|
||||
DType::BF16,
|
||||
)?;
|
||||
|
||||
let hidden = xs_flat.dim(1)?;
|
||||
let ys = ys.reshape((n_tokens, (), hidden))?.sum(D::Minus2)?;
|
||||
if ys.dtype() == original_dtype {
|
||||
Ok(ys)
|
||||
} else {
|
||||
ys.to_dtype(original_dtype)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Module for TpQwen3_5MoeBlock {
|
||||
/// Route + expert dispatch (scatter or fused per the store), then
|
||||
/// the shared expert's sigmoid-gated partial, then ONE AllReduce
|
||||
/// recovering the full activation. The shared expert's per-token
|
||||
/// mix is a replicated scalar, so scaling the partial slice
|
||||
/// commutes with the reduce.
|
||||
fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||||
let (b, l, hidden) = xs.dims3()?;
|
||||
let xs_flat = xs.reshape(((), hidden))?;
|
||||
|
||||
let mut ys = match &self.experts {
|
||||
TpExpertStore::Scatter(experts) => self.forward_scatter(experts, &xs_flat)?,
|
||||
#[cfg(feature = "cuda")]
|
||||
TpExpertStore::Fused {
|
||||
gate_experts,
|
||||
up_experts,
|
||||
down_experts,
|
||||
num_experts,
|
||||
} => self.forward_fused(
|
||||
gate_experts,
|
||||
up_experts,
|
||||
down_experts,
|
||||
*num_experts,
|
||||
&xs_flat,
|
||||
l > 1,
|
||||
)?,
|
||||
};
|
||||
|
||||
if let (Some(shared), Some(gate)) = (&self.shared_expert, &self.shared_expert_gate) {
|
||||
let mix = candle_nn::ops::sigmoid(&gate.forward(&xs_flat)?)?;
|
||||
let shared_out = shared.forward_partial(&xs_flat)?.broadcast_mul(&mix)?;
|
||||
ys = (ys + shared_out)?;
|
||||
}
|
||||
|
||||
#[cfg(feature = "cuda")]
|
||||
if self.needs_reduce {
|
||||
ys = ys.apply_op1_no_bwd(&self.all_reduce)?;
|
||||
}
|
||||
let _ = self.needs_reduce;
|
||||
ys.reshape((b, l, hidden))
|
||||
}
|
||||
}
|
||||
|
||||
/// The FFN slot: dense SwiGLU (Qwen3.6) or the sharded MoE block
|
||||
/// (qwen3_next 80B-A3B, #92) — mirrors the single-GPU `MlpKind`.
|
||||
enum TpMlpKind {
|
||||
Dense(TpQwen3_5MLP),
|
||||
Moe(TpQwen3_5MoeBlock),
|
||||
}
|
||||
|
||||
impl Module for TpMlpKind {
|
||||
fn forward(&self, x: &Tensor) -> candle_core::Result<Tensor> {
|
||||
match self {
|
||||
TpMlpKind::Dense(mlp) => mlp.forward(x),
|
||||
TpMlpKind::Moe(moe) => moe.forward(x),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── decoder layer ──────────────────────────────────────────────────
|
||||
|
||||
enum TpAttentionKind {
|
||||
@@ -745,7 +1300,7 @@ enum TpAttentionKind {
|
||||
pub struct TpQwen3_5DecoderLayer {
|
||||
input_layernorm: Qwen3_5RmsNorm,
|
||||
post_attention_layernorm: Qwen3_5RmsNorm,
|
||||
mlp: TpQwen3_5MLP,
|
||||
mlp: TpMlpKind,
|
||||
attention: TpAttentionKind,
|
||||
}
|
||||
|
||||
@@ -789,7 +1344,25 @@ impl TpQwen3_5DecoderLayer {
|
||||
)?),
|
||||
other => bail!("unknown layer_type '{other}' for layer {layer_idx}"),
|
||||
};
|
||||
let mlp = TpQwen3_5MLP::load(cfg, &vb.pp("mlp"), rank, world_size, comm, quant)?;
|
||||
let mlp = if cfg.layer_uses_moe(layer_idx) {
|
||||
TpMlpKind::Moe(TpQwen3_5MoeBlock::load(
|
||||
cfg,
|
||||
&vb.pp("mlp"),
|
||||
rank,
|
||||
world_size,
|
||||
comm,
|
||||
quant,
|
||||
)?)
|
||||
} else {
|
||||
TpMlpKind::Dense(TpQwen3_5MLP::load(
|
||||
cfg,
|
||||
&vb.pp("mlp"),
|
||||
rank,
|
||||
world_size,
|
||||
comm,
|
||||
quant,
|
||||
)?)
|
||||
};
|
||||
let input_layernorm =
|
||||
Qwen3_5RmsNorm::load(&vb.pp("input_layernorm"), cfg.hidden_size, cfg.rms_norm_eps)?;
|
||||
let post_attention_layernorm = Qwen3_5RmsNorm::load(
|
||||
@@ -841,7 +1414,23 @@ impl TpQwen3_5DecoderLayer {
|
||||
)?),
|
||||
other => bail!("unknown layer_type '{other}' for layer {layer_idx}"),
|
||||
};
|
||||
let mlp = TpQwen3_5MLP::load(cfg, &vb.pp("mlp"), rank, world_size, quant)?;
|
||||
let mlp = if cfg.layer_uses_moe(layer_idx) {
|
||||
TpMlpKind::Moe(TpQwen3_5MoeBlock::load(
|
||||
cfg,
|
||||
&vb.pp("mlp"),
|
||||
rank,
|
||||
world_size,
|
||||
quant,
|
||||
)?)
|
||||
} else {
|
||||
TpMlpKind::Dense(TpQwen3_5MLP::load(
|
||||
cfg,
|
||||
&vb.pp("mlp"),
|
||||
rank,
|
||||
world_size,
|
||||
quant,
|
||||
)?)
|
||||
};
|
||||
let input_layernorm =
|
||||
Qwen3_5RmsNorm::load(&vb.pp("input_layernorm"), cfg.hidden_size, cfg.rms_norm_eps)?;
|
||||
let post_attention_layernorm = Qwen3_5RmsNorm::load(
|
||||
@@ -946,10 +1535,11 @@ impl TpQwen3_5Model {
|
||||
world_size: u32,
|
||||
comm: Arc<Comm>,
|
||||
quant: Option<GgmlDType>,
|
||||
text_prefix: &str,
|
||||
) -> Result<Self> {
|
||||
let dtype = vb.dtype();
|
||||
let device = vb.device().clone();
|
||||
let text_vb = vb.pp("model.language_model");
|
||||
let text_vb = vb.pp(text_prefix);
|
||||
|
||||
let embed_weight = load_replicated(
|
||||
&text_vb.pp("embed_tokens"),
|
||||
@@ -1029,10 +1619,11 @@ impl TpQwen3_5Model {
|
||||
rank: u32,
|
||||
world_size: u32,
|
||||
quant: Option<GgmlDType>,
|
||||
text_prefix: &str,
|
||||
) -> Result<Self> {
|
||||
let dtype = vb.dtype();
|
||||
let device = vb.device().clone();
|
||||
let text_vb = vb.pp("model.language_model");
|
||||
let text_vb = vb.pp(text_prefix);
|
||||
|
||||
let embed_weight = load_replicated(
|
||||
&text_vb.pp("embed_tokens"),
|
||||
@@ -1285,7 +1876,8 @@ impl TpQwen3_5ForCausalLM {
|
||||
quant: Option<GgmlDType>,
|
||||
) -> Result<Self> {
|
||||
let cfg = &config.text_config;
|
||||
let base = TpQwen3_5Model::load(cfg, vb, mmap, rank, world_size, comm, quant)?;
|
||||
let text_prefix = crate::harness::arch::qwen3_5::text_weight_prefix(&config.model_type);
|
||||
let base = TpQwen3_5Model::load(cfg, vb, mmap, rank, world_size, comm, quant, text_prefix)?;
|
||||
let lm_head = build_lm_head(cfg, vb, &base, quant)?;
|
||||
let vision = load_replicated_vision_tower(&config, vb)?;
|
||||
let image_token_id = config.image_token_id;
|
||||
@@ -1309,7 +1901,8 @@ impl TpQwen3_5ForCausalLM {
|
||||
quant: Option<GgmlDType>,
|
||||
) -> Result<Self> {
|
||||
let cfg = &config.text_config;
|
||||
let base = TpQwen3_5Model::load(cfg, vb, mmap, rank, world_size, quant)?;
|
||||
let text_prefix = crate::harness::arch::qwen3_5::text_weight_prefix(&config.model_type);
|
||||
let base = TpQwen3_5Model::load(cfg, vb, mmap, rank, world_size, quant, text_prefix)?;
|
||||
let lm_head = build_lm_head(cfg, vb, &base, quant)?;
|
||||
let vision = load_replicated_vision_tower(&config, vb)?;
|
||||
let image_token_id = config.image_token_id;
|
||||
@@ -1694,3 +2287,208 @@ fn log_construction_complete(
|
||||
"Qwen3-Next model construction complete"
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
#[cfg(not(feature = "cuda"))]
|
||||
use crate::harness::arch::qwen3_5::moe::Qwen3_5MoeBlock;
|
||||
#[cfg(not(feature = "cuda"))]
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Write a tiny MoE-block checkpoint (router + experts + shared
|
||||
/// expert) and return a ShardedVarBuilder over it plus the config.
|
||||
/// Non-cuda only: the cuda `TpQwen3_5MoeBlock::load` takes an NCCL
|
||||
/// `Comm`, which tests cannot construct — the CPU Test job covers
|
||||
/// this; the CUDA job type-checks the cuda variants.
|
||||
#[cfg(not(feature = "cuda"))]
|
||||
fn tiny_moe_fixture(dir: &std::path::Path) -> (TextConfig, std::path::PathBuf) {
|
||||
let dev = Device::Cpu;
|
||||
let randn = |shape: &[usize]| Tensor::randn(0f32, 0.3f32, shape, &dev).unwrap();
|
||||
let (h, inter, n_exp) = (8usize, 4usize, 6usize);
|
||||
|
||||
let mut t: HashMap<String, Tensor> = HashMap::new();
|
||||
t.insert("mlp.gate.weight".into(), randn(&[n_exp, h]));
|
||||
for e in 0..n_exp {
|
||||
t.insert(
|
||||
format!("mlp.experts.{e}.gate_proj.weight"),
|
||||
randn(&[inter, h]),
|
||||
);
|
||||
t.insert(
|
||||
format!("mlp.experts.{e}.up_proj.weight"),
|
||||
randn(&[inter, h]),
|
||||
);
|
||||
t.insert(
|
||||
format!("mlp.experts.{e}.down_proj.weight"),
|
||||
randn(&[h, inter]),
|
||||
);
|
||||
}
|
||||
t.insert(
|
||||
"mlp.shared_expert.gate_proj.weight".into(),
|
||||
randn(&[inter, h]),
|
||||
);
|
||||
t.insert(
|
||||
"mlp.shared_expert.up_proj.weight".into(),
|
||||
randn(&[inter, h]),
|
||||
);
|
||||
t.insert(
|
||||
"mlp.shared_expert.down_proj.weight".into(),
|
||||
randn(&[h, inter]),
|
||||
);
|
||||
t.insert("mlp.shared_expert_gate.weight".into(), randn(&[1, h]));
|
||||
|
||||
let path = dir.join("moe.safetensors");
|
||||
candle_core::safetensors::save(&t, &path).expect("save moe safetensors");
|
||||
|
||||
// Minimal TextConfig via the flat-config parser: only the MoE
|
||||
// fields matter for the block loaders.
|
||||
let cfg = Config::from_config_json(
|
||||
r#"{
|
||||
"model_type": "qwen3_next",
|
||||
"vocab_size": 16, "hidden_size": 8, "intermediate_size": 16,
|
||||
"num_hidden_layers": 1, "num_attention_heads": 2,
|
||||
"num_key_value_heads": 1, "head_dim": 4,
|
||||
"max_position_embeddings": 32, "rms_norm_eps": 1e-6,
|
||||
"num_experts": 6, "num_experts_per_tok": 2,
|
||||
"moe_intermediate_size": 4,
|
||||
"shared_expert_intermediate_size": 4,
|
||||
"norm_topk_prob": true
|
||||
}"#,
|
||||
)
|
||||
.expect("parse tiny moe config")
|
||||
.text_config;
|
||||
(cfg, path)
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "cuda"))]
|
||||
fn vb_over(path: &std::path::Path) -> ShardedVarBuilder {
|
||||
// SAFETY: mmap of a file the test just wrote; nothing mutates it.
|
||||
unsafe {
|
||||
candle_nn::var_builder::ShardedSafeTensors::var_builder(
|
||||
std::slice::from_ref(&path.to_path_buf()),
|
||||
DType::F32,
|
||||
&Device::Cpu,
|
||||
)
|
||||
.expect("build ShardedVarBuilder")
|
||||
}
|
||||
}
|
||||
|
||||
/// world_size = 2 on CPU: the block-end AllReduce is elided, so
|
||||
/// each rank's forward returns its PARTIAL output. Summing the two
|
||||
/// ranks' partials must reproduce the single-GPU output — this
|
||||
/// pins the expert slicing (column gate/up, row down), the
|
||||
/// replicated routing, and the shared-expert partial scaling,
|
||||
/// i.e. everything the real AllReduce would combine.
|
||||
#[cfg(not(feature = "cuda"))]
|
||||
#[test]
|
||||
fn tp_moe_ws2_partials_sum_to_single_gpu_output() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let (cfg, path) = tiny_moe_fixture(dir.path());
|
||||
|
||||
let single =
|
||||
Qwen3_5MoeBlock::load(&cfg, &vb_over(&path).pp("mlp")).expect("single-GPU load");
|
||||
let rank0 = TpQwen3_5MoeBlock::load(&cfg, &vb_over(&path).pp("mlp"), 0, 2, None)
|
||||
.expect("TP rank 0 load");
|
||||
let rank1 = TpQwen3_5MoeBlock::load(&cfg, &vb_over(&path).pp("mlp"), 1, 2, None)
|
||||
.expect("TP rank 1 load");
|
||||
|
||||
let xs = Tensor::randn(0f32, 1.0f32, (1, 4, 8), &Device::Cpu).unwrap();
|
||||
let full: Vec<f32> = single
|
||||
.forward(&xs)
|
||||
.unwrap()
|
||||
.flatten_all()
|
||||
.unwrap()
|
||||
.to_vec1()
|
||||
.unwrap();
|
||||
let p0 = rank0.forward(&xs).unwrap();
|
||||
let p1 = rank1.forward(&xs).unwrap();
|
||||
let summed: Vec<f32> = (p0 + p1).unwrap().flatten_all().unwrap().to_vec1().unwrap();
|
||||
for (i, (x, y)) in full.iter().zip(&summed).enumerate() {
|
||||
assert!(
|
||||
(x - y).abs() < 1e-4,
|
||||
"dim {i}: single {x} vs summed partials {y}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-rank fused-qkvz de-interleave: sharding the fused tensor by
|
||||
/// whole k-head groups and then splitting with per-rank head counts
|
||||
/// must equal the corresponding row-ranges of the full split.
|
||||
#[test]
|
||||
fn per_rank_qkvz_split_matches_full_split_slices() {
|
||||
use crate::harness::arch::qwen3_5::linear_attn::{split_fused_ba, split_fused_qkvz};
|
||||
let dev = Device::Cpu;
|
||||
let (num_k, num_v, head_k, head_v, hidden) = (4usize, 8usize, 3usize, 2usize, 5usize);
|
||||
let r = num_v / num_k;
|
||||
let stride = 2 * head_k + 2 * head_v * r;
|
||||
let fused = Tensor::randn(0f32, 1.0f32, (num_k * stride, hidden), &dev).unwrap();
|
||||
let ba = Tensor::randn(0f32, 1.0f32, (2 * num_v, hidden), &dev).unwrap();
|
||||
|
||||
let (full_qkv, full_z) = split_fused_qkvz(&fused, num_k, num_v, head_k, head_v).unwrap();
|
||||
let (full_b, full_a) = split_fused_ba(&ba, num_k, num_v).unwrap();
|
||||
let key_dim = num_k * head_k;
|
||||
|
||||
let ws = 2usize;
|
||||
for rank in 0..ws {
|
||||
let (pk, pv) = (num_k / ws, num_v / ws);
|
||||
let group_rows = pk * stride;
|
||||
let rank_fused = fused.narrow(0, rank * group_rows, group_rows).unwrap();
|
||||
let (qkv, z) = split_fused_qkvz(&rank_fused, pk, pv, head_k, head_v).unwrap();
|
||||
|
||||
// Expected: the rank's row-ranges of the full [Q|K|V] and Z.
|
||||
let (prk, prv) = (pk * head_k, pv * head_v);
|
||||
let expect_qkv = Tensor::cat(
|
||||
&[
|
||||
full_qkv.narrow(0, rank * prk, prk).unwrap(),
|
||||
full_qkv.narrow(0, key_dim + rank * prk, prk).unwrap(),
|
||||
full_qkv.narrow(0, 2 * key_dim + rank * prv, prv).unwrap(),
|
||||
],
|
||||
0,
|
||||
)
|
||||
.unwrap();
|
||||
let expect_z = full_z.narrow(0, rank * prv, prv).unwrap();
|
||||
let d1: f32 = (qkv - expect_qkv)
|
||||
.unwrap()
|
||||
.abs()
|
||||
.unwrap()
|
||||
.max_all()
|
||||
.unwrap()
|
||||
.to_scalar()
|
||||
.unwrap();
|
||||
let d2: f32 = (z - expect_z)
|
||||
.unwrap()
|
||||
.abs()
|
||||
.unwrap()
|
||||
.max_all()
|
||||
.unwrap()
|
||||
.to_scalar()
|
||||
.unwrap();
|
||||
assert_eq!(d1, 0.0, "rank {rank} qkv slice mismatch");
|
||||
assert_eq!(d2, 0.0, "rank {rank} z slice mismatch");
|
||||
|
||||
// ba: rank's groups are 2r rows each.
|
||||
let rank_ba = ba.narrow(0, rank * pk * 2 * r, pk * 2 * r).unwrap();
|
||||
let (b, a) = split_fused_ba(&rank_ba, pk, pv).unwrap();
|
||||
let expect_b = full_b.narrow(0, rank * pv, pv).unwrap();
|
||||
let expect_a = full_a.narrow(0, rank * pv, pv).unwrap();
|
||||
let d3: f32 = (b - expect_b)
|
||||
.unwrap()
|
||||
.abs()
|
||||
.unwrap()
|
||||
.max_all()
|
||||
.unwrap()
|
||||
.to_scalar()
|
||||
.unwrap();
|
||||
let d4: f32 = (a - expect_a)
|
||||
.unwrap()
|
||||
.abs()
|
||||
.unwrap()
|
||||
.max_all()
|
||||
.unwrap()
|
||||
.to_scalar()
|
||||
.unwrap();
|
||||
assert_eq!(d3, 0.0, "rank {rank} b slice mismatch");
|
||||
assert_eq!(d4, 0.0, "rank {rank} a slice mismatch");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -405,8 +405,10 @@ impl WorkerState {
|
||||
}
|
||||
}
|
||||
}
|
||||
"qwen3_5" => {
|
||||
let cfg: qwen3_5_arch::Config = match serde_json::from_str(&config_json) {
|
||||
"qwen3_5" | "qwen3_next" => {
|
||||
// `from_config_json` normalises the flat qwen3_next
|
||||
// layout (#92) into the nested qwen3_5 shape.
|
||||
let cfg = match qwen3_5_arch::Config::from_config_json(&config_json) {
|
||||
Ok(c) => c,
|
||||
Err(e) => {
|
||||
return WorkerResponse::Error {
|
||||
@@ -437,7 +439,8 @@ impl WorkerState {
|
||||
return WorkerResponse::Error {
|
||||
kind: "unsupported_arch".into(),
|
||||
message: format!(
|
||||
"worker: unsupported model_type '{other}' (supported: qwen3, qwen3_5)"
|
||||
"worker: unsupported model_type '{other}' \
|
||||
(supported: qwen3, qwen3_5, qwen3_next)"
|
||||
),
|
||||
};
|
||||
}
|
||||
|
||||
@@ -84,9 +84,38 @@ pub enum InferenceEvent {
|
||||
/// `output_tokens_details.reasoning_tokens` (responses).
|
||||
/// Zero for non-reasoning models.
|
||||
reasoning_tokens: u32,
|
||||
/// Server-measured prefill/decode timing for the request, or
|
||||
/// `None` on paths that don't measure it (CPU fallback that
|
||||
/// doesn't instrument, tests). Streaming projectors surface
|
||||
/// this as a `helexa_timing` extension on the OpenAI `usage`
|
||||
/// object so the bench harness can compute true prefill vs
|
||||
/// decode tok/s instead of inferring both from client-side
|
||||
/// SSE arrival (#85).
|
||||
timing: Option<FinishTiming>,
|
||||
},
|
||||
}
|
||||
|
||||
/// Server-measured timing for one completed inference, attached to
|
||||
/// [`InferenceEvent::Finish`]. The whole point is to separate the two
|
||||
/// phases the client cannot tell apart from chunk-arrival timing:
|
||||
/// prefill (tokenize + prompt forward pass, ending at the first
|
||||
/// sampled token) and decode (every subsequent token through EOS /
|
||||
/// `max_tokens`).
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct FinishTiming {
|
||||
/// Wall-clock of the prefill phase in milliseconds: from the start
|
||||
/// of the prompt forward pass(es) to the first sampled token.
|
||||
pub prefill_ms: u32,
|
||||
/// Wall-clock of the decode phase in milliseconds: from the first
|
||||
/// sampled token to stream end.
|
||||
pub decode_ms: u32,
|
||||
/// Prompt tokens submitted to the prefill forward pass — the
|
||||
/// denominator for prefill tok/s. With prefix-KV-cache hits (#11)
|
||||
/// the elapsed `prefill_ms` drops while this stays the full prompt
|
||||
/// length, so a high implied rate is itself the cache-hit signal.
|
||||
pub prefill_tokens: u32,
|
||||
}
|
||||
|
||||
/// Why a stream stopped. Stays small on purpose — anything that
|
||||
/// doesn't map cleanly to one of these collapses to [`Stop`].
|
||||
///
|
||||
|
||||
@@ -22,6 +22,6 @@ pub mod openai_chat;
|
||||
pub mod openai_responses;
|
||||
|
||||
pub use event::{
|
||||
FinishReason, InferenceEvent, ReasoningTokenPair, ToolCallTokenPair,
|
||||
FinishReason, FinishTiming, InferenceEvent, ReasoningTokenPair, ToolCallTokenPair,
|
||||
detect_reasoning_token_pair, detect_tool_call_token_pair,
|
||||
};
|
||||
|
||||
@@ -26,11 +26,13 @@
|
||||
//! producer blocks on its own send. The bounded channels
|
||||
//! propagate without us writing any logic.
|
||||
|
||||
use cortex_core::openai::{ChatCompletionChunk, ChunkChoice, CompletionTokensDetails, Usage};
|
||||
use cortex_core::openai::{
|
||||
ChatCompletionChunk, ChunkChoice, CompletionTokensDetails, HelexaTiming, Usage,
|
||||
};
|
||||
use serde_json::json;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use super::event::{FinishReason, InferenceEvent, ReasoningTokenPair};
|
||||
use super::event::{FinishReason, FinishTiming, InferenceEvent, ReasoningTokenPair};
|
||||
|
||||
/// Output channel buffer size. Mirrors the input side's bound; one
|
||||
/// event maps to at most one chunk, so equal capacity keeps the
|
||||
@@ -193,12 +195,14 @@ pub fn project_chat_stream_with(
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
reasoning_tokens,
|
||||
timing,
|
||||
} => {
|
||||
// The finish_reason chunk, then an OpenAI-style
|
||||
// usage-only chunk (`choices: []`, `usage` populated).
|
||||
// Clients (opencode) read this to track context size;
|
||||
// cortex's Anthropic translator also picks `usage` up
|
||||
// for its `message_delta`.
|
||||
// for its `message_delta`. `timing` rides along as the
|
||||
// `helexa_timing` usage extension for the bench harness (#85).
|
||||
vec![
|
||||
final_chunk(&id, created, &model_id, reason),
|
||||
usage_chunk(
|
||||
@@ -208,6 +212,7 @@ pub fn project_chat_stream_with(
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
reasoning_tokens,
|
||||
timing,
|
||||
),
|
||||
]
|
||||
}
|
||||
@@ -334,6 +339,7 @@ fn usage_chunk(
|
||||
prompt_tokens: u32,
|
||||
completion_tokens: u32,
|
||||
reasoning_tokens: u32,
|
||||
timing: Option<FinishTiming>,
|
||||
) -> ChatCompletionChunk {
|
||||
ChatCompletionChunk {
|
||||
id: id.into(),
|
||||
@@ -351,6 +357,14 @@ fn usage_chunk(
|
||||
reasoning_tokens: reasoning_tokens as u64,
|
||||
}),
|
||||
prompt_tokens_details: None,
|
||||
// helexa extension (#85): server-measured prefill/decode
|
||||
// timing for the bench harness. Omitted on paths that don't
|
||||
// measure it so standard clients see unchanged JSON.
|
||||
helexa_timing: timing.map(|t| HelexaTiming {
|
||||
prefill_ms: t.prefill_ms as u64,
|
||||
decode_ms: t.decode_ms as u64,
|
||||
prefill_tokens: t.prefill_tokens as u64,
|
||||
}),
|
||||
}),
|
||||
extra: serde_json::Value::Object(Default::default()),
|
||||
}
|
||||
@@ -391,6 +405,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -413,6 +428,45 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn finish_timing_surfaces_on_usage_chunk() {
|
||||
// O1 (#85) wire contract: a Finish carrying FinishTiming must
|
||||
// surface as `usage.helexa_timing` on the trailing usage chunk,
|
||||
// which is what the bench harness reads to compute true prefill
|
||||
// vs decode tok/s. Absent timing must leave it None.
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
|
||||
let out_rx = project_chat_stream(rx, "id-1".into(), 1700, "m".into());
|
||||
|
||||
tx.send(InferenceEvent::Start).await.unwrap();
|
||||
tx.send(InferenceEvent::Finish {
|
||||
reason: FinishReason::Stop,
|
||||
prompt_tokens: 128,
|
||||
completion_tokens: 64,
|
||||
reasoning_tokens: 0,
|
||||
timing: Some(FinishTiming {
|
||||
prefill_ms: 200,
|
||||
decode_ms: 1500,
|
||||
prefill_tokens: 128,
|
||||
}),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
|
||||
let out = collect(out_rx).await;
|
||||
let usage = out
|
||||
.iter()
|
||||
.find_map(|c| c.usage.as_ref())
|
||||
.expect("usage chunk present");
|
||||
let timing = usage
|
||||
.helexa_timing
|
||||
.as_ref()
|
||||
.expect("helexa_timing populated when Finish carried timing");
|
||||
assert_eq!(timing.prefill_ms, 200);
|
||||
assert_eq!(timing.decode_ms, 1500);
|
||||
assert_eq!(timing.prefill_tokens, 128);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn empty_text_delta_is_dropped() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
|
||||
@@ -434,6 +488,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -496,6 +551,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -547,6 +603,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -592,6 +649,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -635,6 +693,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -662,6 +721,7 @@ mod tests {
|
||||
prompt_tokens: 42,
|
||||
completion_tokens: 5,
|
||||
reasoning_tokens: 2,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -695,6 +755,7 @@ mod tests {
|
||||
prompt_tokens: 10,
|
||||
completion_tokens: 7,
|
||||
reasoning_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -29,9 +29,9 @@
|
||||
|
||||
use cortex_core::openai::{ChatCompletionRequest, ChatMessage, MessageContent};
|
||||
use cortex_core::responses::{
|
||||
OutputTokensDetails, ResponsesContentPart, ResponsesInput, ResponsesInputItem,
|
||||
ResponsesMessageContent, ResponsesOutputContent, ResponsesOutputItem, ResponsesRequest,
|
||||
ResponsesResponse, ResponsesUsage, events,
|
||||
OutputTokensDetails, ResponsesContentPart, ResponsesInput, ResponsesInputElement,
|
||||
ResponsesInputItem, ResponsesMessageContent, ResponsesOutputContent, ResponsesOutputItem,
|
||||
ResponsesRequest, ResponsesResponse, ResponsesUsage, events,
|
||||
};
|
||||
use serde_json::{Value, json};
|
||||
use tokio::sync::mpsc;
|
||||
@@ -109,8 +109,26 @@ pub fn request_to_chat(req: ResponsesRequest) -> Result<ChatCompletionRequest, T
|
||||
});
|
||||
}
|
||||
ResponsesInput::Items(items) => {
|
||||
for item in items {
|
||||
if let Some(msg) = input_item_to_chat(item) {
|
||||
for element in items {
|
||||
let msg = match element {
|
||||
ResponsesInputElement::Typed(item) => input_item_to_chat(item),
|
||||
// Bare `{role, content}` (OpenAI EasyInputMessage —
|
||||
// what litellm/agent-zero emit). `content: null`
|
||||
// (e.g. an assistant turn carrying only tool calls)
|
||||
// collapses to an empty string so the turn is kept.
|
||||
ResponsesInputElement::EasyMessage { role, content } => Some(ChatMessage {
|
||||
role,
|
||||
content: content
|
||||
.map(message_content_to_chat)
|
||||
.unwrap_or_else(|| MessageContent::Text(String::new())),
|
||||
extra: Value::Object(Default::default()),
|
||||
}),
|
||||
// Forward-compat: an item shape we don't model.
|
||||
// Dropped rather than rejected (see
|
||||
// `ResponsesInputElement::Other`).
|
||||
ResponsesInputElement::Other(_) => None,
|
||||
};
|
||||
if let Some(msg) = msg {
|
||||
messages.push(msg);
|
||||
}
|
||||
}
|
||||
@@ -159,11 +177,18 @@ fn input_item_to_chat(item: ResponsesInputItem) -> Option<ChatMessage> {
|
||||
})
|
||||
}
|
||||
ResponsesInputItem::FunctionCallOutput { call_id, output } => {
|
||||
// `output` is either a plain string or an array of content
|
||||
// parts. Render a string as-is; anything else to compact
|
||||
// JSON so the tool result text reaches the model intact.
|
||||
let output_text = match output {
|
||||
Value::String(s) => s,
|
||||
other => other.to_string(),
|
||||
};
|
||||
let mut extra = serde_json::Map::new();
|
||||
extra.insert("tool_call_id".into(), Value::String(call_id));
|
||||
Some(ChatMessage {
|
||||
role: "tool".into(),
|
||||
content: MessageContent::Text(output),
|
||||
content: MessageContent::Text(output_text),
|
||||
extra: Value::Object(extra),
|
||||
})
|
||||
}
|
||||
@@ -192,7 +217,9 @@ fn message_content_to_chat(content: ResponsesMessageContent) -> MessageContent {
|
||||
.filter_map(|p| match p {
|
||||
ResponsesContentPart::InputText { text }
|
||||
| ResponsesContentPart::OutputText { text, .. } => Some(text),
|
||||
ResponsesContentPart::InputImage { .. } => None,
|
||||
ResponsesContentPart::InputImage { .. } | ResponsesContentPart::Unknown => {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n\n");
|
||||
@@ -211,6 +238,7 @@ fn message_content_to_chat(content: ResponsesMessageContent) -> MessageContent {
|
||||
"image_url": { "url": image_url },
|
||||
}));
|
||||
}
|
||||
ResponsesContentPart::Unknown => {}
|
||||
}
|
||||
}
|
||||
MessageContent::Parts(out)
|
||||
@@ -309,6 +337,9 @@ async fn run_projection(
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
reasoning_tokens,
|
||||
// Responses-side `helexa_timing` surfacing not wired yet;
|
||||
// the bench harness reads timing off the chat path (#85).
|
||||
timing: _,
|
||||
} => {
|
||||
finish = Some(reason);
|
||||
// Surface usage on the streaming `response.completed`
|
||||
@@ -535,6 +566,18 @@ mod tests {
|
||||
use super::*;
|
||||
use cortex_core::openai::MessageContent;
|
||||
|
||||
/// Wrap typed items as `input` elements. Most translator tests
|
||||
/// exercise the typed path; the bare easy-message and unknown-item
|
||||
/// paths have dedicated tests below.
|
||||
fn typed_items(items: Vec<ResponsesInputItem>) -> ResponsesInput {
|
||||
ResponsesInput::Items(
|
||||
items
|
||||
.into_iter()
|
||||
.map(ResponsesInputElement::Typed)
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
fn meta() -> ResponseMeta {
|
||||
ResponseMeta {
|
||||
response_id: "resp_1".into(),
|
||||
@@ -614,7 +657,7 @@ mod tests {
|
||||
fn translates_input_items_to_chat_messages() {
|
||||
let req = ResponsesRequest {
|
||||
model: "m".into(),
|
||||
input: ResponsesInput::Items(vec![
|
||||
input: typed_items(vec![
|
||||
ResponsesInputItem::Message {
|
||||
role: "user".into(),
|
||||
content: ResponsesMessageContent::Text("first".into()),
|
||||
@@ -646,7 +689,7 @@ mod tests {
|
||||
fn image_input_translates_to_chat_parts_array() {
|
||||
let req = ResponsesRequest {
|
||||
model: "m".into(),
|
||||
input: ResponsesInput::Items(vec![ResponsesInputItem::Message {
|
||||
input: typed_items(vec![ResponsesInputItem::Message {
|
||||
role: "user".into(),
|
||||
content: ResponsesMessageContent::Parts(vec![
|
||||
ResponsesContentPart::InputText {
|
||||
@@ -687,7 +730,7 @@ mod tests {
|
||||
// it's dropped — but it must not break translation.
|
||||
let req = ResponsesRequest {
|
||||
model: "m".into(),
|
||||
input: ResponsesInput::Items(vec![ResponsesInputItem::Message {
|
||||
input: typed_items(vec![ResponsesInputItem::Message {
|
||||
role: "user".into(),
|
||||
content: ResponsesMessageContent::Parts(vec![
|
||||
ResponsesContentPart::InputText {
|
||||
@@ -729,7 +772,7 @@ mod tests {
|
||||
fn text_only_parts_collapse_to_string() {
|
||||
let req = ResponsesRequest {
|
||||
model: "m".into(),
|
||||
input: ResponsesInput::Items(vec![ResponsesInputItem::Message {
|
||||
input: typed_items(vec![ResponsesInputItem::Message {
|
||||
role: "user".into(),
|
||||
content: ResponsesMessageContent::Parts(vec![
|
||||
ResponsesContentPart::InputText {
|
||||
@@ -759,7 +802,7 @@ mod tests {
|
||||
fn reasoning_items_are_silently_dropped() {
|
||||
let req = ResponsesRequest {
|
||||
model: "m".into(),
|
||||
input: ResponsesInput::Items(vec![
|
||||
input: typed_items(vec![
|
||||
ResponsesInputItem::Reasoning { content: vec![] },
|
||||
ResponsesInputItem::Message {
|
||||
role: "user".into(),
|
||||
@@ -779,6 +822,74 @@ mod tests {
|
||||
assert_eq!(chat.messages[0].role, "user");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bare_easy_messages_translate_like_typed_messages() {
|
||||
// The agent-zero / litellm shape: bare `{role, content}` items
|
||||
// with no `type`. Deserialize from raw JSON (not hand-built)
|
||||
// so this exercises the real parse path end to end.
|
||||
let raw = r#"{
|
||||
"model": "Qwen/Qwen3.6-27B",
|
||||
"store": true,
|
||||
"input": [
|
||||
{"role": "system", "content": "be terse"},
|
||||
{"role": "assistant", "content": "{\"tool_name\":\"response\"}"},
|
||||
{"role": "user", "content": "alpha"}
|
||||
]
|
||||
}"#;
|
||||
let req: ResponsesRequest = serde_json::from_str(raw).unwrap();
|
||||
let chat = request_to_chat(req).unwrap();
|
||||
let roles: Vec<&str> = chat.messages.iter().map(|m| m.role.as_str()).collect();
|
||||
assert_eq!(roles, vec!["system", "assistant", "user"]);
|
||||
assert!(matches!(
|
||||
&chat.messages[2].content,
|
||||
MessageContent::Text(t) if t == "alpha"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn null_content_and_unknown_items_survive_translation() {
|
||||
// An assistant turn with `content: null` is kept (empty text);
|
||||
// an unmodeled item type is dropped, not rejected.
|
||||
let raw = r#"{
|
||||
"model": "m",
|
||||
"input": [
|
||||
{"role": "assistant", "content": null},
|
||||
{"type": "item_reference", "id": "x"},
|
||||
{"role": "user", "content": "go"}
|
||||
]
|
||||
}"#;
|
||||
let req: ResponsesRequest = serde_json::from_str(raw).unwrap();
|
||||
let chat = request_to_chat(req).unwrap();
|
||||
// assistant(null) kept, item_reference dropped, user kept.
|
||||
let roles: Vec<&str> = chat.messages.iter().map(|m| m.role.as_str()).collect();
|
||||
assert_eq!(roles, vec!["assistant", "user"]);
|
||||
assert!(matches!(
|
||||
&chat.messages[0].content,
|
||||
MessageContent::Text(t) if t.is_empty()
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn function_call_output_array_renders_to_text() {
|
||||
// OpenAI allows `function_call_output.output` to be an array of
|
||||
// content parts; the tool result must reach the model as text.
|
||||
let raw = r#"{
|
||||
"model": "m",
|
||||
"input": [
|
||||
{"type": "function_call_output", "call_id": "c1",
|
||||
"output": [{"type": "output_text", "text": "42"}]}
|
||||
]
|
||||
}"#;
|
||||
let req: ResponsesRequest = serde_json::from_str(raw).unwrap();
|
||||
let chat = request_to_chat(req).unwrap();
|
||||
assert_eq!(chat.messages.len(), 1);
|
||||
assert_eq!(chat.messages[0].role, "tool");
|
||||
match &chat.messages[0].content {
|
||||
MessageContent::Text(t) => assert!(t.contains("42"), "got {t:?}"),
|
||||
other => panic!("expected text, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
// ── streaming projector ─────────────────────────────────────────
|
||||
|
||||
async fn collect(mut rx: mpsc::Receiver<ResponseStreamFrame>) -> Vec<ResponseStreamFrame> {
|
||||
@@ -806,6 +917,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -856,6 +968,7 @@ mod tests {
|
||||
prompt_tokens: 30,
|
||||
completion_tokens: 12,
|
||||
reasoning_tokens: 4,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -886,6 +999,7 @@ mod tests {
|
||||
prompt_tokens: 8,
|
||||
completion_tokens: 3,
|
||||
reasoning_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -910,6 +1024,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -956,6 +1071,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
56
crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/config.json
vendored
Normal file
56
crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/config.json
vendored
Normal file
@@ -0,0 +1,56 @@
|
||||
{
|
||||
"architectures": [
|
||||
"Qwen3NextForCausalLM"
|
||||
],
|
||||
"attention_bias": false,
|
||||
"attention_dropout": 0.0,
|
||||
"bos_token_id": null,
|
||||
"decoder_sparse_step": 1,
|
||||
"dtype": "float32",
|
||||
"eos_token_id": null,
|
||||
"head_dim": 32,
|
||||
"hidden_act": "silu",
|
||||
"hidden_size": 64,
|
||||
"initializer_range": 0.02,
|
||||
"intermediate_size": 128,
|
||||
"layer_types": [
|
||||
"linear_attention",
|
||||
"linear_attention",
|
||||
"linear_attention",
|
||||
"full_attention",
|
||||
"linear_attention",
|
||||
"linear_attention",
|
||||
"linear_attention",
|
||||
"full_attention"
|
||||
],
|
||||
"linear_conv_kernel_dim": 4,
|
||||
"linear_key_head_dim": 16,
|
||||
"linear_num_key_heads": 2,
|
||||
"linear_num_value_heads": 4,
|
||||
"linear_value_head_dim": 16,
|
||||
"max_position_embeddings": 512,
|
||||
"mlp_only_layers": [],
|
||||
"model_type": "qwen3_next",
|
||||
"moe_intermediate_size": 32,
|
||||
"norm_topk_prob": true,
|
||||
"num_attention_heads": 4,
|
||||
"num_experts": 16,
|
||||
"num_experts_per_tok": 4,
|
||||
"num_hidden_layers": 8,
|
||||
"num_key_value_heads": 2,
|
||||
"output_router_logits": false,
|
||||
"pad_token_id": null,
|
||||
"partial_rotary_factor": 0.25,
|
||||
"rms_norm_eps": 1e-06,
|
||||
"rope_parameters": {
|
||||
"partial_rotary_factor": 0.25,
|
||||
"rope_theta": 10000000,
|
||||
"rope_type": "default"
|
||||
},
|
||||
"router_aux_loss_coef": 0.001,
|
||||
"shared_expert_intermediate_size": 32,
|
||||
"tie_word_embeddings": false,
|
||||
"transformers_version": "5.9.0",
|
||||
"use_cache": true,
|
||||
"vocab_size": 512
|
||||
}
|
||||
7
crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/generation_config.json
vendored
Normal file
7
crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/generation_config.json
vendored
Normal file
@@ -0,0 +1,7 @@
|
||||
{
|
||||
"_from_model_config": true,
|
||||
"output_attentions": false,
|
||||
"output_hidden_states": false,
|
||||
"transformers_version": "5.9.0",
|
||||
"use_cache": true
|
||||
}
|
||||
BIN
crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/logits.f32
vendored
Normal file
BIN
crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/logits.f32
vendored
Normal file
Binary file not shown.
114
crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/manifest.json
vendored
Normal file
114
crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/manifest.json
vendored
Normal file
@@ -0,0 +1,114 @@
|
||||
{
|
||||
"case": "qwen3_next-tiny",
|
||||
"seed": 92,
|
||||
"token_ids": [
|
||||
3,
|
||||
10,
|
||||
17,
|
||||
24,
|
||||
31,
|
||||
38,
|
||||
45,
|
||||
52,
|
||||
59,
|
||||
66,
|
||||
73,
|
||||
80,
|
||||
87,
|
||||
94,
|
||||
101,
|
||||
108,
|
||||
115,
|
||||
122,
|
||||
129,
|
||||
136,
|
||||
143,
|
||||
150,
|
||||
157,
|
||||
164,
|
||||
171,
|
||||
178,
|
||||
185,
|
||||
192,
|
||||
199,
|
||||
206,
|
||||
213,
|
||||
220,
|
||||
227,
|
||||
234,
|
||||
241,
|
||||
248,
|
||||
255,
|
||||
262,
|
||||
269,
|
||||
276,
|
||||
283,
|
||||
290,
|
||||
297,
|
||||
304,
|
||||
311,
|
||||
318,
|
||||
325,
|
||||
332,
|
||||
339,
|
||||
346,
|
||||
353,
|
||||
360,
|
||||
367,
|
||||
374,
|
||||
381,
|
||||
388,
|
||||
395,
|
||||
402,
|
||||
409,
|
||||
416,
|
||||
423,
|
||||
430,
|
||||
437,
|
||||
444,
|
||||
451,
|
||||
458,
|
||||
465,
|
||||
472,
|
||||
479,
|
||||
486,
|
||||
493,
|
||||
500,
|
||||
507,
|
||||
2,
|
||||
9,
|
||||
16,
|
||||
23,
|
||||
30,
|
||||
37,
|
||||
44,
|
||||
51,
|
||||
58,
|
||||
65,
|
||||
72,
|
||||
79,
|
||||
86,
|
||||
93,
|
||||
100,
|
||||
107,
|
||||
114,
|
||||
121,
|
||||
128,
|
||||
135,
|
||||
142,
|
||||
149,
|
||||
156
|
||||
],
|
||||
"files": {
|
||||
"logits": {
|
||||
"file": "logits.f32",
|
||||
"shape": [
|
||||
512
|
||||
]
|
||||
}
|
||||
},
|
||||
"versions": {
|
||||
"transformers": "5.9.0",
|
||||
"torch": "2.9.1+cu128"
|
||||
}
|
||||
}
|
||||
BIN
crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/model.safetensors
vendored
Normal file
BIN
crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/model.safetensors
vendored
Normal file
Binary file not shown.
122
crates/neuron/tests/qwen3_next_parity.rs
Normal file
122
crates/neuron/tests/qwen3_next_parity.rs
Normal file
@@ -0,0 +1,122 @@
|
||||
//! Numerical parity for the qwen3_next path (#92) against the HF
|
||||
//! transformers reference, via the tiny self-contained fixture
|
||||
//! generated by `script/dump_qwen3_next_tiny.py`.
|
||||
//!
|
||||
//! The fixture directory carries the WHOLE checkpoint (tiny
|
||||
//! random-weight `Qwen3NextForCausalLM`: config.json +
|
||||
//! model.safetensors, a few hundred KB) plus the reference
|
||||
//! final-position logits, so this test needs no snapshot, no env var,
|
||||
//! and runs in CI. It pins: flat-config normalisation, the `model.*`
|
||||
//! weight prefix, the fused `in_proj_qkvz`/`in_proj_ba` de-interleave,
|
||||
//! hybrid full/linear layer interleaving, and the MoE block (routing,
|
||||
//! per-expert SwiGLU, shared expert + sigmoid gate).
|
||||
//!
|
||||
//! Self-skips (with a loud eprintln) while the fixture has not been
|
||||
//! generated yet — regeneration instructions in the script docstring.
|
||||
|
||||
use candle_core::{DType, Device, Tensor};
|
||||
use serde::Deserialize;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Manifest {
|
||||
token_ids: Vec<u32>,
|
||||
files: std::collections::HashMap<String, FileEntry>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct FileEntry {
|
||||
file: String,
|
||||
shape: Vec<usize>,
|
||||
}
|
||||
|
||||
fn fixture_dir() -> PathBuf {
|
||||
Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/numerical/qwen3_next-tiny")
|
||||
}
|
||||
|
||||
fn read_f32(path: &Path) -> Vec<f32> {
|
||||
let bytes = std::fs::read(path).unwrap_or_else(|e| panic!("read {path:?}: {e}"));
|
||||
bytes
|
||||
.chunks_exact(4)
|
||||
.map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tiny_qwen3_next_logits_match_hf_reference() {
|
||||
let dir = fixture_dir();
|
||||
let manifest_path = dir.join("manifest.json");
|
||||
if !manifest_path.exists() {
|
||||
eprintln!(
|
||||
"SKIP qwen3_next parity: fixture not generated yet — run \
|
||||
script/dump_qwen3_next_tiny.py --out {} on a host with \
|
||||
torch + transformers>=4.57",
|
||||
dir.display()
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
let manifest: Manifest =
|
||||
serde_json::from_str(&std::fs::read_to_string(&manifest_path).expect("read manifest"))
|
||||
.expect("parse manifest");
|
||||
let logits_entry = &manifest.files["logits"];
|
||||
let reference = read_f32(&dir.join(&logits_entry.file));
|
||||
assert_eq!(
|
||||
reference.len(),
|
||||
logits_entry.shape.iter().product::<usize>()
|
||||
);
|
||||
|
||||
let config_json = std::fs::read_to_string(dir.join("config.json")).expect("read config.json");
|
||||
let cfg = neuron::harness::arch::qwen3_5::Config::from_config_json(&config_json)
|
||||
.expect("normalise qwen3_next config");
|
||||
|
||||
let dev = Device::Cpu;
|
||||
let st_path = dir.join("model.safetensors");
|
||||
// SAFETY: mmap of committed fixture files; nothing mutates them.
|
||||
let vb = unsafe {
|
||||
candle_nn::var_builder::ShardedSafeTensors::var_builder(
|
||||
std::slice::from_ref(&st_path),
|
||||
DType::F32,
|
||||
&dev,
|
||||
)
|
||||
.expect("build ShardedVarBuilder over fixture checkpoint")
|
||||
};
|
||||
let mut model = neuron::harness::arch::qwen3_5::Qwen3_5ForCausalLM::new(cfg, vb)
|
||||
.expect("load tiny qwen3_next checkpoint through neuron");
|
||||
|
||||
let input = Tensor::new(manifest.token_ids.as_slice(), &dev)
|
||||
.unwrap()
|
||||
.unsqueeze(0)
|
||||
.unwrap();
|
||||
let logits = model.forward(&input, 0).expect("forward");
|
||||
let got: Vec<f32> = logits.flatten_all().unwrap().to_vec1().unwrap();
|
||||
assert_eq!(got.len(), reference.len());
|
||||
|
||||
// f32-vs-f32 through an 8-layer doll-house model: agreement should
|
||||
// be tight (the qwen3_5 text fixtures observe max_abs ≈ 0.000,
|
||||
// cosine ≈ 1.0). Thresholds sit far above rounding noise and far
|
||||
// below any real wiring bug (a swapped de-interleave region, a
|
||||
// topk-before-softmax, a missing shared-expert gate all blow past
|
||||
// them instantly).
|
||||
let max_abs = got
|
||||
.iter()
|
||||
.zip(&reference)
|
||||
.map(|(a, b)| (a - b).abs())
|
||||
.fold(0f32, f32::max);
|
||||
let dot: f64 = got
|
||||
.iter()
|
||||
.zip(&reference)
|
||||
.map(|(a, b)| (*a as f64) * (*b as f64))
|
||||
.sum();
|
||||
let na: f64 = got.iter().map(|a| (*a as f64).powi(2)).sum::<f64>().sqrt();
|
||||
let nb: f64 = reference
|
||||
.iter()
|
||||
.map(|b| (*b as f64).powi(2))
|
||||
.sum::<f64>()
|
||||
.sqrt();
|
||||
let cosine = dot / (na * nb);
|
||||
|
||||
eprintln!("qwen3_next parity: max_abs={max_abs:.6} cosine={cosine:.8}");
|
||||
assert!(max_abs < 1e-3, "max abs diff {max_abs} exceeds 1e-3");
|
||||
assert!(cosine > 0.9999, "cosine {cosine} below 0.9999");
|
||||
}
|
||||
6
data/helexa-upstream-firewalld.xml
Normal file
6
data/helexa-upstream-firewalld.xml
Normal file
@@ -0,0 +1,6 @@
|
||||
<?xml version="1.0" encoding="utf-8"?>
|
||||
<service>
|
||||
<short>helexa-upstream</short>
|
||||
<description>helexa-upstream — mesh account + budget authority API (/authz/v1 for cortex, /web/v1 for the frontend)</description>
|
||||
<port protocol="tcp" port="8090"/>
|
||||
</service>
|
||||
3
data/helexa-upstream-sysusers.conf
Normal file
3
data/helexa-upstream-sysusers.conf
Normal file
@@ -0,0 +1,3 @@
|
||||
g helexa-upstream - -
|
||||
u helexa-upstream - "helexa-upstream authority" /var/lib/helexa-upstream /sbin/nologin
|
||||
m helexa-upstream helexa-upstream
|
||||
22
data/helexa-upstream.service
Normal file
22
data/helexa-upstream.service
Normal file
@@ -0,0 +1,22 @@
|
||||
[Unit]
|
||||
Description=helexa-upstream — mesh account + budget authority (accounts, API keys, allocation ledger, top-up codes)
|
||||
After=network-online.target
|
||||
Wants=network-online.target
|
||||
|
||||
[Service]
|
||||
Type=simple
|
||||
ExecStart=/usr/bin/helexa-upstream serve --config /etc/helexa-upstream/helexa-upstream.toml
|
||||
# HTTP authority for cortex (/authz/v1) and the frontend (/web/v1); restart
|
||||
# unconditionally if it ever exits. It connects out to PostgreSQL (the
|
||||
# system of record) and runs schema migrations on startup.
|
||||
Restart=always
|
||||
RestartSec=10
|
||||
User=helexa-upstream
|
||||
Group=helexa-upstream
|
||||
# Service user home; no local state (PostgreSQL holds everything), but
|
||||
# StateDirectory gives the user a writable, correctly-owned home.
|
||||
StateDirectory=helexa-upstream
|
||||
StateDirectoryMode=0755
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
@@ -23,8 +23,32 @@ db_path = "/var/lib/helexa-bench/bench.sqlite"
|
||||
|
||||
[scenarios]
|
||||
# One chat-latency scenario is generated per size (chat:128, chat:4096).
|
||||
# For a context-length scaling curve (#88), add a ladder up to the model's
|
||||
# limit, e.g. [128, 1024, 4096, 16384, 65536, 131072]; then
|
||||
# `helexa-bench report --scaling` (or GET /api/scaling) shows prefill &
|
||||
# decode tok/s vs context with the decode-flatness verdict. Larger contexts
|
||||
# cost more per sample, so widen the ladder deliberately.
|
||||
prompt_sizes = [128, 4096]
|
||||
max_tokens = 256
|
||||
# Concurrency / agentic-load scenarios (#89): one concurrency:<n> scenario
|
||||
# per level, each firing N simultaneous streams to characterize the real
|
||||
# a0/hermes/opencode fan-out. Empty by default — enable deliberately, since
|
||||
# a burst puts genuine simultaneous load on the serving fleet.
|
||||
# concurrency_levels = [2, 4, 8]
|
||||
# concurrency_prompt_tokens = 512
|
||||
|
||||
# Capability probes (#91): each runs a fixed prompt and stores the full
|
||||
# output for quality scoring — the reasoning/planning axis the speed
|
||||
# scenarios miss. Opt-in (empty by default). After a sweep records them,
|
||||
# `helexa-bench report --capability` lists run ids + artifacts; score with
|
||||
# `helexa-bench score --id <n> --score <x>` (manual now; LLM-judge later).
|
||||
# [[scenarios.capability_probes]]
|
||||
# name = "rust-plan"
|
||||
# max_tokens = 2048
|
||||
# prompt = """
|
||||
# Write an implementation plan for adding rate limiting to an Axum service.
|
||||
# Honor existing conventions, call out trade-offs, and sequence the work.
|
||||
# """
|
||||
|
||||
# Read-only JSON API (consumed by the bench UI + programmatic access),
|
||||
# served alongside the sweep loop by `run` (or standalone via `serve`).
|
||||
|
||||
@@ -32,3 +32,22 @@ F0 scaffold. Theming + i18n (33 languages, usage-ordered selector), the
|
||||
`/mission` page, the chat workspace (Dexie + streaming), and the account
|
||||
dashboard land in subsequent phases — see
|
||||
`~/.claude/plans/we-need-to-plan-modular-graham.md`.
|
||||
|
||||
## Deploy (public beta)
|
||||
|
||||
Build the SPA and serve it from edge nginx on the **same origin** as the
|
||||
two backends — so the browser makes no cross-origin request (no CORS) and
|
||||
the user's API key rides as a first-party bearer.
|
||||
|
||||
```sh
|
||||
npm ci && npm run build # → dist/
|
||||
sudo cp -r dist/* /var/www/helexa.ai/
|
||||
sudo cp deploy/nginx.conf /etc/nginx/conf.d/helexa.ai.conf # adjust upstreams + TLS
|
||||
sudo nginx -t && sudo systemctl reload nginx
|
||||
```
|
||||
|
||||
`deploy/nginx.conf` routes `/` → SPA (history fallback), `/v1` + `/health`
|
||||
→ helexa-router, and `/api/` → helexa-upstream `/web/v1/`. Set
|
||||
`VITE_PUBLIC_BETA=true` at build time for the beta banner. There is **no
|
||||
server-side chat history**: conversations live only in the browser
|
||||
(IndexedDB).
|
||||
|
||||
60
helexa.ai/deploy/nginx.conf
Normal file
60
helexa.ai/deploy/nginx.conf
Normal file
@@ -0,0 +1,60 @@
|
||||
# helexa.ai — edge nginx for the public beta.
|
||||
#
|
||||
# Serves the built SPA (helexa.ai/dist) and reverse-proxies the two
|
||||
# backends on the SAME ORIGIN, so the browser never makes a cross-origin
|
||||
# request: no CORS, and the user's API key rides as a first-party bearer.
|
||||
#
|
||||
# / → static SPA (history fallback to index.html)
|
||||
# /v1, /health → helexa-router (OpenAI-compatible inference data-plane)
|
||||
# /api/ → helexa-upstream /web/v1/ (account control-plane)
|
||||
#
|
||||
# TLS is terminated here (certs omitted — wire up certbot / your CA). The
|
||||
# upstream hosts below are examples; point them at your router/upstream.
|
||||
|
||||
upstream helexa_router { server 127.0.0.1:8088; }
|
||||
upstream helexa_upstream { server 127.0.0.1:8090; }
|
||||
|
||||
server {
|
||||
listen 443 ssl http2;
|
||||
listen [::]:443 ssl http2;
|
||||
server_name helexa.ai;
|
||||
|
||||
# ssl_certificate /etc/letsencrypt/live/helexa.ai/fullchain.pem;
|
||||
# ssl_certificate_key /etc/letsencrypt/live/helexa.ai/privkey.pem;
|
||||
|
||||
root /var/www/helexa.ai;
|
||||
index index.html;
|
||||
|
||||
# Long-cache fingerprinted assets; never cache the HTML shell.
|
||||
location /assets/ {
|
||||
expires 1y;
|
||||
add_header Cache-Control "public, immutable";
|
||||
}
|
||||
|
||||
# Inference data-plane → router. Streaming (SSE): disable buffering so
|
||||
# tokens reach the browser as they arrive.
|
||||
location /v1/ {
|
||||
proxy_pass http://helexa_router;
|
||||
proxy_http_version 1.1;
|
||||
proxy_set_header Host $host;
|
||||
proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
|
||||
proxy_set_header Connection "";
|
||||
proxy_buffering off;
|
||||
proxy_read_timeout 300s;
|
||||
}
|
||||
location = /health {
|
||||
proxy_pass http://helexa_router;
|
||||
}
|
||||
|
||||
# Account control-plane → upstream /web/v1/ (strip the /api prefix).
|
||||
location /api/ {
|
||||
proxy_pass http://helexa_upstream/web/v1/;
|
||||
proxy_set_header Host $host;
|
||||
proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
|
||||
}
|
||||
|
||||
# SPA history fallback: anything else serves index.html.
|
||||
location / {
|
||||
try_files $uri $uri/ /index.html;
|
||||
}
|
||||
}
|
||||
@@ -23,7 +23,7 @@ const ROOT = path.resolve(
|
||||
const RESOURCES_DIR = path.join(ROOT, "src", "i18n", "resources");
|
||||
|
||||
// Namespaces to validate.
|
||||
const NAMESPACES = ["common", "home", "chat"];
|
||||
const NAMESPACES = ["common", "mission", "chat", "account"];
|
||||
|
||||
// Languages to validate should track SUPPORTED_LANGUAGES in src/i18n/languages.ts.
|
||||
// NOTE: This list is intentionally narrower than SUPPORTED_LANGUAGES and does not
|
||||
|
||||
@@ -409,3 +409,15 @@ pre {
|
||||
font-size: 0.8rem;
|
||||
}
|
||||
}
|
||||
|
||||
/* Public-beta banner (F6) — slim accent strip above the header. */
|
||||
.beta-banner {
|
||||
background: var(--bs-warning-bg-subtle, #fff3cd);
|
||||
color: var(--bs-warning-text-emphasis, #664d03);
|
||||
border-bottom: 1px solid var(--bs-warning-border-subtle, #ffe69c);
|
||||
}
|
||||
[data-bs-theme="dark"] .beta-banner {
|
||||
background: rgba(255, 193, 7, 0.12);
|
||||
color: #ffda6a;
|
||||
border-bottom-color: rgba(255, 193, 7, 0.25);
|
||||
}
|
||||
|
||||
@@ -1,34 +1,60 @@
|
||||
import { BrowserRouter, Routes, Route } from "react-router-dom";
|
||||
import { Container } from "react-bootstrap";
|
||||
import ThemeProvider from "./layout/ThemeProvider";
|
||||
import AuthProvider from "./auth/AuthProvider";
|
||||
import RequireAuth from "./auth/RequireAuth";
|
||||
import Header from "./components/Header";
|
||||
import Footer from "./components/Footer";
|
||||
import BetaBanner from "./components/BetaBanner";
|
||||
import Mission from "./pages/Mission";
|
||||
import Chat from "./pages/Chat";
|
||||
import Login from "./pages/auth/Login";
|
||||
import Register from "./pages/auth/Register";
|
||||
import VerifyEmail from "./pages/auth/VerifyEmail";
|
||||
import RequestReset from "./pages/auth/RequestReset";
|
||||
import ResetPassword from "./pages/auth/ResetPassword";
|
||||
import Dashboard from "./pages/account/Dashboard";
|
||||
import ApiKeys from "./pages/account/ApiKeys";
|
||||
import "./App.css";
|
||||
|
||||
// F1 composition root: theme + router + layout shell. The chat workspace
|
||||
// (`/`, F3), `/mission` (F2), and the auth/account routes (F4) replace these
|
||||
// placeholders in later phases.
|
||||
function Placeholder({ title }: { title: string }) {
|
||||
return (
|
||||
<Container className="py-5 flex-grow-1">
|
||||
<h1 className="mb-2">{title}</h1>
|
||||
<p className="text-muted">helexa public beta — coming online.</p>
|
||||
</Container>
|
||||
);
|
||||
}
|
||||
|
||||
// Composition root: theme → router → auth → layout shell. `/` is the chat
|
||||
// workspace (F3); `/mission` the EU-sovereignty narrative (F2); the auth +
|
||||
// account routes (F4) follow, with /account guarded.
|
||||
export default function App() {
|
||||
return (
|
||||
<ThemeProvider>
|
||||
<BrowserRouter>
|
||||
<div className="d-flex flex-column min-vh-100">
|
||||
<Header />
|
||||
<Routes>
|
||||
<Route path="/" element={<Placeholder title="Chat" />} />
|
||||
<Route path="/mission" element={<Placeholder title="Mission" />} />
|
||||
</Routes>
|
||||
<Footer />
|
||||
</div>
|
||||
<AuthProvider>
|
||||
<div className="d-flex flex-column min-vh-100">
|
||||
<BetaBanner />
|
||||
<Header />
|
||||
<Routes>
|
||||
<Route path="/" element={<Chat />} />
|
||||
<Route path="/mission" element={<Mission />} />
|
||||
<Route path="/login" element={<Login />} />
|
||||
<Route path="/register" element={<Register />} />
|
||||
<Route path="/verify" element={<VerifyEmail />} />
|
||||
<Route path="/forgot" element={<RequestReset />} />
|
||||
<Route path="/reset" element={<ResetPassword />} />
|
||||
<Route
|
||||
path="/account"
|
||||
element={
|
||||
<RequireAuth>
|
||||
<Dashboard />
|
||||
</RequireAuth>
|
||||
}
|
||||
/>
|
||||
<Route
|
||||
path="/account/keys"
|
||||
element={
|
||||
<RequireAuth>
|
||||
<ApiKeys />
|
||||
</RequireAuth>
|
||||
}
|
||||
/>
|
||||
</Routes>
|
||||
<Footer />
|
||||
</div>
|
||||
</AuthProvider>
|
||||
</BrowserRouter>
|
||||
</ThemeProvider>
|
||||
);
|
||||
|
||||
213
helexa.ai/src/api/account.ts
Normal file
213
helexa.ai/src/api/account.ts
Normal file
@@ -0,0 +1,213 @@
|
||||
// Account API client over helexa-upstream's /web/v1 (B4/B5). The browser
|
||||
// calls a same-origin `/api` prefix (vite-proxied in dev, nginx-routed in
|
||||
// prod). A MockAccountApi behind VITE_USE_MOCK_ACCOUNT_API lets the
|
||||
// dashboard be built/demoed before the upstream service is reachable.
|
||||
|
||||
import {
|
||||
ApiError,
|
||||
type AccountBalance,
|
||||
type ApiKeySummary,
|
||||
type CreatedKey,
|
||||
type Session,
|
||||
} from "./types";
|
||||
|
||||
export interface AccountApi {
|
||||
register(email: string, password: string, fingerprint?: string): Promise<void>;
|
||||
verify(token: string): Promise<void>;
|
||||
login(email: string, password: string): Promise<Session>;
|
||||
requestReset(email: string): Promise<void>;
|
||||
confirmReset(token: string, newPassword: string): Promise<void>;
|
||||
account(token: string): Promise<AccountBalance>;
|
||||
listKeys(token: string): Promise<ApiKeySummary[]>;
|
||||
createKey(
|
||||
token: string,
|
||||
label: string,
|
||||
limitKind: "percent" | "hardcap",
|
||||
limitValue: number,
|
||||
): Promise<CreatedKey>;
|
||||
archiveKey(token: string, id: string): Promise<void>;
|
||||
updateKeyLimit(
|
||||
token: string,
|
||||
id: string,
|
||||
limitKind: "percent" | "hardcap",
|
||||
limitValue: number,
|
||||
): Promise<void>;
|
||||
redeem(token: string, code: string): Promise<AccountBalance>;
|
||||
}
|
||||
|
||||
const BASE = (import.meta.env.VITE_ACCOUNT_BASE_URL || "/api").replace(/\/$/, "");
|
||||
|
||||
async function call<T>(
|
||||
path: string,
|
||||
init: RequestInit & { token?: string } = {},
|
||||
): Promise<T> {
|
||||
const headers: Record<string, string> = { "content-type": "application/json" };
|
||||
if (init.token) headers.authorization = `Bearer ${init.token}`;
|
||||
let resp: Response;
|
||||
try {
|
||||
resp = await fetch(`${BASE}${path}`, { ...init, headers });
|
||||
} catch {
|
||||
throw new ApiError(0, "network_error", "Could not reach the account service.");
|
||||
}
|
||||
if (resp.status === 204) return undefined as T;
|
||||
let body: unknown = null;
|
||||
try {
|
||||
body = await resp.json();
|
||||
} catch {
|
||||
/* empty body */
|
||||
}
|
||||
if (!resp.ok) {
|
||||
const err = (body as { error?: { code?: string; message?: string } })?.error;
|
||||
throw new ApiError(resp.status, err?.code ?? "error", err?.message ?? "Request failed.");
|
||||
}
|
||||
return body as T;
|
||||
}
|
||||
|
||||
class RealAccountApi implements AccountApi {
|
||||
async register(email: string, password: string, fingerprint?: string) {
|
||||
await call("/register", {
|
||||
method: "POST",
|
||||
body: JSON.stringify({ email, password, fingerprint }),
|
||||
});
|
||||
}
|
||||
async verify(token: string) {
|
||||
await call("/verify", { method: "POST", body: JSON.stringify({ token }) });
|
||||
}
|
||||
login(email: string, password: string) {
|
||||
return call<Session>("/login", {
|
||||
method: "POST",
|
||||
body: JSON.stringify({ email, password }),
|
||||
});
|
||||
}
|
||||
async requestReset(email: string) {
|
||||
await call("/password-reset/request", {
|
||||
method: "POST",
|
||||
body: JSON.stringify({ email }),
|
||||
});
|
||||
}
|
||||
async confirmReset(token: string, newPassword: string) {
|
||||
await call("/password-reset/confirm", {
|
||||
method: "POST",
|
||||
body: JSON.stringify({ token, new_password: newPassword }),
|
||||
});
|
||||
}
|
||||
account(token: string) {
|
||||
return call<AccountBalance>("/account", { token });
|
||||
}
|
||||
listKeys(token: string) {
|
||||
return call<{ keys: ApiKeySummary[] }>("/keys", { token }).then((r) => r.keys);
|
||||
}
|
||||
createKey(token: string, label: string, limit_kind: "percent" | "hardcap", limit_value: number) {
|
||||
return call<CreatedKey>("/keys", {
|
||||
method: "POST",
|
||||
token,
|
||||
body: JSON.stringify({ label, limit_kind, limit_value }),
|
||||
});
|
||||
}
|
||||
async archiveKey(token: string, id: string) {
|
||||
await call(`/keys/${id}/archive`, { method: "POST", token, body: "{}" });
|
||||
}
|
||||
async updateKeyLimit(
|
||||
token: string,
|
||||
id: string,
|
||||
limit_kind: "percent" | "hardcap",
|
||||
limit_value: number,
|
||||
) {
|
||||
await call(`/keys/${id}/limit`, {
|
||||
method: "PATCH",
|
||||
token,
|
||||
body: JSON.stringify({ limit_kind, limit_value }),
|
||||
});
|
||||
}
|
||||
redeem(token: string, code: string) {
|
||||
return call<AccountBalance>("/redeem", {
|
||||
method: "POST",
|
||||
token,
|
||||
body: JSON.stringify({ code }),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// ── Mock (VITE_USE_MOCK_ACCOUNT_API) ────────────────────────────────
|
||||
// Minimal in-memory account so the dashboard is fully developable offline.
|
||||
|
||||
class MockAccountApi implements AccountApi {
|
||||
private total = 1_000_000;
|
||||
private spent = 0;
|
||||
private reserved = 0;
|
||||
private keys: ApiKeySummary[] = [];
|
||||
private seq = 1;
|
||||
|
||||
async register() {}
|
||||
async verify() {}
|
||||
async login(): Promise<Session> {
|
||||
return { token: "mock-token", expires_in: 604800 };
|
||||
}
|
||||
async requestReset() {}
|
||||
async confirmReset() {}
|
||||
async account(): Promise<AccountBalance> {
|
||||
return {
|
||||
account_id: "mock-account",
|
||||
allocation_total: this.total,
|
||||
allocation_spent: this.spent,
|
||||
allocation_reserved: this.reserved,
|
||||
};
|
||||
}
|
||||
async listKeys(): Promise<ApiKeySummary[]> {
|
||||
return [...this.keys];
|
||||
}
|
||||
async createKey(
|
||||
_t: string,
|
||||
label: string,
|
||||
limit_kind: "percent" | "hardcap",
|
||||
limit_value: number,
|
||||
): Promise<CreatedKey> {
|
||||
const id = `mock-${this.seq++}`;
|
||||
const prefix = `sk-helexa-mock${this.seq}`;
|
||||
this.keys.push({
|
||||
id,
|
||||
prefix,
|
||||
label,
|
||||
status: "active",
|
||||
limit_kind,
|
||||
limit_value,
|
||||
spent: 0,
|
||||
reserved: 0,
|
||||
created_at: new Date().toISOString(),
|
||||
});
|
||||
return { id, key: `${prefix}-RAWSECRETSHOWNONCE`, prefix, limit_kind, limit_value };
|
||||
}
|
||||
async archiveKey(_t: string, id: string) {
|
||||
const k = this.keys.find((x) => x.id === id);
|
||||
if (k) k.status = "archived";
|
||||
}
|
||||
async updateKeyLimit(
|
||||
_t: string,
|
||||
id: string,
|
||||
limit_kind: "percent" | "hardcap",
|
||||
limit_value: number,
|
||||
) {
|
||||
const k = this.keys.find((x) => x.id === id);
|
||||
if (k) {
|
||||
k.limit_kind = limit_kind;
|
||||
k.limit_value = limit_value;
|
||||
}
|
||||
}
|
||||
async redeem(_t: string, code: string): Promise<AccountBalance> {
|
||||
if (!code.startsWith("helexa-topup-")) {
|
||||
throw new ApiError(400, "bad_request", "invalid or already-redeemed code");
|
||||
}
|
||||
this.total += 500_000;
|
||||
return this.account();
|
||||
}
|
||||
}
|
||||
|
||||
let instance: AccountApi | null = null;
|
||||
export function accountApi(): AccountApi {
|
||||
if (!instance) {
|
||||
instance = import.meta.env.VITE_USE_MOCK_ACCOUNT_API
|
||||
? new MockAccountApi()
|
||||
: new RealAccountApi();
|
||||
}
|
||||
return instance;
|
||||
}
|
||||
45
helexa.ai/src/api/types.ts
Normal file
45
helexa.ai/src/api/types.ts
Normal file
@@ -0,0 +1,45 @@
|
||||
// Wire types for the helexa-upstream /web/v1 account API (B4/B5).
|
||||
|
||||
export interface ApiKeySummary {
|
||||
id: string;
|
||||
prefix: string;
|
||||
label: string;
|
||||
status: "active" | "archived";
|
||||
limit_kind: "percent" | "hardcap";
|
||||
limit_value: number;
|
||||
spent: number;
|
||||
reserved: number;
|
||||
created_at: string;
|
||||
}
|
||||
|
||||
export interface CreatedKey {
|
||||
id: string;
|
||||
/** Raw secret — shown exactly once at creation. */
|
||||
key: string;
|
||||
prefix: string;
|
||||
limit_kind: "percent" | "hardcap";
|
||||
limit_value: number;
|
||||
}
|
||||
|
||||
export interface AccountBalance {
|
||||
account_id: string;
|
||||
allocation_total: number;
|
||||
allocation_spent: number;
|
||||
allocation_reserved: number;
|
||||
}
|
||||
|
||||
export interface Session {
|
||||
token: string;
|
||||
expires_in: number;
|
||||
}
|
||||
|
||||
/** Typed error carrying the backend's machine-readable code. */
|
||||
export class ApiError extends Error {
|
||||
code: string;
|
||||
status: number;
|
||||
constructor(status: number, code: string, message: string) {
|
||||
super(message);
|
||||
this.code = code;
|
||||
this.status = status;
|
||||
}
|
||||
}
|
||||
76
helexa.ai/src/auth/AuthProvider.tsx
Normal file
76
helexa.ai/src/auth/AuthProvider.tsx
Normal file
@@ -0,0 +1,76 @@
|
||||
import { useEffect, useState, type ReactNode } from "react";
|
||||
import { accountApi } from "../api/account";
|
||||
import { claimAnonymousData } from "../data/repositories";
|
||||
import { getFingerprint } from "../lib/fingerprint";
|
||||
import { AuthContext } from "./context";
|
||||
|
||||
const TOKEN_KEY = "helexa.token";
|
||||
const EMAIL_KEY = "helexa.email";
|
||||
|
||||
export default function AuthProvider({ children }: { children: ReactNode }) {
|
||||
const [token, setToken] = useState<string | null>(() =>
|
||||
localStorage.getItem(TOKEN_KEY),
|
||||
);
|
||||
const [email, setEmail] = useState<string | null>(() =>
|
||||
localStorage.getItem(EMAIL_KEY),
|
||||
);
|
||||
const [accountId, setAccountId] = useState<string | null>(null);
|
||||
|
||||
// Resolve the account id for an existing session (page reload) so the chat
|
||||
// workspace can scope its IndexedDB owner without a fresh login.
|
||||
useEffect(() => {
|
||||
if (!token || accountId) return;
|
||||
accountApi()
|
||||
.account(token)
|
||||
.then((a) => setAccountId(a.account_id))
|
||||
.catch(() => {
|
||||
/* token may be stale; chat falls back to anon until re-login */
|
||||
});
|
||||
}, [token, accountId]);
|
||||
|
||||
async function login(em: string, password: string): Promise<void> {
|
||||
const api = accountApi();
|
||||
const session = await api.login(em, password);
|
||||
localStorage.setItem(TOKEN_KEY, session.token);
|
||||
localStorage.setItem(EMAIL_KEY, em);
|
||||
setToken(session.token);
|
||||
setEmail(em);
|
||||
// Claim anonymous local history into the account (stays client-side).
|
||||
try {
|
||||
const acct = await api.account(session.token);
|
||||
setAccountId(acct.account_id);
|
||||
await claimAnonymousData(acct.account_id);
|
||||
} catch {
|
||||
/* non-fatal */
|
||||
}
|
||||
}
|
||||
|
||||
async function register(em: string, password: string): Promise<void> {
|
||||
const fingerprint = await getFingerprint();
|
||||
await accountApi().register(em, password, fingerprint);
|
||||
}
|
||||
|
||||
function logout(): void {
|
||||
localStorage.removeItem(TOKEN_KEY);
|
||||
localStorage.removeItem(EMAIL_KEY);
|
||||
setToken(null);
|
||||
setEmail(null);
|
||||
setAccountId(null);
|
||||
}
|
||||
|
||||
return (
|
||||
<AuthContext.Provider
|
||||
value={{
|
||||
token,
|
||||
email,
|
||||
accountId,
|
||||
status: token ? "authed" : "anon",
|
||||
login,
|
||||
register,
|
||||
logout,
|
||||
}}
|
||||
>
|
||||
{children}
|
||||
</AuthContext.Provider>
|
||||
);
|
||||
}
|
||||
14
helexa.ai/src/auth/RequireAuth.tsx
Normal file
14
helexa.ai/src/auth/RequireAuth.tsx
Normal file
@@ -0,0 +1,14 @@
|
||||
import { type ReactNode } from "react";
|
||||
import { Navigate, useLocation } from "react-router-dom";
|
||||
import { useAuth } from "./context";
|
||||
|
||||
/** Route guard: redirect unauthenticated users to /login?next=…. */
|
||||
export default function RequireAuth({ children }: { children: ReactNode }) {
|
||||
const { status } = useAuth();
|
||||
const location = useLocation();
|
||||
if (status !== "authed") {
|
||||
const next = encodeURIComponent(location.pathname + location.search);
|
||||
return <Navigate to={`/login?next=${next}`} replace />;
|
||||
}
|
||||
return <>{children}</>;
|
||||
}
|
||||
26
helexa.ai/src/auth/context.ts
Normal file
26
helexa.ai/src/auth/context.ts
Normal file
@@ -0,0 +1,26 @@
|
||||
import { createContext, useContext } from "react";
|
||||
|
||||
export interface AuthContextValue {
|
||||
token: string | null;
|
||||
email: string | null;
|
||||
/** The signed-in account id (for the Dexie owner + usage queries). */
|
||||
accountId: string | null;
|
||||
status: "anon" | "authed";
|
||||
login: (email: string, password: string) => Promise<void>;
|
||||
register: (email: string, password: string) => Promise<void>;
|
||||
logout: () => void;
|
||||
}
|
||||
|
||||
export const AuthContext = createContext<AuthContextValue>({
|
||||
token: null,
|
||||
email: null,
|
||||
accountId: null,
|
||||
status: "anon",
|
||||
login: async () => {},
|
||||
register: async () => {},
|
||||
logout: () => {},
|
||||
});
|
||||
|
||||
export function useAuth(): AuthContextValue {
|
||||
return useContext(AuthContext);
|
||||
}
|
||||
35
helexa.ai/src/components/BetaBanner.tsx
Normal file
35
helexa.ai/src/components/BetaBanner.tsx
Normal file
@@ -0,0 +1,35 @@
|
||||
import { useState } from "react";
|
||||
import { useTranslation } from "react-i18next";
|
||||
|
||||
/**
|
||||
* Slim public-beta notice shown above the header when VITE_PUBLIC_BETA is
|
||||
* set. Dismissible for the session (sessionStorage) so it doesn't nag, but
|
||||
* returns on the next visit while the beta lasts.
|
||||
*/
|
||||
const SHOWN = import.meta.env.VITE_PUBLIC_BETA === "true";
|
||||
const DISMISS_KEY = "helexa.betaDismissed";
|
||||
|
||||
export default function BetaBanner() {
|
||||
const { t } = useTranslation("common");
|
||||
const [hidden, setHidden] = useState(
|
||||
() => sessionStorage.getItem(DISMISS_KEY) === "1",
|
||||
);
|
||||
if (!SHOWN || hidden) return null;
|
||||
return (
|
||||
<div className="beta-banner d-flex align-items-center justify-content-center gap-2 px-3 py-1 small">
|
||||
<span>
|
||||
<strong>{t("beta.tag")}</strong> {t("beta.message")}
|
||||
</span>
|
||||
<button
|
||||
type="button"
|
||||
className="btn-close btn-close-white ms-2"
|
||||
style={{ fontSize: "0.6rem" }}
|
||||
aria-label={t("beta.dismiss")}
|
||||
onClick={() => {
|
||||
sessionStorage.setItem(DISMISS_KEY, "1");
|
||||
setHidden(true);
|
||||
}}
|
||||
/>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -6,11 +6,12 @@ import { useTheme } from "../layout/theme";
|
||||
import { useTranslation } from "react-i18next";
|
||||
import { AUTONYM_MAP, type LanguageCode, isRtlLanguage } from "../i18n/languages";
|
||||
import { getLanguageOptionsByUsage } from "../i18n/translation-priority";
|
||||
import { useAuth } from "../auth/context";
|
||||
|
||||
/**
|
||||
* Top navigation: brand, primary routes (chat at `/`, `/mission`), an
|
||||
* auth-aware cluster (stubbed until F4 wires sessions), the theme toggle,
|
||||
* and the language selector.
|
||||
* auth-aware cluster (Account/Sign out when signed in, else Sign in/up),
|
||||
* the theme toggle, and the language selector.
|
||||
*
|
||||
* The language picker is ordered by **estimated usage**
|
||||
* (getLanguageOptionsByUsage), not alphabetically — a deliberate choice that
|
||||
@@ -21,6 +22,7 @@ import { getLanguageOptionsByUsage } from "../i18n/translation-priority";
|
||||
const Header: React.FC = () => {
|
||||
const { theme, toggleTheme } = useTheme();
|
||||
const { t, i18n } = useTranslation("common");
|
||||
const { status, logout } = useAuth();
|
||||
|
||||
const currentLanguage: LanguageCode = (i18n.language.split("-")[0] ||
|
||||
"en") as LanguageCode;
|
||||
@@ -75,13 +77,31 @@ const Header: React.FC = () => {
|
||||
</Nav>
|
||||
|
||||
<div className="d-flex align-items-center gap-2">
|
||||
{/* Auth cluster — plain links until F4 wires session state. */}
|
||||
<NavLink to="/login" className="nav-link">
|
||||
{t("nav.login")}
|
||||
</NavLink>
|
||||
<NavLink to="/register" className="nav-link">
|
||||
{t("nav.register")}
|
||||
</NavLink>
|
||||
{/* Auth-aware cluster. */}
|
||||
{status === "authed" ? (
|
||||
<>
|
||||
<NavLink to="/account" className="nav-link">
|
||||
{t("nav.account")}
|
||||
</NavLink>
|
||||
<Button
|
||||
size="sm"
|
||||
variant="outline-secondary"
|
||||
onClick={logout}
|
||||
className="me-1"
|
||||
>
|
||||
{t("nav.logout")}
|
||||
</Button>
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
<NavLink to="/login" className="nav-link">
|
||||
{t("nav.login")}
|
||||
</NavLink>
|
||||
<NavLink to="/register" className="nav-link">
|
||||
{t("nav.register")}
|
||||
</NavLink>
|
||||
</>
|
||||
)}
|
||||
|
||||
<Button
|
||||
size="sm"
|
||||
|
||||
72
helexa.ai/src/data/db.ts
Normal file
72
helexa.ai/src/data/db.ts
Normal file
@@ -0,0 +1,72 @@
|
||||
// IndexedDB (Dexie) — the ONLY home for chat history and project
|
||||
// organisation. Nothing here is ever sent to a server (#69/#F3): the mesh
|
||||
// serves inference, but conversations live exclusively in the browser.
|
||||
//
|
||||
// `owner` namespaces data: `"anon"` for the fingerprinted anonymous visitor,
|
||||
// or an account id once signed in. On login, anonymous data can be claimed
|
||||
// into the account (F4) — still purely client-side.
|
||||
|
||||
import Dexie, { type Table } from "dexie";
|
||||
|
||||
export interface Project {
|
||||
id: string;
|
||||
owner: string;
|
||||
name: string;
|
||||
createdAt: number;
|
||||
updatedAt: number;
|
||||
archived: boolean;
|
||||
sortOrder: number;
|
||||
}
|
||||
|
||||
export interface Conversation {
|
||||
id: string;
|
||||
owner: string;
|
||||
projectId: string | null; // null → "Unsorted"
|
||||
title: string;
|
||||
model: string;
|
||||
createdAt: number;
|
||||
updatedAt: number;
|
||||
pinned: boolean;
|
||||
}
|
||||
|
||||
export type MessageRole = "system" | "user" | "assistant";
|
||||
export type MessageStatus = "complete" | "streaming" | "error";
|
||||
|
||||
export interface Message {
|
||||
id: string;
|
||||
conversationId: string;
|
||||
role: MessageRole;
|
||||
content: string;
|
||||
createdAt: number;
|
||||
status: MessageStatus;
|
||||
errorCode?: string;
|
||||
promptTokens?: number;
|
||||
completionTokens?: number;
|
||||
}
|
||||
|
||||
/** Small key/value store: fingerprint, active conversation, anon usage. */
|
||||
export interface Meta {
|
||||
key: string;
|
||||
value: unknown;
|
||||
}
|
||||
|
||||
class HelexaDB extends Dexie {
|
||||
projects!: Table<Project, string>;
|
||||
conversations!: Table<Conversation, string>;
|
||||
messages!: Table<Message, string>;
|
||||
meta!: Table<Meta, string>;
|
||||
|
||||
constructor() {
|
||||
super("helexa");
|
||||
this.version(1).stores({
|
||||
// Indexes only — Dexie stores the whole object. Compound indexes
|
||||
// drive the common queries (by owner, by conversation in time order).
|
||||
projects: "id, owner, [owner+archived], updatedAt",
|
||||
conversations: "id, owner, projectId, [owner+projectId], updatedAt",
|
||||
messages: "id, conversationId, [conversationId+createdAt]",
|
||||
meta: "key",
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
export const db = new HelexaDB();
|
||||
154
helexa.ai/src/data/repositories.ts
Normal file
154
helexa.ai/src/data/repositories.ts
Normal file
@@ -0,0 +1,154 @@
|
||||
// Typed CRUD + queries over the Dexie store. UI components use the
|
||||
// `useLiveQuery` hook (dexie-react-hooks) with the list helpers here so the
|
||||
// sidebar/thread react to writes automatically.
|
||||
|
||||
import Dexie from "dexie";
|
||||
import {
|
||||
db,
|
||||
type Conversation,
|
||||
type Message,
|
||||
type MessageRole,
|
||||
type Project,
|
||||
} from "./db";
|
||||
|
||||
function uuid(): string {
|
||||
return crypto.randomUUID();
|
||||
}
|
||||
function now(): number {
|
||||
return Date.now();
|
||||
}
|
||||
|
||||
// ── projects ────────────────────────────────────────────────────────
|
||||
|
||||
export async function listProjects(owner: string): Promise<Project[]> {
|
||||
const rows = await db.projects.where({ owner }).toArray();
|
||||
return rows
|
||||
.filter((p) => !p.archived)
|
||||
.sort((a, b) => a.sortOrder - b.sortOrder || a.createdAt - b.createdAt);
|
||||
}
|
||||
|
||||
export async function createProject(owner: string, name: string): Promise<string> {
|
||||
const id = uuid();
|
||||
const ts = now();
|
||||
await db.projects.add({
|
||||
id,
|
||||
owner,
|
||||
name,
|
||||
createdAt: ts,
|
||||
updatedAt: ts,
|
||||
archived: false,
|
||||
sortOrder: ts,
|
||||
});
|
||||
return id;
|
||||
}
|
||||
|
||||
export async function renameProject(id: string, name: string): Promise<void> {
|
||||
await db.projects.update(id, { name, updatedAt: now() });
|
||||
}
|
||||
|
||||
export async function archiveProject(id: string): Promise<void> {
|
||||
// Detach its conversations to "Unsorted" so nothing is orphaned.
|
||||
await db.transaction("rw", db.projects, db.conversations, async () => {
|
||||
await db.projects.update(id, { archived: true, updatedAt: now() });
|
||||
const convs = await db.conversations.where({ projectId: id }).toArray();
|
||||
await Promise.all(
|
||||
convs.map((c) => db.conversations.update(c.id, { projectId: null })),
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
// ── conversations ───────────────────────────────────────────────────
|
||||
|
||||
export async function listConversations(owner: string): Promise<Conversation[]> {
|
||||
const rows = await db.conversations.where({ owner }).toArray();
|
||||
return rows.sort(
|
||||
(a, b) => Number(b.pinned) - Number(a.pinned) || b.updatedAt - a.updatedAt,
|
||||
);
|
||||
}
|
||||
|
||||
export async function createConversation(
|
||||
owner: string,
|
||||
model: string,
|
||||
projectId: string | null = null,
|
||||
title = "New chat",
|
||||
): Promise<string> {
|
||||
const id = uuid();
|
||||
const ts = now();
|
||||
await db.conversations.add({
|
||||
id,
|
||||
owner,
|
||||
projectId,
|
||||
title,
|
||||
model,
|
||||
createdAt: ts,
|
||||
updatedAt: ts,
|
||||
pinned: false,
|
||||
});
|
||||
return id;
|
||||
}
|
||||
|
||||
export async function renameConversation(id: string, title: string): Promise<void> {
|
||||
await db.conversations.update(id, { title, updatedAt: now() });
|
||||
}
|
||||
|
||||
export async function moveConversation(
|
||||
id: string,
|
||||
projectId: string | null,
|
||||
): Promise<void> {
|
||||
await db.conversations.update(id, { projectId, updatedAt: now() });
|
||||
}
|
||||
|
||||
export async function deleteConversation(id: string): Promise<void> {
|
||||
await db.transaction("rw", db.conversations, db.messages, async () => {
|
||||
await db.messages.where({ conversationId: id }).delete();
|
||||
await db.conversations.delete(id);
|
||||
});
|
||||
}
|
||||
|
||||
// ── messages ────────────────────────────────────────────────────────
|
||||
|
||||
export async function listMessages(conversationId: string): Promise<Message[]> {
|
||||
return db.messages
|
||||
.where("[conversationId+createdAt]")
|
||||
.between([conversationId, Dexie.minKey], [conversationId, Dexie.maxKey])
|
||||
.toArray();
|
||||
}
|
||||
|
||||
export async function addMessage(
|
||||
conversationId: string,
|
||||
role: MessageRole,
|
||||
content: string,
|
||||
status: Message["status"] = "complete",
|
||||
): Promise<string> {
|
||||
const id = uuid();
|
||||
await db.messages.add({ id, conversationId, role, content, createdAt: now(), status });
|
||||
await db.conversations.update(conversationId, { updatedAt: now() });
|
||||
return id;
|
||||
}
|
||||
|
||||
export async function appendToMessage(id: string, delta: string): Promise<void> {
|
||||
const msg = await db.messages.get(id);
|
||||
if (!msg) return;
|
||||
await db.messages.update(id, { content: msg.content + delta });
|
||||
}
|
||||
|
||||
export async function finalizeMessage(
|
||||
id: string,
|
||||
patch: Partial<Pick<Message, "status" | "errorCode" | "promptTokens" | "completionTokens">>,
|
||||
): Promise<void> {
|
||||
await db.messages.update(id, patch);
|
||||
}
|
||||
|
||||
/** Rewrite all `anon` data to `accountId` on first login (stays local). */
|
||||
export async function claimAnonymousData(accountId: string): Promise<void> {
|
||||
await db.transaction("rw", db.projects, db.conversations, async () => {
|
||||
const projects = await db.projects.where({ owner: "anon" }).toArray();
|
||||
await Promise.all(
|
||||
projects.map((p) => db.projects.update(p.id, { owner: accountId })),
|
||||
);
|
||||
const convs = await db.conversations.where({ owner: "anon" }).toArray();
|
||||
await Promise.all(
|
||||
convs.map((c) => db.conversations.update(c.id, { owner: accountId })),
|
||||
);
|
||||
});
|
||||
}
|
||||
@@ -10,134 +10,166 @@ import type { LanguageCode } from "./languages";
|
||||
// Core languages
|
||||
import enCommon from "./resources/en/common.json";
|
||||
import ruCommon from "./resources/ru/common.json";
|
||||
import enHome from "./resources/en/home.json";
|
||||
import ruHome from "./resources/ru/home.json";
|
||||
import enMission from "./resources/en/mission.json";
|
||||
import ruMission from "./resources/ru/mission.json";
|
||||
import enChat from "./resources/en/chat.json";
|
||||
import enAccount from "./resources/en/account.json";
|
||||
import ruChat from "./resources/ru/chat.json";
|
||||
import ruAccount from "./resources/ru/account.json";
|
||||
|
||||
// Scandinavian & Nordic languages
|
||||
import daCommon from "./resources/da/common.json";
|
||||
import daHome from "./resources/da/home.json";
|
||||
import daMission from "./resources/da/mission.json";
|
||||
import daChat from "./resources/da/chat.json";
|
||||
import daAccount from "./resources/da/account.json";
|
||||
|
||||
import fiCommon from "./resources/fi/common.json";
|
||||
import fiHome from "./resources/fi/home.json";
|
||||
import fiMission from "./resources/fi/mission.json";
|
||||
import fiChat from "./resources/fi/chat.json";
|
||||
import fiAccount from "./resources/fi/account.json";
|
||||
|
||||
import noCommon from "./resources/no/common.json";
|
||||
import noHome from "./resources/no/home.json";
|
||||
import noMission from "./resources/no/mission.json";
|
||||
import noChat from "./resources/no/chat.json";
|
||||
import noAccount from "./resources/no/account.json";
|
||||
|
||||
import svCommon from "./resources/sv/common.json";
|
||||
import svHome from "./resources/sv/home.json";
|
||||
import svMission from "./resources/sv/mission.json";
|
||||
import svChat from "./resources/sv/chat.json";
|
||||
import svAccount from "./resources/sv/account.json";
|
||||
|
||||
import bgCommon from "./resources/bg/common.json";
|
||||
import bgHome from "./resources/bg/home.json";
|
||||
import bgMission from "./resources/bg/mission.json";
|
||||
import bgChat from "./resources/bg/chat.json";
|
||||
import bgAccount from "./resources/bg/account.json";
|
||||
|
||||
import etCommon from "./resources/et/common.json";
|
||||
import etHome from "./resources/et/home.json";
|
||||
import etMission from "./resources/et/mission.json";
|
||||
import etChat from "./resources/et/chat.json";
|
||||
import etAccount from "./resources/et/account.json";
|
||||
|
||||
// African & MENA languages
|
||||
import swCommon from "./resources/sw/common.json";
|
||||
import swHome from "./resources/sw/home.json";
|
||||
import swMission from "./resources/sw/mission.json";
|
||||
import swChat from "./resources/sw/chat.json";
|
||||
import swAccount from "./resources/sw/account.json";
|
||||
|
||||
import arCommon from "./resources/ar/common.json";
|
||||
import arHome from "./resources/ar/home.json";
|
||||
import arMission from "./resources/ar/mission.json";
|
||||
import arChat from "./resources/ar/chat.json";
|
||||
import arAccount from "./resources/ar/account.json";
|
||||
|
||||
import faCommon from "./resources/fa/common.json";
|
||||
import faHome from "./resources/fa/home.json";
|
||||
import faMission from "./resources/fa/mission.json";
|
||||
import faChat from "./resources/fa/chat.json";
|
||||
import faAccount from "./resources/fa/account.json";
|
||||
|
||||
import haCommon from "./resources/ha/common.json";
|
||||
import haHome from "./resources/ha/home.json";
|
||||
import haMission from "./resources/ha/mission.json";
|
||||
import haChat from "./resources/ha/chat.json";
|
||||
import haAccount from "./resources/ha/account.json";
|
||||
|
||||
import amCommon from "./resources/am/common.json";
|
||||
import amHome from "./resources/am/home.json";
|
||||
import amMission from "./resources/am/mission.json";
|
||||
import amChat from "./resources/am/chat.json";
|
||||
import amAccount from "./resources/am/account.json";
|
||||
|
||||
import yoCommon from "./resources/yo/common.json";
|
||||
import yoHome from "./resources/yo/home.json";
|
||||
import yoMission from "./resources/yo/mission.json";
|
||||
import yoChat from "./resources/yo/chat.json";
|
||||
import yoAccount from "./resources/yo/account.json";
|
||||
|
||||
import zuCommon from "./resources/zu/common.json";
|
||||
import zuHome from "./resources/zu/home.json";
|
||||
import zuMission from "./resources/zu/mission.json";
|
||||
import zuChat from "./resources/zu/chat.json";
|
||||
import zuAccount from "./resources/zu/account.json";
|
||||
|
||||
// Darija (Moroccan Arabic)
|
||||
import maCommon from "./resources/ma/common.json";
|
||||
import maHome from "./resources/ma/home.json";
|
||||
import maMission from "./resources/ma/mission.json";
|
||||
import maChat from "./resources/ma/chat.json";
|
||||
import maAccount from "./resources/ma/account.json";
|
||||
|
||||
// European / other languages
|
||||
import esCommon from "./resources/es/common.json";
|
||||
import esHome from "./resources/es/home.json";
|
||||
import esMission from "./resources/es/mission.json";
|
||||
import esChat from "./resources/es/chat.json";
|
||||
import esAccount from "./resources/es/account.json";
|
||||
|
||||
import frCommon from "./resources/fr/common.json";
|
||||
import frHome from "./resources/fr/home.json";
|
||||
import frMission from "./resources/fr/mission.json";
|
||||
import frChat from "./resources/fr/chat.json";
|
||||
import frAccount from "./resources/fr/account.json";
|
||||
|
||||
import deCommon from "./resources/de/common.json";
|
||||
import deHome from "./resources/de/home.json";
|
||||
import deMission from "./resources/de/mission.json";
|
||||
import deChat from "./resources/de/chat.json";
|
||||
import deAccount from "./resources/de/account.json";
|
||||
|
||||
import elCommon from "./resources/el/common.json";
|
||||
import elHome from "./resources/el/home.json";
|
||||
import elMission from "./resources/el/mission.json";
|
||||
import elChat from "./resources/el/chat.json";
|
||||
import elAccount from "./resources/el/account.json";
|
||||
|
||||
import itCommon from "./resources/it/common.json";
|
||||
import itHome from "./resources/it/home.json";
|
||||
import itMission from "./resources/it/mission.json";
|
||||
import itChat from "./resources/it/chat.json";
|
||||
import itAccount from "./resources/it/account.json";
|
||||
|
||||
import heCommon from "./resources/he/common.json";
|
||||
import heHome from "./resources/he/home.json";
|
||||
import heMission from "./resources/he/mission.json";
|
||||
import heChat from "./resources/he/chat.json";
|
||||
import heAccount from "./resources/he/account.json";
|
||||
|
||||
import ptCommon from "./resources/pt/common.json";
|
||||
import ptHome from "./resources/pt/home.json";
|
||||
import ptMission from "./resources/pt/mission.json";
|
||||
import ptChat from "./resources/pt/chat.json";
|
||||
import ptAccount from "./resources/pt/account.json";
|
||||
|
||||
import roCommon from "./resources/ro/common.json";
|
||||
import roHome from "./resources/ro/home.json";
|
||||
import roMission from "./resources/ro/mission.json";
|
||||
import roChat from "./resources/ro/chat.json";
|
||||
import roAccount from "./resources/ro/account.json";
|
||||
|
||||
import kaCommon from "./resources/ka/common.json";
|
||||
import kaHome from "./resources/ka/home.json";
|
||||
import kaMission from "./resources/ka/mission.json";
|
||||
import kaChat from "./resources/ka/chat.json";
|
||||
import kaAccount from "./resources/ka/account.json";
|
||||
|
||||
import trCommon from "./resources/tr/common.json";
|
||||
import trHome from "./resources/tr/home.json";
|
||||
import trMission from "./resources/tr/mission.json";
|
||||
import trChat from "./resources/tr/chat.json";
|
||||
import trAccount from "./resources/tr/account.json";
|
||||
|
||||
import plCommon from "./resources/pl/common.json";
|
||||
import plHome from "./resources/pl/home.json";
|
||||
import plMission from "./resources/pl/mission.json";
|
||||
import plChat from "./resources/pl/chat.json";
|
||||
import plAccount from "./resources/pl/account.json";
|
||||
|
||||
import ukCommon from "./resources/uk/common.json";
|
||||
import ukHome from "./resources/uk/home.json";
|
||||
import ukMission from "./resources/uk/mission.json";
|
||||
import ukChat from "./resources/uk/chat.json";
|
||||
import ukAccount from "./resources/uk/account.json";
|
||||
|
||||
import nlCommon from "./resources/nl/common.json";
|
||||
import nlHome from "./resources/nl/home.json";
|
||||
import nlMission from "./resources/nl/mission.json";
|
||||
import nlChat from "./resources/nl/chat.json";
|
||||
import nlAccount from "./resources/nl/account.json";
|
||||
|
||||
import srCommon from "./resources/sr/common.json";
|
||||
import srHome from "./resources/sr/home.json";
|
||||
import srMission from "./resources/sr/mission.json";
|
||||
import srChat from "./resources/sr/chat.json";
|
||||
import srAccount from "./resources/sr/account.json";
|
||||
|
||||
import kkCommon from "./resources/kk/common.json";
|
||||
import kkHome from "./resources/kk/home.json";
|
||||
import kkMission from "./resources/kk/mission.json";
|
||||
import kkChat from "./resources/kk/chat.json";
|
||||
import kkAccount from "./resources/kk/account.json";
|
||||
|
||||
import uzCommon from "./resources/uz/common.json";
|
||||
import uzHome from "./resources/uz/home.json";
|
||||
import uzMission from "./resources/uz/mission.json";
|
||||
import uzChat from "./resources/uz/chat.json";
|
||||
import uzAccount from "./resources/uz/account.json";
|
||||
|
||||
/**
|
||||
* Application translation resources, split by language and namespace.
|
||||
@@ -149,167 +181,199 @@ import uzChat from "./resources/uz/chat.json";
|
||||
const resources: Resource = {
|
||||
en: {
|
||||
common: enCommon,
|
||||
home: enHome,
|
||||
mission: enMission,
|
||||
chat: enChat,
|
||||
account: enAccount,
|
||||
},
|
||||
ru: {
|
||||
common: ruCommon,
|
||||
home: ruHome,
|
||||
mission: ruMission,
|
||||
chat: ruChat,
|
||||
account: ruAccount,
|
||||
},
|
||||
bg: {
|
||||
common: bgCommon,
|
||||
home: bgHome,
|
||||
mission: bgMission,
|
||||
chat: bgChat,
|
||||
account: bgAccount,
|
||||
},
|
||||
da: {
|
||||
common: daCommon,
|
||||
home: daHome,
|
||||
mission: daMission,
|
||||
chat: daChat,
|
||||
account: daAccount,
|
||||
},
|
||||
et: {
|
||||
common: etCommon,
|
||||
home: etHome,
|
||||
mission: etMission,
|
||||
chat: etChat,
|
||||
account: etAccount,
|
||||
},
|
||||
fi: {
|
||||
common: fiCommon,
|
||||
home: fiHome,
|
||||
mission: fiMission,
|
||||
chat: fiChat,
|
||||
account: fiAccount,
|
||||
},
|
||||
kk: {
|
||||
common: kkCommon,
|
||||
home: kkHome,
|
||||
mission: kkMission,
|
||||
chat: kkChat,
|
||||
account: kkAccount,
|
||||
},
|
||||
uz: {
|
||||
common: uzCommon,
|
||||
home: uzHome,
|
||||
mission: uzMission,
|
||||
chat: uzChat,
|
||||
account: uzAccount,
|
||||
},
|
||||
|
||||
// African & MENA languages (LTR unless marked RTL via isRtlLanguage)
|
||||
sw: {
|
||||
common: swCommon,
|
||||
home: swHome,
|
||||
mission: swMission,
|
||||
chat: swChat,
|
||||
account: swAccount,
|
||||
},
|
||||
ar: {
|
||||
common: arCommon,
|
||||
home: arHome,
|
||||
mission: arMission,
|
||||
chat: arChat,
|
||||
account: arAccount,
|
||||
},
|
||||
fa: {
|
||||
common: faCommon,
|
||||
home: faHome,
|
||||
mission: faMission,
|
||||
chat: faChat,
|
||||
account: faAccount,
|
||||
},
|
||||
ha: {
|
||||
common: haCommon,
|
||||
home: haHome,
|
||||
mission: haMission,
|
||||
chat: haChat,
|
||||
account: haAccount,
|
||||
},
|
||||
am: {
|
||||
common: amCommon,
|
||||
home: amHome,
|
||||
mission: amMission,
|
||||
chat: amChat,
|
||||
account: amAccount,
|
||||
},
|
||||
yo: {
|
||||
common: yoCommon,
|
||||
home: yoHome,
|
||||
mission: yoMission,
|
||||
chat: yoChat,
|
||||
account: yoAccount,
|
||||
},
|
||||
zu: {
|
||||
common: zuCommon,
|
||||
home: zuHome,
|
||||
mission: zuMission,
|
||||
chat: zuChat,
|
||||
account: zuAccount,
|
||||
},
|
||||
ma: {
|
||||
common: maCommon,
|
||||
home: maHome,
|
||||
mission: maMission,
|
||||
chat: maChat,
|
||||
account: maAccount,
|
||||
},
|
||||
|
||||
// European & other languages
|
||||
es: {
|
||||
common: esCommon,
|
||||
home: esHome,
|
||||
mission: esMission,
|
||||
chat: esChat,
|
||||
account: esAccount,
|
||||
},
|
||||
fr: {
|
||||
common: frCommon,
|
||||
home: frHome,
|
||||
mission: frMission,
|
||||
chat: frChat,
|
||||
account: frAccount,
|
||||
},
|
||||
de: {
|
||||
common: deCommon,
|
||||
home: deHome,
|
||||
mission: deMission,
|
||||
chat: deChat,
|
||||
account: deAccount,
|
||||
},
|
||||
el: {
|
||||
common: elCommon,
|
||||
home: elHome,
|
||||
mission: elMission,
|
||||
chat: elChat,
|
||||
account: elAccount,
|
||||
},
|
||||
it: {
|
||||
common: itCommon,
|
||||
home: itHome,
|
||||
mission: itMission,
|
||||
chat: itChat,
|
||||
account: itAccount,
|
||||
},
|
||||
he: {
|
||||
common: heCommon,
|
||||
home: heHome,
|
||||
mission: heMission,
|
||||
chat: heChat,
|
||||
account: heAccount,
|
||||
},
|
||||
pt: {
|
||||
common: ptCommon,
|
||||
home: ptHome,
|
||||
mission: ptMission,
|
||||
chat: ptChat,
|
||||
account: ptAccount,
|
||||
},
|
||||
ro: {
|
||||
common: roCommon,
|
||||
home: roHome,
|
||||
mission: roMission,
|
||||
chat: roChat,
|
||||
account: roAccount,
|
||||
},
|
||||
ka: {
|
||||
common: kaCommon,
|
||||
home: kaHome,
|
||||
mission: kaMission,
|
||||
chat: kaChat,
|
||||
account: kaAccount,
|
||||
},
|
||||
tr: {
|
||||
common: trCommon,
|
||||
home: trHome,
|
||||
mission: trMission,
|
||||
chat: trChat,
|
||||
account: trAccount,
|
||||
},
|
||||
pl: {
|
||||
common: plCommon,
|
||||
home: plHome,
|
||||
mission: plMission,
|
||||
chat: plChat,
|
||||
account: plAccount,
|
||||
},
|
||||
uk: {
|
||||
common: ukCommon,
|
||||
home: ukHome,
|
||||
mission: ukMission,
|
||||
chat: ukChat,
|
||||
account: ukAccount,
|
||||
},
|
||||
nl: {
|
||||
common: nlCommon,
|
||||
home: nlHome,
|
||||
mission: nlMission,
|
||||
chat: nlChat,
|
||||
account: nlAccount,
|
||||
},
|
||||
sr: {
|
||||
common: srCommon,
|
||||
home: srHome,
|
||||
mission: srMission,
|
||||
chat: srChat,
|
||||
account: srAccount,
|
||||
},
|
||||
no: {
|
||||
common: noCommon,
|
||||
home: noHome,
|
||||
mission: noMission,
|
||||
chat: noChat,
|
||||
account: noAccount,
|
||||
},
|
||||
sv: {
|
||||
common: svCommon,
|
||||
home: svHome,
|
||||
mission: svMission,
|
||||
chat: svChat,
|
||||
account: svAccount,
|
||||
},
|
||||
};
|
||||
|
||||
@@ -335,7 +399,7 @@ i18n.use(initReactI18next).init({
|
||||
lng: browserLang,
|
||||
fallbackLng: "en",
|
||||
supportedLngs: SUPPORTED_LANGUAGES,
|
||||
ns: ["common", "home", "chat"],
|
||||
ns: ["common", "mission", "chat", "account"],
|
||||
defaultNS: "common",
|
||||
// Because we control the keys and interpolate only simple values.
|
||||
interpolation: {
|
||||
|
||||
71
helexa.ai/src/i18n/resources/am/account.json
Normal file
71
helexa.ai/src/i18n/resources/am/account.json
Normal file
@@ -0,0 +1,71 @@
|
||||
{
|
||||
"login": {
|
||||
"title": "Sign in",
|
||||
"email": "Email",
|
||||
"password": "Password",
|
||||
"submit": "Sign in",
|
||||
"noAccount": "No account? Sign up"
|
||||
},
|
||||
"register": {
|
||||
"title": "Create your account",
|
||||
"email": "Email",
|
||||
"password": "Password",
|
||||
"submit": "Sign up",
|
||||
"haveAccount": "Already have an account? Sign in",
|
||||
"checkEmail": "Almost there — check your email to verify your account."
|
||||
},
|
||||
"verify": {
|
||||
"verifying": "Verifying…",
|
||||
"ok": "Email verified. You can now sign in.",
|
||||
"failed": "This verification link is invalid or has expired.",
|
||||
"toLogin": "Go to sign in"
|
||||
},
|
||||
"reset": {
|
||||
"requestTitle": "Reset your password",
|
||||
"email": "Email",
|
||||
"requestSubmit": "Send reset link",
|
||||
"requestDone": "If that email has an account, a reset link is on its way.",
|
||||
"confirmTitle": "Choose a new password",
|
||||
"newPassword": "New password",
|
||||
"confirmSubmit": "Set password",
|
||||
"ok": "Password updated. You can now sign in."
|
||||
},
|
||||
"dashboard": {
|
||||
"title": "Account",
|
||||
"balance": "Allocation",
|
||||
"total": "Total",
|
||||
"spent": "Spent",
|
||||
"reserved": "Reserved",
|
||||
"remaining": "Remaining",
|
||||
"manageKeys": "Manage API keys",
|
||||
"redeemTitle": "Redeem a top-up code",
|
||||
"redeemPlaceholder": "helexa-topup-…",
|
||||
"redeem": "Redeem",
|
||||
"redeemed": "Code redeemed.",
|
||||
"logout": "Sign out"
|
||||
},
|
||||
"keys": {
|
||||
"title": "API keys",
|
||||
"create": "Create key",
|
||||
"label": "Label",
|
||||
"limitKind": "Limit",
|
||||
"percent": "% of allocation",
|
||||
"hardcap": "Hard cap (tokens)",
|
||||
"value": "Value",
|
||||
"none": "No keys yet.",
|
||||
"createdTitle": "Your new API key",
|
||||
"createdWarn": "Copy it now — you won't see it again.",
|
||||
"copy": "Copy",
|
||||
"copied": "Copied",
|
||||
"archive": "Archive",
|
||||
"save": "Save",
|
||||
"status": "Status",
|
||||
"usage": "Used",
|
||||
"useForChat": "Use for chat on this device",
|
||||
"usedForChat": "Enabled for chat ✓"
|
||||
},
|
||||
"error": {
|
||||
"generic": "Something went wrong.",
|
||||
"unauthorized": "Please sign in again."
|
||||
}
|
||||
}
|
||||
@@ -5,5 +5,17 @@
|
||||
"transcriptPlaceholder": "የውይይቱ ሪኮርድ እዚህ ይታያል። የሞዴሉን እና የተጠቃሚውን መልዕክቶች በሚንቀሳቀስ ኮንቴይነር ውስጥ ያቀርቡ፣ приወይም በዙር ዙር በመከፈል ማቅረብ ይችላሉ።",
|
||||
"inputPlaceholder": "ውይይትን ለመጀምር መልዕክት ይፃፉ…",
|
||||
"send": "መላክ",
|
||||
"clear": "ማጽዳት"
|
||||
"clear": "ማጽዳት",
|
||||
"newChat": "New chat",
|
||||
"newProject": "New project",
|
||||
"newProjectName": "New project",
|
||||
"unsorted": "Unsorted",
|
||||
"emptyState": "Start a conversation. Your history stays in this browser.",
|
||||
"anonBanner": "You have reached the anonymous limit. Sign up for a free allocation.",
|
||||
"signUp": "Sign up",
|
||||
"stop": "Stop",
|
||||
"topUp": "Top up",
|
||||
"rateLimited": "Rate limited — wait a moment and retry.",
|
||||
"needsKey": "Create an API key and enable it for chat to send as yourself.",
|
||||
"manageKeysLink": "Manage keys"
|
||||
}
|
||||
|
||||
@@ -59,5 +59,10 @@
|
||||
},
|
||||
"footer": {
|
||||
"copyright": "© {{year}} helexa.ai"
|
||||
},
|
||||
"beta": {
|
||||
"tag": "Public beta",
|
||||
"message": "helexa is in open beta — expect rough edges. Your chats stay in your browser.",
|
||||
"dismiss": "Dismiss"
|
||||
}
|
||||
}
|
||||
|
||||
71
helexa.ai/src/i18n/resources/ar/account.json
Normal file
71
helexa.ai/src/i18n/resources/ar/account.json
Normal file
@@ -0,0 +1,71 @@
|
||||
{
|
||||
"login": {
|
||||
"title": "Sign in",
|
||||
"email": "Email",
|
||||
"password": "Password",
|
||||
"submit": "Sign in",
|
||||
"noAccount": "No account? Sign up"
|
||||
},
|
||||
"register": {
|
||||
"title": "Create your account",
|
||||
"email": "Email",
|
||||
"password": "Password",
|
||||
"submit": "Sign up",
|
||||
"haveAccount": "Already have an account? Sign in",
|
||||
"checkEmail": "Almost there — check your email to verify your account."
|
||||
},
|
||||
"verify": {
|
||||
"verifying": "Verifying…",
|
||||
"ok": "Email verified. You can now sign in.",
|
||||
"failed": "This verification link is invalid or has expired.",
|
||||
"toLogin": "Go to sign in"
|
||||
},
|
||||
"reset": {
|
||||
"requestTitle": "Reset your password",
|
||||
"email": "Email",
|
||||
"requestSubmit": "Send reset link",
|
||||
"requestDone": "If that email has an account, a reset link is on its way.",
|
||||
"confirmTitle": "Choose a new password",
|
||||
"newPassword": "New password",
|
||||
"confirmSubmit": "Set password",
|
||||
"ok": "Password updated. You can now sign in."
|
||||
},
|
||||
"dashboard": {
|
||||
"title": "Account",
|
||||
"balance": "Allocation",
|
||||
"total": "Total",
|
||||
"spent": "Spent",
|
||||
"reserved": "Reserved",
|
||||
"remaining": "Remaining",
|
||||
"manageKeys": "Manage API keys",
|
||||
"redeemTitle": "Redeem a top-up code",
|
||||
"redeemPlaceholder": "helexa-topup-…",
|
||||
"redeem": "Redeem",
|
||||
"redeemed": "Code redeemed.",
|
||||
"logout": "Sign out"
|
||||
},
|
||||
"keys": {
|
||||
"title": "API keys",
|
||||
"create": "Create key",
|
||||
"label": "Label",
|
||||
"limitKind": "Limit",
|
||||
"percent": "% of allocation",
|
||||
"hardcap": "Hard cap (tokens)",
|
||||
"value": "Value",
|
||||
"none": "No keys yet.",
|
||||
"createdTitle": "Your new API key",
|
||||
"createdWarn": "Copy it now — you won't see it again.",
|
||||
"copy": "Copy",
|
||||
"copied": "Copied",
|
||||
"archive": "Archive",
|
||||
"save": "Save",
|
||||
"status": "Status",
|
||||
"usage": "Used",
|
||||
"useForChat": "Use for chat on this device",
|
||||
"usedForChat": "Enabled for chat ✓"
|
||||
},
|
||||
"error": {
|
||||
"generic": "Something went wrong.",
|
||||
"unauthorized": "Please sign in again."
|
||||
}
|
||||
}
|
||||
@@ -5,5 +5,17 @@
|
||||
"transcriptPlaceholder": "سجل المحادثة سيظهر هنا. اعرض رسائل النموذج والمستخدم في حاوية قابلة للتمرير، ويمكنك تجميعها حسب أدوار الحوار إذا رغبت.",
|
||||
"inputPlaceholder": "اكتب رسالة لبدء المحادثة…",
|
||||
"send": "إرسال",
|
||||
"clear": "مسح"
|
||||
"clear": "مسح",
|
||||
"newChat": "New chat",
|
||||
"newProject": "New project",
|
||||
"newProjectName": "New project",
|
||||
"unsorted": "Unsorted",
|
||||
"emptyState": "Start a conversation. Your history stays in this browser.",
|
||||
"anonBanner": "You have reached the anonymous limit. Sign up for a free allocation.",
|
||||
"signUp": "Sign up",
|
||||
"stop": "Stop",
|
||||
"topUp": "Top up",
|
||||
"rateLimited": "Rate limited — wait a moment and retry.",
|
||||
"needsKey": "Create an API key and enable it for chat to send as yourself.",
|
||||
"manageKeysLink": "Manage keys"
|
||||
}
|
||||
|
||||
@@ -59,5 +59,10 @@
|
||||
},
|
||||
"footer": {
|
||||
"copyright": "© {{year}} helexa.ai"
|
||||
},
|
||||
"beta": {
|
||||
"tag": "Public beta",
|
||||
"message": "helexa is in open beta — expect rough edges. Your chats stay in your browser.",
|
||||
"dismiss": "Dismiss"
|
||||
}
|
||||
}
|
||||
|
||||
71
helexa.ai/src/i18n/resources/bg/account.json
Normal file
71
helexa.ai/src/i18n/resources/bg/account.json
Normal file
@@ -0,0 +1,71 @@
|
||||
{
|
||||
"login": {
|
||||
"title": "Sign in",
|
||||
"email": "Email",
|
||||
"password": "Password",
|
||||
"submit": "Sign in",
|
||||
"noAccount": "No account? Sign up"
|
||||
},
|
||||
"register": {
|
||||
"title": "Create your account",
|
||||
"email": "Email",
|
||||
"password": "Password",
|
||||
"submit": "Sign up",
|
||||
"haveAccount": "Already have an account? Sign in",
|
||||
"checkEmail": "Almost there — check your email to verify your account."
|
||||
},
|
||||
"verify": {
|
||||
"verifying": "Verifying…",
|
||||
"ok": "Email verified. You can now sign in.",
|
||||
"failed": "This verification link is invalid or has expired.",
|
||||
"toLogin": "Go to sign in"
|
||||
},
|
||||
"reset": {
|
||||
"requestTitle": "Reset your password",
|
||||
"email": "Email",
|
||||
"requestSubmit": "Send reset link",
|
||||
"requestDone": "If that email has an account, a reset link is on its way.",
|
||||
"confirmTitle": "Choose a new password",
|
||||
"newPassword": "New password",
|
||||
"confirmSubmit": "Set password",
|
||||
"ok": "Password updated. You can now sign in."
|
||||
},
|
||||
"dashboard": {
|
||||
"title": "Account",
|
||||
"balance": "Allocation",
|
||||
"total": "Total",
|
||||
"spent": "Spent",
|
||||
"reserved": "Reserved",
|
||||
"remaining": "Remaining",
|
||||
"manageKeys": "Manage API keys",
|
||||
"redeemTitle": "Redeem a top-up code",
|
||||
"redeemPlaceholder": "helexa-topup-…",
|
||||
"redeem": "Redeem",
|
||||
"redeemed": "Code redeemed.",
|
||||
"logout": "Sign out"
|
||||
},
|
||||
"keys": {
|
||||
"title": "API keys",
|
||||
"create": "Create key",
|
||||
"label": "Label",
|
||||
"limitKind": "Limit",
|
||||
"percent": "% of allocation",
|
||||
"hardcap": "Hard cap (tokens)",
|
||||
"value": "Value",
|
||||
"none": "No keys yet.",
|
||||
"createdTitle": "Your new API key",
|
||||
"createdWarn": "Copy it now — you won't see it again.",
|
||||
"copy": "Copy",
|
||||
"copied": "Copied",
|
||||
"archive": "Archive",
|
||||
"save": "Save",
|
||||
"status": "Status",
|
||||
"usage": "Used",
|
||||
"useForChat": "Use for chat on this device",
|
||||
"usedForChat": "Enabled for chat ✓"
|
||||
},
|
||||
"error": {
|
||||
"generic": "Something went wrong.",
|
||||
"unauthorized": "Please sign in again."
|
||||
}
|
||||
}
|
||||
@@ -5,5 +5,17 @@
|
||||
"transcriptPlaceholder": "Тук ще се показва историята на чата. Визуализирай съобщенията от модела и потребителя в превъртащ се контейнер, по желание групирани по ход.",
|
||||
"inputPlaceholder": "Напиши съобщение, за да започнеш разговора…",
|
||||
"send": "Изпрати",
|
||||
"clear": "Изчисти"
|
||||
"clear": "Изчисти",
|
||||
"newChat": "New chat",
|
||||
"newProject": "New project",
|
||||
"newProjectName": "New project",
|
||||
"unsorted": "Unsorted",
|
||||
"emptyState": "Start a conversation. Your history stays in this browser.",
|
||||
"anonBanner": "You have reached the anonymous limit. Sign up for a free allocation.",
|
||||
"signUp": "Sign up",
|
||||
"stop": "Stop",
|
||||
"topUp": "Top up",
|
||||
"rateLimited": "Rate limited — wait a moment and retry.",
|
||||
"needsKey": "Create an API key and enable it for chat to send as yourself.",
|
||||
"manageKeysLink": "Manage keys"
|
||||
}
|
||||
|
||||
@@ -59,5 +59,10 @@
|
||||
},
|
||||
"footer": {
|
||||
"copyright": "© {{year}} helexa.ai"
|
||||
},
|
||||
"beta": {
|
||||
"tag": "Public beta",
|
||||
"message": "helexa is in open beta — expect rough edges. Your chats stay in your browser.",
|
||||
"dismiss": "Dismiss"
|
||||
}
|
||||
}
|
||||
|
||||
71
helexa.ai/src/i18n/resources/da/account.json
Normal file
71
helexa.ai/src/i18n/resources/da/account.json
Normal file
@@ -0,0 +1,71 @@
|
||||
{
|
||||
"login": {
|
||||
"title": "Sign in",
|
||||
"email": "Email",
|
||||
"password": "Password",
|
||||
"submit": "Sign in",
|
||||
"noAccount": "No account? Sign up"
|
||||
},
|
||||
"register": {
|
||||
"title": "Create your account",
|
||||
"email": "Email",
|
||||
"password": "Password",
|
||||
"submit": "Sign up",
|
||||
"haveAccount": "Already have an account? Sign in",
|
||||
"checkEmail": "Almost there — check your email to verify your account."
|
||||
},
|
||||
"verify": {
|
||||
"verifying": "Verifying…",
|
||||
"ok": "Email verified. You can now sign in.",
|
||||
"failed": "This verification link is invalid or has expired.",
|
||||
"toLogin": "Go to sign in"
|
||||
},
|
||||
"reset": {
|
||||
"requestTitle": "Reset your password",
|
||||
"email": "Email",
|
||||
"requestSubmit": "Send reset link",
|
||||
"requestDone": "If that email has an account, a reset link is on its way.",
|
||||
"confirmTitle": "Choose a new password",
|
||||
"newPassword": "New password",
|
||||
"confirmSubmit": "Set password",
|
||||
"ok": "Password updated. You can now sign in."
|
||||
},
|
||||
"dashboard": {
|
||||
"title": "Account",
|
||||
"balance": "Allocation",
|
||||
"total": "Total",
|
||||
"spent": "Spent",
|
||||
"reserved": "Reserved",
|
||||
"remaining": "Remaining",
|
||||
"manageKeys": "Manage API keys",
|
||||
"redeemTitle": "Redeem a top-up code",
|
||||
"redeemPlaceholder": "helexa-topup-…",
|
||||
"redeem": "Redeem",
|
||||
"redeemed": "Code redeemed.",
|
||||
"logout": "Sign out"
|
||||
},
|
||||
"keys": {
|
||||
"title": "API keys",
|
||||
"create": "Create key",
|
||||
"label": "Label",
|
||||
"limitKind": "Limit",
|
||||
"percent": "% of allocation",
|
||||
"hardcap": "Hard cap (tokens)",
|
||||
"value": "Value",
|
||||
"none": "No keys yet.",
|
||||
"createdTitle": "Your new API key",
|
||||
"createdWarn": "Copy it now — you won't see it again.",
|
||||
"copy": "Copy",
|
||||
"copied": "Copied",
|
||||
"archive": "Archive",
|
||||
"save": "Save",
|
||||
"status": "Status",
|
||||
"usage": "Used",
|
||||
"useForChat": "Use for chat on this device",
|
||||
"usedForChat": "Enabled for chat ✓"
|
||||
},
|
||||
"error": {
|
||||
"generic": "Something went wrong.",
|
||||
"unauthorized": "Please sign in again."
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user