Compare commits

..

63 Commits

Author SHA1 Message Date
daeab6be12 fix(neuron): TP unload with Arc-shared pool mutex (#98)
All checks were successful
CI / Format (push) Successful in 8s
CI / Clippy (push) Successful in 2m18s
CI / Test (push) Successful in 6m0s
CI / CUDA type-check (push) Successful in 12m24s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Format (pull_request) Successful in 8s
CI / Clippy (pull_request) Successful in 2m22s
CI / Test (pull_request) Successful in 5m32s
CI / CUDA type-check (pull_request) Successful in 12m40s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
cfg(cuda)-only: the unload path moved the pool out of the (previously
owned) mutex; with TpLoadedModel.pool behind Arc for the engine's
owned guard, recover sole ownership via Arc::try_unwrap —
Arc::try_unwrap(tp) succeeding already implies the engine is idle, so
this is the expected case; the fallback (engine guard mid-release
race) unloads through the lock and lets the pool drop with its last
reference.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-03 17:04:43 +03:00
2a366a1c17 feat(neuron): engine backend seam — TP models batch through the engine (#98)
Some checks failed
CI / Format (push) Successful in 8s
CI / Clippy (push) Successful in 2m18s
CI / Test (push) Successful in 6m8s
CI / CUDA type-check (push) Failing after 12m48s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
The engine loop now drives either serving substrate through a
BackendConfig/ActiveSession seam:

- Single (unchanged semantics): Arc'd clones of the LoadedModel
  pieces; active phase holds inference_lock.
- Tp: a Weak<TpLoadedModel>, upgraded only while the batch is active —
  an idle engine never keeps an unloaded model alive, and an unload
  mid-batch completes when the batch drains. The active phase holds
  the pool mutex (TpLoadedModel.pool is now Arc<Mutex<WorkerPool>> so
  the guard can be owned), which is exactly the serialization the
  per-request TP path used — vision and non-streaming requests wait
  their turn as before.

Backend ops (op_prefill/op_step/op_assemble/op_extract/
op_snapshot_seq/op_drop_snap) dispatch per substrate: the TP arms
reuse restore_or_clear_tp/chunked_prefill_tp/store_prefix_snapshot_tp
verbatim and mint pool-wide snapshot ids from next_snapshot_id so all
ranks key identically.

Wiring: TpLoadedModel.engine is a OnceLock set right after the Arc is
built (the Weak needs the Arc to exist); spawn gated on the same
conditions as single-GPU (max_in_flight > 1, snapshot-capable,
NEURON_BATCHING != 0 — default config keeps everything off).
inference_tp_stream routes text streams through the engine when
present.

Engine gold test (single backend) unchanged and green; the TP arms
are cfg(cuda) — branch CUDA type-check + live smoke on beast are the
gates.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-03 16:46:00 +03:00
4c3a2735b0 feat(neuron): TP rank plumbing for batched decode (#98)
All checks were successful
CI / Format (push) Successful in 8s
CI / Clippy (push) Successful in 2m19s
CI / Test (push) Successful in 5m42s
CI / CUDA type-check (push) Successful in 12m53s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
TP mirror of the batched-decode primitives, spanning all three tiers:

- arch: TpQwen3_5Model/ForCausalLM forward_batch_decode +
  batch_decode_mask (same per-row batch_cos_sin rope path and
  one-gap-per-row mask as single-GPU; the shared attention_context
  core already handles the batched shapes).
- RPC: GenerateStepBatch (per-row tokens + broadcast geometry),
  AssembleKvBatch (pool-minted snapshot ids), ExtractKvRows
  (pre-minted ids so every rank keys identically); responses
  KvBatchAssembled { padded_len } / KvRowsExtracted { bytes }. No
  protocol version to bump — the RPC evolves by adding variants
  (subprocess is spawned from the same binary; unknown ops fail as
  bad_request).
- subprocess worker: handlers derive per-row positions + padding mask
  locally and discard logits (the NCCL collectives are the point),
  and run assemble_batch/extract_row against the rank's own shard
  snapshots — the snapshot types and geometry are identical across
  ranks, only head counts shard.
- leader: Job::{TpForwardLogitsBatch, TpAssembleKvBatch,
  TpExtractKvRows} + dispatch handlers against the TP slab + handle
  wrappers.
- pool: WorkerPool::{generate_step_batch, assemble_kv_batch,
  extract_kv_rows} with the same fan-out → leader → always-drain
  shape as generate_step (incl. the #17 watchdog on the batched
  step); assembly asserts every rank agrees on padded_len.

Nearly all of this is cfg(cuda) — the branch CUDA type-check is the
real gate. Engine backend seam (TpLoadedModel wiring) follows in the
next commit.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-03 16:23:50 +03:00
7454ec7744 Merge pull request 'feat(neuron): continuous batching — batched decode + engine (#98, slices 1–3b)' (#121) from feat/98-batched-decode into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 10s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m2s
build-prerelease / Test (push) Successful in 6m24s
build-prerelease / Build neuron-blackwell (push) Successful in 8m47s
build-prerelease / Build helexa-bench binary (push) Has been skipped
build-prerelease / Build helexa-upstream binary (push) Has been skipped
build-prerelease / Package helexa-bench RPM (push) Has been skipped
build-prerelease / Package helexa-upstream RPM (push) Has been skipped
build-prerelease / Build neuron-ampere (push) Successful in 2m6s
build-prerelease / Build neuron-ada (push) Successful in 2m8s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 42s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 48s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 2m53s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 20s
2026-07-03 12:39:08 +00:00
d3ef9513f8 fix(neuron): as_deref for Arc'd prefix_cache in cuda-only worker path (#98)
All checks were successful
CI / Format (push) Successful in 7s
CI / Clippy (push) Successful in 2m19s
CI / Test (push) Successful in 6m0s
CI / CUDA type-check (push) Successful in 12m36s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Format (pull_request) Successful in 7s
CI / Clippy (pull_request) Successful in 3m4s
CI / Test (pull_request) Successful in 6m16s
CI / CUDA type-check (pull_request) Successful in 12m38s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
candle.rs:2441 sits behind cfg(cuda) — the CPU build can't see it, and
the CUDA type-check caught the Option<&Arc<Mutex<…>>> vs
Option<&Mutex<…>> mismatch introduced by the LoadedModel Arc change.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-03 15:24:00 +03:00
ec7a5b750a feat(neuron): lockstep batched decode engine (#98)
Some checks failed
CI / Format (push) Successful in 7s
CI / Clippy (push) Successful in 2m14s
CI / Test (push) Successful in 5m25s
CI / CUDA type-check (push) Failing after 12m25s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
Slice 3b: per-model engine task multiplexing concurrent text chat
streams through one (B,1) forward per decode step, replacing the
per-request inference_lock serialization when [admission]
max_in_flight > 1 on a snapshot-capable worker-path model.

- harness/engine.rs: engine loop (join → B=1 prefill via the existing
  chunked-prefill + prefix-cache paths → snapshot → rebatch via
  ExtractKvRows/AssembleKvBatch; lockstep ForwardLogitsBatch steps;
  per-slot CPU sampling with per-slot LogitsProcessor + repeat-penalty
  history). Per-slot router tasks own the incremental detokenizer and
  the reasoning/tool-call state machine (same logic as route_token!)
  and emit InferenceEvents — DecodeStream's tokenizer borrow lives
  inside the router's async block, and slow consumers decouple from
  the lockstep loop.
- The engine holds the model's inference_lock while it has active
  slots, so vision and non-streaming requests (which keep the direct
  path) still serialize safely against the batch; the lock releases
  when the batch drains.
- Fatal worker errors fail all active slots, mark the model poisoned
  on device-fault classification, and stop the engine (later submits
  fail fast).
- LoadedModel: poisoned/inference_lock/prefix_cache/prefill_rate move
  behind Arc (shared with the engine without keeping the model alive);
  new engine: Option<EngineHandle> field. Engine activates only when
  max_in_flight > 1 (config default 1 = existing behaviour everywhere)
  and NEURON_BATCHING != 0 (kill switch).
- Prefill helpers (restore_or_clear/chunked_prefill/store_prefix/
  emit_delta) ungated from cfg(cuda) so the engine and its CPU tests
  compile everywhere.

Gold test: three greedy requests submitted concurrently (ragged
prompts, mid-flight joins, rebatching) produce byte-identical text,
token counts, and finish reasons to the same requests run
sequentially through the same engine, on the tiny fixture.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-03 15:03:43 +03:00
7f2fcc3527 feat(neuron): batch-row extraction for rebatch on join/leave (#98)
All checks were successful
CI / Format (push) Successful in 8s
CI / Clippy (push) Successful in 2m27s
CI / Test (push) Successful in 5m41s
CI / CUDA type-check (push) Successful in 12m3s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
Slice 3a: extract_row slices one row of a batched cache state back
into a contiguous single-sequence snapshot — attention KV keeps the
prefix [0, prefix_len) plus the lockstep decode columns [padded_len,
padded_len + steps) and drops the padding gap; GDN rows are sliced
whole and deep-copied. A join/leave rebatches by extracting every
surviving row and re-running assemble_batch, preserving the
one-gap-per-row invariant without any scatter kernels.

Worker side: Job::ExtractKvRows captures the live state once and
stores each extracted row in the snapshot slab; composes with
AssembleKvBatch / DropKvSnapshot for the full rebatch.

Tests: arch-level round-trip (batched decode → extract → solo B=1
continuation matches pure-sequential) and a worker end-to-end leave
path (3-row batch → extract 2 survivors → re-assemble → continued
lockstep decode matches the sequential reference).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-03 14:46:49 +03:00
9ac1a1586c feat(neuron): device-worker jobs for batched decode (#98)
Some checks failed
CI / Test (push) Waiting to run
CI / CUDA type-check (push) Waiting to run
CI / Format (push) Successful in 8s
CI / Clippy (push) Successful in 2m17s
CI / Build cortex SRPM (push) Has been cancelled
CI / Build neuron SRPM (push) Has been cancelled
CI / Publish cortex to COPR (push) Has been cancelled
CI / Publish neuron to COPR (push) Has been cancelled
CI / Bump version in source (push) Has been cancelled
Slice 2: AssembleKvBatch + ForwardLogitsBatch route the batched-decode
primitives through the per-device worker thread.

- AssembleKvBatch assembles stored per-sequence KV snapshots into one
  batched (B,…) live state via snapshot::assemble_batch and installs
  it through the existing restore path; replies the padded KV length.
- ForwardLogitsBatch runs one lockstep decode step: builds the (B,1)
  input on the worker's device, derives per-row positions and the
  padding mask from prefix_lens/padded_len/step, and replies one CPU
  [vocab] logits row per batch row — same tensors-never-escape
  contract as ForwardLogits.
- ModelArch::forward_batch_decode / batch_decode_mask dispatch (qwen3_5
  only; other archs error as defence in depth).
- Dense safetensors loads now pick f32 on non-CUDA devices: candle's
  CPU backend has no bf16 matmul, so the CPU fallback (and the CPU
  test worker) was broken at first forward under the hardcoded BF16.

End-to-end test: tiny qwen3_next fixture loaded through the worker,
three ragged sequences prefilled+snapshotted, assembled, and every
ForwardLogitsBatch row checked against the sequential ForwardLogits
reference at each step.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-03 14:42:06 +03:00
05b6c70a3a feat(neuron): batched lockstep decode groundwork for qwen3_5 (#98)
Some checks failed
CI / CUDA type-check (push) Waiting to run
CI / Format (push) Successful in 8s
CI / Clippy (push) Successful in 2m23s
CI / Test (push) Successful in 5m51s
CI / Build cortex SRPM (push) Has been cancelled
CI / Build neuron SRPM (push) Has been cancelled
CI / Publish cortex to COPR (push) Has been cancelled
CI / Publish neuron to COPR (push) Has been cancelled
CI / Bump version in source (push) Has been cancelled
Slice 1 of continuous batching (F4d): arch-level batched decode with
per-row positions, assembled from per-sequence prefix snapshots.

- rope: batch_cos_sin gathers cos/sin at arbitrary per-row positions
  ((B,1,half)); apply_cos_sin dispatches on cos rank — rank-3 takes a
  local GLM rotate-half broadcast path (rope_slow only broadcasts one
  position table across the batch).
- snapshot: assemble_batch cats per-sequence KvCacheSnapshots into one
  (B,…) batched state — attention K/V right-padded along the sequence
  axis (keys are stored post-RoPE, so padding is position-inert), GDN
  conv/recurrent states cat on dim 0. Text-only: rope_delta must be 0.
- model: forward_batch_decode((B,1), positions, mask) lockstep decode
  step returning (B,1,vocab); batch_decode_mask builds the (B,1,1,S)
  additive mask covering each row's padding gap [prefix_len,
  padded_len), None when lengths are uniform.

Gold test: three ragged sequences decoded in one lockstep batch match
the same sequences decoded sequentially, hidden-for-hidden at every
step, on the tiny CPU fixture. Plus rope gather/uniform-equivalence
and mask-geometry units.

No serving-path changes yet — the engine loop (slice 3) and worker
job plumbing (slice 2) build on these primitives.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-03 14:34:31 +03:00
9c037c5a0b Merge pull request 'perf(neuron): route only prefill through FlashAttention (#95)' (#120) from fix/95-flash-prefill-only into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 11s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Package helexa-bench RPM (push) Has been skipped
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m27s
build-prerelease / Test (push) Successful in 6m43s
build-prerelease / Build neuron-blackwell (push) Successful in 9m0s
build-prerelease / Build helexa-upstream binary (push) Has been skipped
build-prerelease / Package helexa-upstream RPM (push) Has been skipped
build-prerelease / Build neuron-ampere (push) Successful in 2m11s
build-prerelease / Build neuron-ada (push) Successful in 2m12s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 47s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 48s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 2m54s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 21s
2026-07-03 09:29:51 +00:00
cb6ddbbd98 perf(neuron): route only prefill through FlashAttention (#95)
All checks were successful
CI / Format (push) Successful in 9s
CI / Clippy (push) Successful in 2m17s
CI / Test (push) Successful in 5m55s
CI / CUDA type-check (push) Successful in 12m39s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Format (pull_request) Successful in 7s
CI / Clippy (pull_request) Successful in 2m21s
CI / Test (pull_request) Successful in 6m59s
CI / CUDA type-check (pull_request) Successful in 12m39s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
On-beast A/B (27B, 30,419-token prompt, 2x RTX 5090, e824ea3):
prefill 24.8s eager → 22.1s flash (~11%, diluted by the hybrid's
~25% full-attention layer ratio), but decode REGRESSED ~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 byte-identical in both modes (parity gate passed,
incl. chunked-prefill causal alignment).

Restrict the flash dispatch to q_len > 1: prefill keeps the win,
decode keeps the eager path.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-03 12:15:03 +03:00
e824ea3274 Merge pull request 'perf(neuron): FlashAttention on the qwen3_5 full-attention layers (#95)' (#119) from feat/95-flash-attn into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 12s
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m25s
build-prerelease / Test (push) Successful in 7m24s
build-prerelease / Build neuron-blackwell (push) Successful in 8m52s
build-prerelease / Build helexa-bench binary (push) Successful in 1m32s
build-prerelease / Build cortex binary (push) Successful in 1m54s
build-prerelease / Build neuron-ada (push) Successful in 2m7s
build-prerelease / Build neuron-ampere (push) Successful in 2m8s
build-prerelease / Build helexa-upstream binary (push) Successful in 2m11s
build-prerelease / Package helexa-bench RPM (push) Successful in 23s
build-prerelease / Package cortex RPM (push) Successful in 25s
build-prerelease / Package helexa-upstream RPM (push) Successful in 22s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 43s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 48s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 2m51s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 22s
2026-07-03 08:47:24 +00:00
9139b07d71 fix(neuron): candle-flash-attn must be a direct optional dependency (#95)
All checks were successful
CI / Format (push) Successful in 7s
CI / Clippy (push) Successful in 2m18s
CI / Test (push) Successful in 5m38s
CI / CUDA type-check (push) Successful in 12m35s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Format (pull_request) Successful in 8s
CI / Clippy (pull_request) Successful in 2m14s
CI / Test (pull_request) Successful in 6m59s
CI / CUDA type-check (pull_request) Successful in 12m35s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
The flash-attn feature only forwarded to candle-transformers, which
enables flash inside ITS models — candle_flash_attn was not nameable
from neuron's own attention core (E0433 in the CUDA type-check).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-03 11:31:28 +03:00
4762537225 perf(neuron): FlashAttention on the qwen3_5 full-attention layers (#95)
Some checks failed
CI / Format (push) Successful in 13s
CI / Clippy (push) Successful in 2m21s
CI / Test (push) Successful in 5m27s
CI / CUDA type-check (push) Failing after 12m37s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
F4a. The hybrid's full-attention layers computed eager
matmul→softmax→matmul with materialised GQA-repeated K/V and O(L²)
score/mask tensors — the dominant term in prefill at agentic context
sizes (measured: ~30 s TTFT at 16k tokens, the trigger for the
opencode timeout spiral).

- attention_context(): shared attention core for the single-GPU and
  TP paths. With the flash-attn feature on a CUDA device in f16/bf16
  it dispatches to candle_flash_attn::flash_attn — GQA native (no
  repeated-K/V materialisation), causality as a kernel flag, no score
  matrix. Everything else (CPU, f32, feature off) keeps the eager
  fallback, byte-identical to the previous math — the qwen3_next HF
  parity fixture pins this in CI.
- Chunked prefill correctness rides flash-attention v2.1+ causal
  semantics (bottom-right alignment for seqlen_q != seqlen_k); the
  invariant and the caller's mask contract are documented on the
  helper. On-beast flash-vs-eager greedy parity is the merge gate.
- NEURON_FLASH_ATTN=0 forces eager at runtime — A/B lever + rollback.
- Feature enabled for the blackwell flavour build only (beast carries
  the agentic prefill pain; ada/ampere follow once the win is
  measured). CUDA type-check now covers cuda,flash-attn; timeout
  raised for the first uncached kernel compile.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-03 11:14:47 +03:00
9b635bfe5c Merge pull request 'fix(helexa-bench): thinking-model streams misread as dead + impossible decode rates (#117)' (#118) from fix/117-bench-thinking-metrics into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 13s
build-prerelease / Build neuron-blackwell (push) Has been skipped
build-prerelease / Build neuron-ampere (push) Has been skipped
build-prerelease / Build neuron-ada (push) Has been skipped
build-prerelease / Package helexa-neuron-ada RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been skipped
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Successful in 1m32s
build-prerelease / Package helexa-bench RPM (push) Successful in 21s
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m25s
build-prerelease / Test (push) Successful in 7m5s
build-prerelease / Build helexa-upstream binary (push) Has been skipped
build-prerelease / Package helexa-upstream RPM (push) Has been skipped
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 13s
2026-07-02 20:10:21 +00:00
28aa06b768 fix(helexa-bench): thinking-model streams misread as dead + impossible decode rates (#117)
All checks were successful
CI / Format (push) Successful in 7s
CI / CUDA type-check (push) Successful in 1m37s
CI / Clippy (push) Successful in 2m17s
CI / Test (push) Successful in 5m49s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Format (pull_request) Successful in 7s
CI / CUDA type-check (pull_request) Successful in 1m35s
CI / Clippy (pull_request) Successful in 2m26s
CI / Test (pull_request) Successful in 6m50s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
The scenario stream reader only counted `content` deltas: reasoning
models (Qwen3-Next-Thinking; Qwen3 family with thinking on) stream
reasoning_content first — for their whole budget on small-max_tokens
scenarios — so the harness saw a dead stream ("no content chunks
received"), failed the cell, and re-retried it every sweep forever
(measured=0 skipped=15 failed=27, for hours, live 2026-07-02).
Meanwhile decode_tps divided usage.completion_tokens (reasoning
INCLUSIVE) by the visible-content-only chunk window — reporting
244 tok/s for a 1.7B on a 3060 whose true rate is ~39.

- Liveness now counts any generated delta (content or
  reasoning_content); a cell fails only when the stream is truly
  dead. ttft_s = first generated delta — unchanged semantics for
  non-thinking models (their first delta is content) and honest
  engine latency for thinking models.
- decode_tps prefers the server-measured prefill/decode split (#85),
  which the reader was already capturing but not using; the client
  fallback divides the CHUNK count by the chunk window — one frame,
  never mixed.
- Captured artifacts remain visible-content-only.

Cells are keyed by the target neuron's build SHA, so shipping this
mid-collection does not re-key the in-progress F3 rotation.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-02 23:03:15 +03:00
fc95faaed6 Merge pull request 'fix(neuron): admission queue leaked slots when clients cancelled mid-wait' (#115) from fix/admission-cancel-leak into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 12s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Has been skipped
build-prerelease / Package helexa-bench RPM (push) Has been skipped
build-prerelease / Build neuron-blackwell (push) Successful in 1m18s
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m22s
build-prerelease / Test (push) Successful in 7m20s
build-prerelease / Build helexa-upstream binary (push) Has been skipped
build-prerelease / Package helexa-upstream RPM (push) Has been skipped
build-prerelease / Build neuron-ada (push) Successful in 1m57s
build-prerelease / Build neuron-ampere (push) Successful in 2m7s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 46s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 44s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 46s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 16s
2026-07-02 13:36:11 +00:00
2df829ed73 fix(neuron): admission queue leaked slots when clients cancelled mid-wait
All checks were successful
CI / Format (push) Successful in 7s
CI / CUDA type-check (push) Successful in 1m39s
CI / Clippy (push) Successful in 2m25s
CI / Test (push) Successful in 5m48s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Format (pull_request) Successful in 7s
CI / CUDA type-check (pull_request) Successful in 1m35s
CI / Clippy (pull_request) Successful in 2m15s
CI / Test (pull_request) Successful in 7m10s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
AdmissionController::enter() incremented the pending + per-principal
counts, then awaited the in-flight semaphore. That await is where a
client disconnect lands: axum drops the request future, and neither
the success path nor the timeout branch's rollback ever ran — every
abandoned wait permanently leaked one pending slot and one
per-principal count.

Under a client retry storm (opencode re-sending after aborting), the
leaks ratchet pending to max_pending, after which every request gets
an instant QueueFull 429; the leaked per-principal counts alone can
pin a principal at its fair-share cap forever. Observed live
2026-07-02 during the implementation-eval runs: queue_depth stuck at 1
on an idle model, ~70ms 429s at the gateway, 36 aborted attempts and
507k prompt tokens burned against 38k generated in one 40-minute
window.

The reservation is now a RAII PendingReservation taken before the
await: admitted permits carry it for the request lifetime, and any
other exit — timeout, or the future being dropped mid-queue — rolls
the counts back in Drop. Regression test aborts two queued waiters
and asserts both the queue depth and the principal count recover.

Follow-up (separate issue): abandon chunked prefill early when the
streaming client is gone — needs a clean sentinel distinct from the
device-fault path so it never trips poison/auto-recovery.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-02 16:28:50 +03:00
f611e3479b Merge pull request 'fix(neuron): truncated-mid-think non-streaming responses are all reasoning (#112)' (#113) from fix/112-prompt-opened-reasoning into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 12s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Has been skipped
build-prerelease / Package helexa-bench RPM (push) Has been skipped
build-prerelease / Build neuron-blackwell (push) Successful in 1m19s
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m37s
build-prerelease / Test (push) Successful in 7m30s
build-prerelease / Build helexa-upstream binary (push) Has been skipped
build-prerelease / Package helexa-upstream RPM (push) Has been skipped
build-prerelease / Build neuron-ada (push) Successful in 1m58s
build-prerelease / Build neuron-ampere (push) Successful in 2m5s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 45s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 44s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 48s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 14s
2026-07-02 09:04:50 +00:00
41ec05e40c fix(neuron): truncated-mid-think non-streaming responses are all reasoning (#112)
All checks were successful
CI / Format (push) Successful in 9s
CI / CUDA type-check (push) Successful in 1m37s
CI / Clippy (push) Successful in 2m16s
CI / Test (push) Successful in 5m29s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Format (pull_request) Successful in 7s
CI / CUDA type-check (pull_request) Successful in 1m36s
CI / Clippy (pull_request) Successful in 2m20s
CI / Test (pull_request) Successful in 7m19s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
When the chat template force-opens the think block in the generation
prompt (Qwen3-Next-80B-A3B-Thinking ends with
'<|im_start|>assistant\n<think>\n') and the generation exhausts
max_tokens before emitting </think>, the non-streaming
split_off_reasoning had no close-marker anchor and returned the whole
chain-of-thought as content. Thread the existing
prompt_opens_reasoning result (already used by the streaming state
machine) into the split: no close marker + prompt-opened block means
content is empty and every generated token counts as reasoning.

Closes #112

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-02 11:57:52 +03:00
3242cce39b chore(bench): capability probes need 4096 tokens for the Thinking variant (#94)
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 11s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Has been skipped
build-prerelease / Package helexa-bench RPM (push) Has been skipped
build-prerelease / Build helexa-upstream binary (push) Has been skipped
build-prerelease / Package helexa-upstream RPM (push) Has been skipped
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m22s
build-prerelease / Test (push) Successful in 6m43s
build-prerelease / Build neuron-blackwell (push) Has been skipped
build-prerelease / Build neuron-ampere (push) Has been skipped
build-prerelease / Build neuron-ada (push) Has been skipped
build-prerelease / Package helexa-neuron-ada RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been skipped
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been skipped
The 80B-A3B-Thinking template force-opens a think block, so its
token budget goes to reasoning first — at 2048 the post-think answer
would be empty or truncated (the 27B and Coder-Next already ran to
~1930 tokens). The next neuron build SHA re-collects every model's
capability cells at 4096, giving F3 a consistent comparison set.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-02 10:47:18 +03:00
35041f11de Merge pull request 'fix(neuron): route the fused MoE block in the activation dtype (#92)' (#111) from fix/92-fused-router-dtype into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 12s
build-prerelease / Build neuron-blackwell (push) Successful in 1m29s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Has been skipped
build-prerelease / Package helexa-bench RPM (push) Has been skipped
build-prerelease / Build neuron-ampere (push) Successful in 2m10s
build-prerelease / Build neuron-ada (push) Successful in 2m10s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m40s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 47s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 43s
build-prerelease / Test (push) Successful in 5m43s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 1m3s
build-prerelease / Build helexa-upstream binary (push) Has been skipped
build-prerelease / Package helexa-upstream RPM (push) Has been skipped
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 18s
2026-07-02 04:34:40 +00:00
40fecb2cf7 fix(neuron): route the fused MoE block in the activation dtype (#92)
All checks were successful
CI / Format (push) Successful in 7s
CI / CUDA type-check (push) Successful in 1m38s
CI / Clippy (push) Successful in 2m17s
CI / Test (push) Successful in 5m43s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Format (pull_request) Successful in 7s
CI / CUDA type-check (pull_request) Successful in 1m39s
CI / Clippy (pull_request) Successful in 2m20s
CI / Test (pull_request) Successful in 6m25s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
The fused path fed f32 activations into the replicated bf16 router —
a dtype-mismatched matmul that failed the first live fused request on
the 80B and (correctly) tripped the poison/auto-recovery machinery.
Route in the original activation dtype like route_scatter does; only
the softmax and downstream weights are f32 (which the grouped-GEMM
kernel requires for topk_weights anyway).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-02 07:27:43 +03:00
8d50918a1a Merge pull request 'perf(neuron): fused grouped-GEMM dispatch for the qwen3_next MoE block (#92)' (#110) from feat/92-moe-fused into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 13s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Build neuron-blackwell (push) Successful in 1m23s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m47s
build-prerelease / Test (push) Successful in 6m15s
build-prerelease / Build helexa-bench binary (push) Has been skipped
build-prerelease / Build helexa-upstream binary (push) Has been skipped
build-prerelease / Package helexa-bench RPM (push) Has been skipped
build-prerelease / Package helexa-upstream RPM (push) Has been skipped
build-prerelease / Build neuron-ada (push) Successful in 2m2s
build-prerelease / Build neuron-ampere (push) Successful in 2m11s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 44s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 48s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 50s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 15s
2026-07-02 03:53:11 +00:00
29a7054c23 perf(neuron): fused grouped-GEMM dispatch for the qwen3_next MoE block (#92)
All checks were successful
CI / Format (push) Successful in 8s
CI / CUDA type-check (push) Successful in 1m39s
CI / Clippy (push) Successful in 2m17s
CI / Test (push) Successful in 5m33s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Format (pull_request) Successful in 7s
CI / CUDA type-check (pull_request) Successful in 1m38s
CI / Clippy (pull_request) Successful in 2m33s
CI / Test (pull_request) Successful in 5m56s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
F1 slice 4 — replaces the correctness-first scatter loop on CUDA:

- TpExpertStore::Fused holds each projection as ONE stacked per-rank
  QTensor ([E, out/ws, in]) built at load: read every expert's rank
  slice, stack, ISQ the stack in a single parallel pass per
  projection (~0.5 GB transient bf16 per stack at 80B dims).
- forward_fused ports candle-transformers' FusedMoeGGUF::forward onto
  the per-rank stacks via candle-nn's moe_gemm_gguf grouped-GEMM
  kernels: routing, index sort, and all expert GEMMs stay on-device —
  three kernel launches per layer regardless of top-k, no GPU→CPU
  routing sync. gate/up run unweighted (tokens×topk rows); the down
  GEMM folds the routing weights in-kernel; the (tokens, topk, hidden)
  view sums over topk. Output is the rank's partial; the shared expert
  and single block-end AllReduce are unchanged.
- Store chosen at load: CUDA + ISQ with a kernel-supported GGML dtype
  (q2k-q6k, q8_0) + GGUF block alignment on both GEMM K dims;
  NEURON_MOE_FUSED=0 forces scatter as escape hatch and A/B lever.
  CPU/non-cuda builds keep scatter (the CI parity anchor).

Motivation: the 80B-A3B live smoke measured 4.3 tok/s decode on the
scatter path (host routing sync + ~30 tiny GEMV launches per layer
per token) vs ~27 tok/s for the dense 27B. Target: close the ~6x gap.
On-beast A/B (fused vs NEURON_MOE_FUSED=0 scatter, greedy-token
equivalence + decode tok/s) follows this merge.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-02 06:46:06 +03:00
7489a49d24 fix(ci): render RPM release stamps in UTC regardless of committer TZ
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 12s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m21s
build-prerelease / Build neuron-blackwell (push) Successful in 1m40s
build-prerelease / Build helexa-bench binary (push) Successful in 1m36s
build-prerelease / Build cortex binary (push) Successful in 2m2s
build-prerelease / Build neuron-ampere (push) Successful in 2m25s
build-prerelease / Build neuron-ada (push) Successful in 2m27s
build-prerelease / Build helexa-upstream binary (push) Successful in 2m23s
build-prerelease / Test (push) Successful in 5m32s
build-prerelease / Package helexa-bench RPM (push) Successful in 20s
build-prerelease / Package helexa-upstream RPM (push) Successful in 20s
build-prerelease / Package cortex RPM (push) Successful in 25s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 44s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 44s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 49s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 18s
git log --date=format: renders each commit's timestamp in that
commit's OWN recorded timezone. Local commits (UTC+3) and Gitea
server-side merge commits (UTC) therefore produced non-monotonic
release stamps: the 2026-07-01 bench-config commit stamped 235959
(local) while the later F1 merges stamped 2123xx/2154xx (UTC), so RPM
EVR comparison ranked the older build newest. The deploy gate
correctly detected the mismatch (it compares by buildTime) but dnf
upgrade refused the "downgrade" and silently restarted the stale
binary fleet-wide.

TZ=UTC0 + --date=format-local: pins the stamp to UTC for every
commit. Hosts currently pinned on the 235959 build self-heal on the
first build after 2026-07-02 00:00 UTC; beast is downgraded manually
(it needs the F1 #92 build for the 80B bring-up).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-02 01:29:46 +03:00
341f97f705 Merge pull request 'feat(neuron): TP expert sharding for the qwen3_next MoE block (#92)' (#109) from feat/92-qwen3_5-moe into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 12s
build-prerelease / Build neuron-blackwell (push) Successful in 1m26s
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m17s
build-prerelease / Test (push) Successful in 9m12s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Has been skipped
build-prerelease / Package helexa-bench RPM (push) Has been skipped
build-prerelease / Build helexa-upstream binary (push) Has been skipped
build-prerelease / Package helexa-upstream RPM (push) Has been skipped
build-prerelease / Build neuron-ada (push) Successful in 7m24s
build-prerelease / Build neuron-ampere (push) Successful in 7m39s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 42s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 45s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 47s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 15s
2026-07-01 21:54:11 +00:00
fc56eafbf7 fix(ci): deploy LLM probe retries on 429/503 instead of hard-failing
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 12s
build-prerelease / Build neuron-blackwell (push) Has been skipped
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m51s
build-prerelease / Build neuron-ampere (push) Has been skipped
build-prerelease / Build neuron-ada (push) Has been skipped
build-prerelease / Package helexa-neuron-ada RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been skipped
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Package helexa-bench RPM (push) Has been skipped
build-prerelease / Build helexa-upstream binary (push) Has been skipped
build-prerelease / Package helexa-upstream RPM (push) Has been skipped
build-prerelease / Test (push) Successful in 5m54s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been skipped
The post-deploy probe races real traffic by construction: a deploy
publishes a new build SHA, which is exactly what triggers helexa-bench
to re-sweep every scenario against the restarted neuron — and the
capability probes (#91) hold the batch-1 model for minutes of
generation. Admission control answers concurrent requests with 429/503
+ Retry-After (#53/#63); the old single-shot `curl -fsS` treated that
honest backpressure as deploy failure (run 694, deploy-neurons/benjy).

The probe now retries on 429/503 honoring Retry-After (min 5s,
default 15s) within a ~6-minute budget; any other error or a wrong
response still fails immediately.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-02 00:44:14 +03:00
a157c8fa27 fix(neuron): gate TP MoE load tests to non-cuda builds (#92)
All checks were successful
CI / Format (push) Successful in 8s
CI / CUDA type-check (push) Successful in 1m41s
CI / Clippy (push) Successful in 2m16s
CI / Test (push) Successful in 5m58s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Format (pull_request) Successful in 15s
CI / CUDA type-check (pull_request) Successful in 1m37s
CI / Clippy (pull_request) Successful in 2m22s
CI / Test (pull_request) Successful in 9m35s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
The CUDA type-check compiles test targets too, where
TpQwen3_5MoeBlock::load takes an NCCL Comm the tests cannot
construct (E0061). The CPU Test job runs the partial-sum parity
test; the CUDA job type-checks the cuda load/reduce variants.
Also drops an unused binding only the cuda test build flagged.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-02 00:42:45 +03:00
201cc54a7b feat(neuron): TP expert sharding for the qwen3_next MoE block (#92)
Some checks failed
CI / Format (push) Successful in 6s
CI / CUDA type-check (push) Failing after 1m31s
CI / Clippy (push) Successful in 2m19s
CI / Test (push) Successful in 5m35s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
F1 slice 3 — the MoE FFN under tensor parallelism:

- TpQwen3_5MoeBlock: per-expert sharding along moe_intermediate_size
  (gate/up column-sliced, down input-sliced) with forward_partial
  outputs, deliberately NOT composed from RowParallelLinear — that
  embeds an AllReduce per call, and top-10 routing would mean ten
  collectives per layer. Partial sums from routed experts and the
  shared expert add linearly, so the whole block does exactly ONE
  AllReduce at the end, preserving the one-reduce-per-FFN pattern.
- Router + shared_expert_gate are replicated (tiny): every rank
  computes identical top-k assignments with zero communication, via
  the route_scatter helper now shared with the single-GPU block.
- TpMlpKind { Dense, Moe } decoder dispatch on layer_uses_moe,
  mirroring the single-GPU MlpKind; per-rank ISQ applies to expert
  slices through the existing MaybeQuantLinear path.
- TP fused-checkpoint loading: K- and V-heads shard uniformly
  together, so each rank owns a contiguous span of whole k-head
  groups of the fused in_proj_qkvz/in_proj_ba tensors — a uniform
  dim-0 shard is exactly the rank's groups, then the per-rank
  de-interleave (split_fused_qkvz/ba with per-rank head counts)
  restores the contiguous [Q|K|V]+Z / B+A layout. Auto-detected;
  Qwen3.6's separate-tensor path untouched.
- text_weight_prefix threaded through the TP model load (qwen3_next
  keeps the text core at model.*); the slice-1 TP MoE guard removed.

Validation (CPU, runs in CI): with the block-end AllReduce elided,
rank-0 + rank-1 partial outputs at world_size=2 sum exactly to the
single-GPU block's output (which itself has exact HF parity from
slice 2) — pinning expert slicing, replicated routing, and
shared-expert partial scaling. A second test pins the per-rank
group-span de-interleave against row-slices of the full split.
The cfg(cuda) load/reduce variants are validated by the CUDA
type-check in CI; live TP-2 smoke lands with F2's checkpoint.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-02 00:34:07 +03:00
781d92d9a3 Merge pull request 'feat(neuron): qwen3_next config + MoE FFN block, single-GPU path with exact HF parity (#92)' (#108) from feat/92-qwen3_5-moe into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 14s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Package helexa-bench RPM (push) Has been skipped
build-prerelease / Build neuron-blackwell (push) Successful in 1m27s
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m31s
build-prerelease / Test (push) Successful in 6m52s
build-prerelease / Build helexa-upstream binary (push) Has been skipped
build-prerelease / Package helexa-upstream RPM (push) Has been skipped
build-prerelease / Build neuron-ada (push) Successful in 2m4s
build-prerelease / Build neuron-ampere (push) Successful in 2m10s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 46s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 50s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 58s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 16s
2026-07-01 21:23:23 +00:00
9bf13f09dd feat(neuron): qwen3_next MoE FFN block, single-GPU path + HF parity (#92)
All checks were successful
CI / Format (push) Successful in 8s
CI / CUDA type-check (push) Successful in 1m38s
CI / Clippy (push) Successful in 2m17s
CI / Test (push) Successful in 5m17s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Format (pull_request) Successful in 7s
CI / CUDA type-check (pull_request) Successful in 1m38s
CI / Clippy (pull_request) Successful in 2m15s
CI / Test (pull_request) Successful in 6m54s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
F1 slice 2 — the MoE block itself, CPU/single-GPU:

- arch/qwen3_5/moe.rs: Qwen3_5MoeBlock — top-k router (upstream
  softmax-then-topk order, renorm iff norm_topk_prob), per-expert
  SwiGLU (reusing Qwen3_5MLP at moe_intermediate_size), and the
  always-on shared expert mixed via sigmoid(shared_expert_gate).
  Correctness-first host-side scatter dispatch; the fused grouped-GEMM
  path is slice 4 behind the same forward signature.
- decoder.rs: MlpKind { Dense, Moe } dispatch on layer_uses_moe,
  mirroring AttentionKind.
- linear_attn.rs: fused-checkpoint support — qwen3_next stores
  in_proj_qkvz / in_proj_ba interleaved per key-head group (upstream
  fix_query_key_value_ordering layout); split_fused_qkvz/ba
  de-interleave once at load into the contiguous [Q|K|V] + Z / B + A
  layout the forward path (incl. the conv channels) already uses.
  Auto-detected via contains_tensor, so Qwen3.6 checkpoints are
  untouched.
- mod.rs: text_weight_prefix() — qwen3_next checkpoints put the text
  core at `model.*`, Qwen3.6 at `model.language_model.*`; the slice-1
  single-GPU MoE guard is removed (TP guard stays until slice 3).

Validation:
- qwen3_next_parity integration test replays a committed tiny
  random-weight HF Qwen3NextForCausalLM checkpoint (generated by
  script/dump_qwen3_next_tiny.py on beast: transformers 5.9.0,
  torch 2.9.1) through neuron's full load path: max_abs 0.000000,
  cosine 1.00000000 at f32 — exact parity, pinning the config
  normalisation, weight prefix, qkvz/ba de-interleave, hybrid layer
  interleaving, and the whole MoE block against upstream.
- Unit tests: scatter forward vs per-token dense reference,
  no-shared-expert/no-renorm behaviour, fused-split round-trip, and a
  flat-layout end-to-end structural load.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-02 00:16:19 +03:00
63fb8c6e76 chore(bench): enable 27B baseline capability + concurrency scenarios (#94)
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 11s
build-prerelease / Build neuron-blackwell (push) Successful in 1m36s
build-prerelease / Build neuron-ampere (push) Successful in 2m11s
build-prerelease / Build neuron-ada (push) Successful in 2m7s
build-prerelease / Build helexa-bench binary (push) Successful in 2m43s
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m15s
build-prerelease / Build cortex binary (push) Successful in 3m34s
build-prerelease / Build helexa-upstream binary (push) Successful in 3m38s
build-prerelease / Package helexa-bench RPM (push) Successful in 20s
build-prerelease / Package cortex RPM (push) Successful in 24s
build-prerelease / Package helexa-upstream RPM (push) Successful in 22s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 49s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 50s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 58s
build-prerelease / Test (push) Successful in 7m18s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 17s
Opt in the bob bench host to the concurrency (#89) and capability-probe
(#91) scenario families so the Qwen3.6-27B baseline carries
p95-under-concurrency and scored planning-quality data before the F3
A/B decision gate (#94) needs it for comparison against the 80B-A3B
variants. Already synced to bob and live; this is the tracked record.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-01 23:59:59 +03:00
a1426f177c feat(neuron): accept qwen3_next configs into the qwen3_5 arch (#92)
All checks were successful
CI / Format (push) Successful in 8s
CI / CUDA type-check (push) Successful in 1m59s
CI / Clippy (push) Successful in 2m26s
CI / Test (push) Successful in 6m42s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
F1 slice 1 — config + normalization, no kernel work yet:

- TextConfig gains the MoE hyperparameters (num_experts,
  num_experts_per_tok, moe_intermediate_size,
  shared_expert_intermediate_size, decoder_sparse_step,
  mlp_only_layers, norm_topk_prob), all defaulting to the dense case
  so existing Qwen3.6 configs parse unchanged, plus layer_uses_moe()
  implementing the upstream sparse-layer selection.
- Config::from_config_json normalises the flat qwen3_next layout
  (Qwen3-Next-80B-A3B family) into the nested qwen3_5 shape: wraps
  flat hyperparameters, nests flat rope fields into rope_parameters,
  forces attn_output_gate (unconditional in upstream qwen3_next), and
  derives layer_types from full_attention_interval via the upstream
  (i+1) % interval convention when absent.
- All four config parse sites (dense CPU, dense CUDA worker, TP
  leader, TP worker subprocess) route through from_config_json;
  "qwen3_next" added to DENSE/TP_SUPPORTED_MODEL_TYPES and the
  dispatch match arms.
- Explicit num_experts > 0 guards in Qwen3_5Model::load and
  TpQwen3_5Model::load so a premature 80B load fails with a clear
  "MoE not implemented yet (#92)" instead of a cryptic missing-tensor
  error. Slices 2/3 remove them.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-01 23:59:44 +03:00
a9f5ac7ab4 Merge pull request 'feat(cortex): re-expose flat max_model_len / max_input_tokens / max_output_tokens on /v1/models (#78)' (#107) from feat/78-max-model-len into main
Some checks failed
build-prerelease / Lint (fmt + clippy) (push) Blocked by required conditions
build-prerelease / Build neuron-blackwell (push) Blocked by required conditions
build-prerelease / Resolve version stamps + change detection (push) Successful in 13s
build-prerelease / Test (push) Has been cancelled
build-prerelease / Build cortex binary (push) Has been cancelled
build-prerelease / Build helexa-bench binary (push) Has been cancelled
build-prerelease / Build helexa-upstream binary (push) Has been cancelled
build-prerelease / Build neuron-ampere (push) Has been cancelled
build-prerelease / Build neuron-ada (push) Has been cancelled
build-prerelease / Package cortex RPM (push) Has been cancelled
build-prerelease / Package helexa-bench RPM (push) Has been cancelled
build-prerelease / Package helexa-upstream RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ada RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been cancelled
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been cancelled
2026-07-01 20:58:08 +00:00
0126144a93 ci: retrigger after runner died mid-job (#78)
All checks were successful
CI / Format (push) Successful in 8s
CI / CUDA type-check (pull_request) Successful in 1m32s
CI / Format (pull_request) Successful in 6s
CI / CUDA type-check (push) Successful in 1m38s
CI / Clippy (pull_request) Successful in 3m12s
CI / Clippy (push) Successful in 3m29s
CI / Test (push) Successful in 10m39s
CI / Test (pull_request) Successful in 10m23s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-01 23:44:01 +03:00
869033d08e feat(cortex): re-expose flat max_model_len/max_input_tokens/max_output_tokens on /v1/models (#78)
Some checks failed
CI / Format (push) Successful in 13s
CI / CUDA type-check (pull_request) Successful in 1m26s
CI / Format (pull_request) Successful in 1m35s
CI / Clippy (pull_request) Successful in 5m19s
CI / Test (push) Successful in 10m37s
CI / Test (pull_request) Failing after 10m25s
CI / Clippy (push) Failing after 14m38s
CI / CUDA type-check (push) Failing after 16m6s
CI / Build cortex SRPM (push) Has been cancelled
CI / Build neuron SRPM (push) Has been cancelled
CI / Publish cortex to COPR (push) Has been cancelled
CI / Publish neuron to COPR (push) Has been cancelled
CI / Bump version in source (push) Has been cancelled
CI / Build cortex SRPM (pull_request) Has been cancelled
CI / Build neuron SRPM (pull_request) Has been cancelled
CI / Publish cortex to COPR (pull_request) Has been cancelled
CI / Publish neuron to COPR (pull_request) Has been cancelled
CI / Bump version in source (pull_request) Has been cancelled
Hermes Agent (and the wider vLLM-convention client ecosystem) probes
/v1/models for flat context-window keys and cannot see helexa's
limit.context — it fell back to a hardcoded catalogue guess that only
matched by luck. Re-add the flat fields additively, derived from the
settled `limit` at serialization time (cortex list_models and the
router's federation aggregate), omitted when the window is genuinely
unknown. `limit` stays the opencode-oriented source of truth.

Flips the old regression guard that asserted max_model_len must not
appear — the removal it guarded was based on the wrong assumption that
the field had no consumer.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx
2026-07-01 23:02:09 +03:00
7c5011e5d3 docs: remove milestone-b hand-off from source (moved to gitignored doc/plan/)
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 37s
build-prerelease / Build neuron-blackwell (push) Has been skipped
build-prerelease / Build neuron-ampere (push) Has been skipped
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Build helexa-upstream binary (push) Has been skipped
build-prerelease / Package helexa-upstream RPM (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Successful in 2m3s
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m16s
build-prerelease / Package helexa-bench RPM (push) Successful in 1m24s
build-prerelease / Test (push) Successful in 6m37s
build-prerelease / Build neuron-ada (push) Has been skipped
build-prerelease / Package helexa-neuron-ada RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been skipped
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 51s
The hand-off brief was committed unintentionally; it belongs in the local
doc/plan/ notes folder (gitignored), not in source. Content preserved
locally at doc/plan/milestone-b-epic-84.md.

Reverts d1366614.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VrJ4i3pfLRSTM76o3ofnVq
2026-06-27 15:24:23 +03:00
d1366614f9 docs: Milestone B (epic #84) session hand-off brief
Some checks failed
build-prerelease / Test (push) Blocked by required conditions
build-prerelease / Build helexa-bench binary (push) Blocked by required conditions
build-prerelease / Resolve version stamps + change detection (push) Successful in 35s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m49s
build-prerelease / Build neuron-blackwell (push) Has been skipped
build-prerelease / Build neuron-ampere (push) Has been skipped
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Build helexa-upstream binary (push) Has been skipped
build-prerelease / Package helexa-upstream RPM (push) Has been skipped
build-prerelease / Build neuron-ada (push) Has been cancelled
build-prerelease / Package helexa-bench RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ada RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been cancelled
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been cancelled
Hand-off for the next session opening the Reasoning-frontier milestone
(80B-A3B MoE). Captures: the strategic rationale, decisions already made
(vision via cold-swap; variant-agnostic MoE; measure-first), what
Milestone A shipped and how to use the bench for F3, the F1–F4e todo with
dependencies, an F1 deep-dive (the qwen3_5+MoE fusion, files, HF-reference
validation), and the branch→local-CI→CUDA-gate→merge rhythm.

Docs-only; no cfg(cuda) impact.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VrJ4i3pfLRSTM76o3ofnVq
2026-06-27 15:19:40 +03:00
1884535542 Merge pull request 'feat(helexa-bench): capability probe scenario + quality scoring (#91)' (#106) from feat/91-bench-capability into main
Some checks failed
build-prerelease / Resolve version stamps + change detection (push) Successful in 44s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Build helexa-upstream binary (push) Has been skipped
build-prerelease / Package helexa-upstream RPM (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Successful in 2m17s
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m58s
build-prerelease / Package helexa-bench RPM (push) Successful in 1m19s
build-prerelease / Test (push) Successful in 7m34s
build-prerelease / Build neuron-blackwell (push) Has been skipped
build-prerelease / Build neuron-ampere (push) Has been cancelled
build-prerelease / Build neuron-ada (push) Has been cancelled
build-prerelease / Package helexa-neuron-ada RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been cancelled
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been cancelled
2026-06-27 12:10:20 +00:00
140f9124ee feat(helexa-bench): capability probe scenario + quality scoring (#91)
All checks were successful
CI / Format (push) Successful in 43s
CI / CUDA type-check (push) Successful in 2m35s
CI / Clippy (push) Successful in 2m46s
CI / Test (push) Successful in 5m54s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Format (pull_request) Successful in 43s
CI / CUDA type-check (pull_request) Successful in 1m53s
CI / Clippy (pull_request) Successful in 2m57s
CI / Test (pull_request) Successful in 7m42s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
The six perf scenarios measure speed/resources; this measures the axis
they miss — reasoning/planning quality — so the frontier A/B (F3) can
pick on capability, not just throughput.

Per the chosen approach: store the artifact always, with schema for BOTH
a manual score and a future LLM-judge; start manual.

- scenario: CapabilityScenario (capability:<name>) runs a fixed prompt
  and captures the full output text (stream_and_measure gains a
  capture_text path); opt-in via config.capability_probes (empty
  default — long outputs, deliberate).
- store: three additive columns (artifact, quality_score, scorer);
  capability_runs(unscored_only) worklist + set_score(id, score, scorer).
  Drill-down RunRow omits the large artifact column.
- cli: `helexa-bench score --id <n> --score <x> [--scorer ...]` (manual);
  `report --capability` (per-model median score + per-run artifact
  snippets); GET /api/capability. LLM-judge deferred (schema ready).
- example config documents an implementation-planning probe.

Tests: artifact storage + scoring lifecycle, capability scenario built
from config, capability markdown (median + snippet).

Part of the Performance observability epic (#83), O7 — completes the
milestone. Feeds the F3 frontier A/B decision gate (#94).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VrJ4i3pfLRSTM76o3ofnVq
2026-06-27 15:03:13 +03:00
3b77dbfa32 Merge pull request 'feat(helexa-bench): cold-load / model-swap cost measurement (#90)' (#105) from feat/90-bench-swap-cost into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 41s
build-prerelease / Build neuron-blackwell (push) Has been skipped
build-prerelease / Build neuron-ampere (push) Has been skipped
build-prerelease / Build neuron-ada (push) Has been skipped
build-prerelease / Package helexa-neuron-ada RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been skipped
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Successful in 2m29s
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m26s
build-prerelease / Package helexa-bench RPM (push) Successful in 1m22s
build-prerelease / Test (push) Successful in 7m24s
build-prerelease / Build helexa-upstream binary (push) Has been skipped
build-prerelease / Package helexa-upstream RPM (push) Has been skipped
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 56s
2026-06-27 11:22:36 +00:00
842caf5d4d feat(helexa-bench): cold-load / model-swap cost measurement (#90)
All checks were successful
CI / Format (push) Successful in 45s
CI / CUDA type-check (push) Successful in 1m39s
CI / Clippy (push) Successful in 2m56s
CI / Test (push) Successful in 5m36s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Format (pull_request) Successful in 41s
CI / CUDA type-check (pull_request) Successful in 1m37s
CI / Clippy (pull_request) Successful in 3m8s
CI / Test (pull_request) Successful in 6m33s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
The vision cold-swap policy (F4e) needs real numbers for unload→reload
latency and the cold first-request after a reload — not guesses.

Adds a deliberate `swap-cost` subcommand (NOT the continuous sweep,
since unloading takes the live model offline for the reload). For each
neuron target's warm models: unload, time the synchronous reload, then
time a cold first request. Recorded under scenario "swap" so it tracks
per build like everything else.

- client: unload_model, load_model, and spec_from_info (reconstructs a
  ModelSpec from /models — TP inferred from device count, quant left
  None for neuron to resolve).
- sweep: Sweeper::swap_cost_once + measure_swap (build_record gains a
  swap-timing param; existing sweep call sites pass None).
- scenario: cold_probe — a single timed request reusing the SSE core.
- store: two additive columns (swap_unload_ms, swap_load_ms) + a
  swap_costs() pivot (reload latency + cold first-request medians).
- report: `report --swap` view + render_swap_json; GET /api/swap.
- cli: `helexa-bench swap-cost` with an explicit offline warning.

Tests: swap_costs pivot + scenario exclusion, swap markdown render.

Part of the Performance observability epic (#83), O6. Feeds the vision
cold-swap policy (#99 / F4e).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VrJ4i3pfLRSTM76o3ofnVq
2026-06-27 14:15:42 +03:00
a488ade675 Merge pull request 'feat(helexa-bench): context-length scaling view (#88)' (#104) from feat/88-bench-context-scaling into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 42s
build-prerelease / Build neuron-blackwell (push) Has been skipped
build-prerelease / Build neuron-ampere (push) Has been skipped
build-prerelease / Build neuron-ada (push) Has been skipped
build-prerelease / Package helexa-neuron-ada RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been skipped
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Successful in 2m25s
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m13s
build-prerelease / Package helexa-bench RPM (push) Successful in 1m20s
build-prerelease / Test (push) Successful in 6m38s
build-prerelease / Build helexa-upstream binary (push) Has been skipped
build-prerelease / Package helexa-upstream RPM (push) Has been skipped
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 49s
2026-06-27 11:03:35 +00:00
d472b6428a feat(helexa-bench): context-length scaling view (#88)
All checks were successful
CI / Format (push) Successful in 38s
CI / CUDA type-check (push) Successful in 1m49s
CI / Clippy (push) Successful in 2m44s
CI / Test (push) Successful in 6m5s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Format (pull_request) Successful in 38s
CI / CUDA type-check (pull_request) Successful in 1m36s
CI / Clippy (pull_request) Successful in 3m12s
CI / Test (pull_request) Successful in 6m14s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
The chat:<n> cells already capture prefill & decode tok/s per context
(via #85/#86); this pivots them into a per-(target,model) scaling curve
and computes decode-flatness — decode tok/s at the largest context ÷ the
smallest. ~1.0 confirms the Gated-DeltaNet O(1)-in-sequence-length
decode; a sharp drop locates where the model stops scaling for free.

- store: Store::scaling() pivots the latest-build chat:<n> report cells
  into ScalingCurve/ScalingPoint, ordered by context, with the flatness
  ratio (concurrency:<n> and other scenarios are excluded).
- report: render_scaling_markdown (one block per model: prefill/decode
  tok/s vs ctx + a flatness verdict) and render_scaling_json.
- cli: `helexa-bench report --scaling` selects the view.
- api: GET /api/scaling for the bench UI.
- example config documents widening prompt_sizes into a scaling ladder.

No new request shape — reuses the chat-latency measurement points, so a
denser curve is just more prompt_sizes entries (operators widen the
ladder deliberately; large contexts cost more per sample).

Tests: scaling pivot + flatness + scenario exclusion, markdown render.

Part of the Performance observability epic (#83), O4. Validates the
long-context property that makes the 80B-A3B frontier model (#84) viable.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VrJ4i3pfLRSTM76o3ofnVq
2026-06-27 13:56:07 +03:00
0488768afc Merge pull request 'feat(helexa-bench): concurrency / agentic-load scenario (#89)' (#103) from feat/89-bench-concurrency into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 38s
build-prerelease / Build neuron-blackwell (push) Has been skipped
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m36s
build-prerelease / Test (push) Successful in 6m34s
build-prerelease / Build neuron-ampere (push) Has been skipped
build-prerelease / Build neuron-ada (push) Has been skipped
build-prerelease / Package helexa-neuron-ada RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been skipped
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Build helexa-upstream binary (push) Has been skipped
build-prerelease / Package helexa-upstream RPM (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Successful in 2m13s
build-prerelease / Package helexa-bench RPM (push) Successful in 1m27s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 59s
2026-06-27 10:25:27 +00:00
209bf58207 feat(helexa-bench): concurrency / agentic-load scenario (#89)
All checks were successful
CI / Format (push) Successful in 37s
CI / CUDA type-check (push) Successful in 1m37s
CI / Clippy (push) Successful in 2m46s
CI / Test (push) Successful in 6m17s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Format (pull_request) Successful in 38s
CI / CUDA type-check (pull_request) Successful in 1m36s
CI / Clippy (pull_request) Successful in 3m5s
CI / Test (pull_request) Successful in 6m16s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
Batch-1 single-request timing can't characterize the real workload —
a0/hermes/opencode each fan out many agentic requests per user turn.

Add a ConcurrencyScenario (concurrency:<n>) that fires N simultaneous
streams and records, per burst: aggregate node throughput (total tokens
/ burst window), within-burst TTFT p95 (the tail under load), median
admission queue-wait — derived as TTFT − server prefill_ms (#85), no
/health poll needed — and the count of streams shed by admission
(429/503). On a batch-1 server, throughput stays flat while queue-wait
and p95 inflate with N; that gap is the evidence for/against continuous
batching (#98).

- scenario: ConcurrencyScenario + shared chat_payload helper; reuses
  stream_and_measure per stream via join_all.
- config: concurrency_levels (default empty — opt-in, a burst loads the
  fleet) + concurrency_prompt_tokens (512); example config documents it.
- store: four additive columns (same migration pattern); aggregate
  surfaces concurrency, within-burst p95, queue-wait, rejected medians.
- report: markdown gains conc / queue ms / rej columns; JSON the full set.

Tests: median/percentile/admission-detect helpers, config builds
concurrency scenarios, store surfaces burst metrics.

Part of the Performance observability epic (#83), O5. Produces the
evidence base for continuous batching (#98 / F4d).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VrJ4i3pfLRSTM76o3ofnVq
2026-06-27 13:17:10 +03:00
1ea7425585 Merge pull request 'feat(helexa-bench): per-run VRAM high-water + GPU util/temp from /health (#87)' (#102) from feat/87-bench-vram-telemetry into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 44s
build-prerelease / Build neuron-blackwell (push) Has been skipped
build-prerelease / Build neuron-ampere (push) Has been skipped
build-prerelease / Build neuron-ada (push) Has been skipped
build-prerelease / Package helexa-neuron-ada RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been skipped
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been skipped
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Build helexa-upstream binary (push) Has been skipped
build-prerelease / Package helexa-upstream RPM (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Successful in 2m4s
build-prerelease / Package helexa-bench RPM (push) Successful in 1m18s
build-prerelease / Lint (fmt + clippy) (push) Successful in 4m24s
build-prerelease / Test (push) Successful in 7m22s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 50s
2026-06-27 09:57:58 +00:00
a415454962 feat(helexa-bench): per-run VRAM high-water + GPU util/temp from /health (#87)
All checks were successful
CI / Format (push) Successful in 37s
CI / CUDA type-check (push) Successful in 1m37s
CI / Clippy (push) Successful in 2m46s
CI / Test (push) Successful in 6m0s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Format (pull_request) Successful in 39s
CI / CUDA type-check (pull_request) Successful in 1m37s
CI / Clippy (pull_request) Successful in 2m58s
CI / Test (pull_request) Successful in 7m18s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
Bench had no measured VRAM headroom — only the anecdotal "~2/3 used".
Sample neuron's GET /health (DeviceHealth: vram_used/free, util%, temp)
right after each measured run and record node-sum VRAM used plus the
hottest device's utilization and temperature.

- client: fetch_health (neuron only; soft None on transport failure so a
  flaky /health never fails a measurement).
- sweep: HealthAgg folds a /health snapshot to node totals; sampled
  post-run (the recent decode-VRAM peak) and threaded into the record.
  /health is ~5s-cached neuron-side, so this is a coarse high-water proxy,
  documented as such.
- store: three additive columns (same PRAGMA-guarded migration);
  aggregate emits vram_used median + node vram_total (summed from the
  discovery devices in gpus_json) for real headroom, plus util/temp
  medians.
- report: markdown gains a "VRAM (GB)" used/total column; JSON gains the
  used/total + util/temp medians.

Tests: VRAM/headroom aggregation, gpu_total_vram summation.

Part of the Performance observability epic (#83). Completes the p1-now
trio (O1-O3); unblocks the enriched-27B baseline capture on beast.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VrJ4i3pfLRSTM76o3ofnVq
2026-06-27 12:50:53 +03:00
06a36566d1 Merge pull request 'feat(helexa-bench): percentiles + prefill/decode split in store & report (#86)' (#101) from feat/86-bench-percentiles into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 31s
build-prerelease / Build neuron-blackwell (push) Successful in 1m39s
build-prerelease / Build neuron-ada (push) Successful in 2m15s
build-prerelease / Build neuron-ampere (push) Successful in 2m18s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 1m52s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 1m49s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 2m15s
build-prerelease / Build cortex binary (push) Successful in 2m39s
build-prerelease / Lint (fmt + clippy) (push) Successful in 2m51s
build-prerelease / Package cortex RPM (push) Successful in 1m15s
build-prerelease / Test (push) Successful in 6m52s
build-prerelease / Build helexa-bench binary (push) Successful in 2m12s
build-prerelease / Build helexa-upstream binary (push) Successful in 2m52s
build-prerelease / Package helexa-bench RPM (push) Successful in 1m19s
build-prerelease / Package helexa-upstream RPM (push) Successful in 1m18s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 54s
2026-06-27 09:08:53 +00:00
4cb52e3144 Merge pull request 'feat(neuron): server-measured prefill/decode timing on Finish (#85)' (#100) from feat/85-prefill-decode-timing into main
Some checks failed
build-prerelease / Resolve version stamps + change detection (push) Waiting to run
build-prerelease / Lint (fmt + clippy) (push) Has been cancelled
build-prerelease / Test (push) Has been cancelled
build-prerelease / Build cortex binary (push) Has been cancelled
build-prerelease / Build helexa-bench binary (push) Has been cancelled
build-prerelease / Build helexa-upstream binary (push) Has been cancelled
build-prerelease / Build neuron-blackwell (push) Has been cancelled
build-prerelease / Build neuron-ampere (push) Has been cancelled
build-prerelease / Build neuron-ada (push) Has been cancelled
build-prerelease / Package cortex RPM (push) Has been cancelled
build-prerelease / Package helexa-bench RPM (push) Has been cancelled
build-prerelease / Package helexa-upstream RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ada RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been cancelled
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been cancelled
2026-06-27 09:08:15 +00:00
afc1f7a706 feat(helexa-bench): percentiles + prefill/decode split in store & report (#86)
All checks were successful
CI / Format (push) Successful in 38s
CI / CUDA type-check (push) Successful in 1m37s
CI / Clippy (push) Successful in 3m13s
CI / Test (push) Successful in 6m24s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Format (pull_request) Successful in 47s
CI / CUDA type-check (pull_request) Successful in 1m32s
CI / Clippy (pull_request) Successful in 3m11s
CI / Test (pull_request) Successful in 6m35s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
Bench reported only a single median per cell, hiding tail latency and
unable to record the server-measured prefill/decode split now emitted
on `usage.helexa_timing` (#85).

- scenario: parse `usage.helexa_timing` into ScenarioMetrics
  (prefill_ms, decode_ms, prefill_tokens).
- store: persist the three columns (additive PRAGMA-guarded migration
  via ensure_columns, so pre-#85 DBs backfill as NULL); aggregate now
  emits p50/p95/p99 for TTFT and total (nearest-rank) plus a
  prefill-tok/s median derived from the split.
- report: markdown gains prefill tok/s, TTFT p95, total p95 columns;
  JSON gains p95/p99 + prefill_ms/decode_ms/prefill_tps medians.

Tests: nearest-rank percentile, idempotent backfill migration, and a
report cell asserting percentiles + prefill split.

Part of the Performance observability epic (#83). Stacked on #85.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VrJ4i3pfLRSTM76o3ofnVq
2026-06-27 11:58:36 +03:00
6f956dfda3 fix(neuron): hoist TP prefill/decode timers out of 'work block (#85)
All checks were successful
CI / Format (push) Successful in 35s
CI / CUDA type-check (push) Successful in 1m32s
CI / Clippy (push) Successful in 3m17s
CI / Test (push) Successful in 7m6s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Format (pull_request) Successful in 39s
CI / CUDA type-check (pull_request) Successful in 1m49s
CI / Clippy (pull_request) Successful in 2m58s
CI / Test (pull_request) Successful in 7m3s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
The TP streaming producer builds its Finish event after the 'work
labelled block exits (inside `if failure.is_none()`), but prefill_elapsed
and decode_start were declared inside that block — so the CUDA type-check
failed with E0425 (the CPU build doesn't compile this cfg(cuda) path).

Hoist `prefill_ms_measured: u32` and `decode_start: Option<Instant>`
above the block; populate them at the prefill→decode boundary inside;
read them at the terminal Finish. The worker and local producers were
unaffected (their timers already share the Finish scope).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VrJ4i3pfLRSTM76o3ofnVq
2026-06-27 11:58:26 +03:00
6e0f15c888 feat(neuron): server-measured prefill/decode timing on Finish (#85)
Some checks failed
CI / CUDA type-check (push) Failing after 1m40s
CI / Format (push) Successful in 42s
CI / Clippy (push) Successful in 2m51s
CI / Test (push) Successful in 6m16s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
The harness emitted only token counts on InferenceEvent::Finish; all
timing was client-side SSE arrival, so bench "TTFT" conflated tokenize
+ prefill and decode tok/s was a window estimate.

Add FinishTiming { prefill_ms, decode_ms, prefill_tokens } to the
Finish event, populated by all three streaming producers (TP, worker,
and local CPU paths), and surface it on the OpenAI chat
`usage.helexa_timing` extension so helexa-bench can compute true
prefill vs decode tok/s. cortex forwards usage verbatim, so the field
survives proxying. Non-streaming and Responses paths carry None for
now (bench reads the streaming chat path).

Keystone for the Performance observability epic (#83).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VrJ4i3pfLRSTM76o3ofnVq
2026-06-27 11:41:17 +03:00
66eb9f558f Merge pull request 'fix(neuron): surface reasoning_tokens in non-streaming /v1/responses usage' (#82) from fix/responses-usage-reasoning-tokens into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 47s
build-prerelease / Build neuron-blackwell (push) Successful in 1m27s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Has been skipped
build-prerelease / Package helexa-bench RPM (push) Has been skipped
build-prerelease / Build neuron-ada (push) Successful in 2m4s
build-prerelease / Build neuron-ampere (push) Successful in 2m14s
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m14s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 1m41s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 1m57s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 1m52s
build-prerelease / Test (push) Successful in 6m41s
build-prerelease / Build helexa-upstream binary (push) Has been skipped
build-prerelease / Package helexa-upstream RPM (push) Has been skipped
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 55s
2026-06-26 18:30:05 +00:00
f96a2e7ed3 fix(neuron): surface reasoning_tokens in non-streaming /v1/responses usage
All checks were successful
CI / Format (push) Successful in 35s
CI / CUDA type-check (push) Successful in 1m38s
CI / Clippy (push) Successful in 2m42s
CI / Test (push) Successful in 6m36s
CI / Build cortex SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Format (pull_request) Successful in 36s
CI / CUDA type-check (pull_request) Successful in 1m35s
CI / Clippy (pull_request) Successful in 2m54s
CI / Test (pull_request) Successful in 6m33s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
The non-streaming responses handler hardcoded `output_tokens_details: None`,
so the reasoning sub-count that the chat path now computes (via
`split_off_reasoning`, which strips the `<think>` span and counts it into
`completion_tokens_details.reasoning_tokens`) never reached the Responses-API
usage object. The streaming responses path already emits it.

Carry `chat usage.completion_tokens_details.reasoning_tokens` through to
`ResponsesUsage.output_tokens_details.reasoning_tokens`, so streaming and
non-streaming `/v1/responses` report reasoning accounting identically.
`output_tokens` still counts every generated token (reasoning included);
`reasoning_tokens` is the additive sub-count, per OpenAI's shape.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RLyKaJVFDAYnAGiLvVrK8K
2026-06-26 21:21:59 +03:00
b17b555a3d Merge pull request 'fix(neuron): strip &lt;think&gt; reasoning from non-streaming completions' (#81) from fix/responses-nonstreaming-reasoning-leak into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 46s
build-prerelease / Build neuron-blackwell (push) Successful in 1m27s
build-prerelease / Build cortex binary (push) Has been skipped
build-prerelease / Build helexa-bench binary (push) Has been skipped
build-prerelease / Package cortex RPM (push) Has been skipped
build-prerelease / Package helexa-bench RPM (push) Has been skipped
build-prerelease / Build neuron-ada (push) Successful in 2m3s
build-prerelease / Build neuron-ampere (push) Successful in 2m12s
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m6s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 1m40s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 1m51s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 1m47s
build-prerelease / Test (push) Successful in 6m26s
build-prerelease / Build helexa-upstream binary (push) Has been skipped
build-prerelease / Package helexa-upstream RPM (push) Has been skipped
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 54s
2026-06-26 17:47:45 +00:00
13daf95514 fix(neuron): strip <think> reasoning from non-streaming completions
All checks were successful
CI / Format (push) Successful in 39s
CI / CUDA type-check (push) Successful in 1m41s
CI / Clippy (push) Successful in 2m47s
CI / Test (push) Successful in 5m51s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Format (pull_request) Successful in 40s
CI / CUDA type-check (pull_request) Successful in 1m38s
CI / Clippy (pull_request) Successful in 2m55s
CI / Test (pull_request) Successful in 6m20s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
The non-streaming inference path (single-GPU `chat_completion` and the TP
`chat_completion_tp_inner`) returned the model's full decode — reasoning
preamble + `</think>` + answer — as the assistant `content`. The streaming
path already drops reasoning (emits it as ReasoningDelta, which the chat and
Responses projectors discard), so the two transports disagreed: a streaming
client saw the clean answer, a non-streaming client saw the chain-of-thought
glued to the front of it.

This is exactly what broke agent-zero v2.0 on `/v1/responses` (non-streaming):
the model produced a correct in-band JSON answer after `</think>`, but a0's
parser saw the `<think>` preamble first and rejected the turn as "misformat,
no valid tool request found". a0 v2.1 happens to tolerate the preamble, but
any strict non-streaming client (and a0 v2.0) does not — and reasoning tokens
leaking into `content` also misreport as visible output.

Add `split_off_reasoning(generated_ids, reasoning_pair)`: if the model
declares a reasoning marker pair and its close token (`</think>`) appears in
the output, return only the tokens after the last close marker as content and
count the rest as reasoning. The chat template injects the *opening* marker
into the prompt, so the generated tokens carry the close marker but not the
open one — splitting on the close-token id (not a decoded string) is robust to
tokenizer byte-fallback. Non-reasoning models, thinking-disabled requests, and
generations truncated mid-reasoning have no close token and pass through
unchanged.

Both non-streaming sites now decode only the answer span and populate
`completion_tokens_details.reasoning_tokens` (the streaming-only accounting
gap noted at the call sites, #64). Unit tests cover the strip, no-marker,
no-pair, close-at-end, and multiple-close cases.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RLyKaJVFDAYnAGiLvVrK8K
2026-06-26 20:39:41 +03:00
319b01e0b2 Merge pull request 'fix(neuron): accept bare {role,content} input on /v1/responses (agent-zero)' (#80) from fix/responses-easy-message-input into main
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 49s
build-prerelease / Build neuron-blackwell (push) Successful in 1m29s
build-prerelease / Build helexa-bench binary (push) Successful in 2m21s
build-prerelease / Build cortex binary (push) Successful in 3m23s
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m50s
build-prerelease / Package helexa-bench RPM (push) Successful in 1m39s
build-prerelease / Test (push) Successful in 6m43s
build-prerelease / Package cortex RPM (push) Successful in 1m30s
build-prerelease / Build neuron-ada (push) Successful in 2m7s
build-prerelease / Build neuron-ampere (push) Successful in 2m15s
build-prerelease / Build helexa-upstream binary (push) Successful in 2m42s
build-prerelease / Package helexa-upstream RPM (push) Successful in 1m20s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 1m39s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 1m42s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 1m46s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 51s
2026-06-26 16:53:08 +00:00
6731adca51 fix(neuron): accept bare {role,content} input on /v1/responses (agent-zero)
All checks were successful
CI / Format (push) Successful in 41s
CI / CUDA type-check (push) Successful in 1m40s
CI / Clippy (push) Successful in 2m45s
CI / Test (push) Successful in 6m14s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
CI / Format (pull_request) Successful in 39s
CI / CUDA type-check (pull_request) Successful in 1m36s
CI / Clippy (pull_request) Successful in 3m18s
CI / Test (pull_request) Successful in 7m6s
CI / Build cortex SRPM (pull_request) Has been skipped
CI / Publish cortex to COPR (pull_request) Has been skipped
CI / Build neuron SRPM (pull_request) Has been skipped
CI / Publish neuron to COPR (pull_request) Has been skipped
CI / Bump version in source (pull_request) Has been skipped
agent-zero (via litellm) drives the OpenAI Responses API and sends
`input` items in the "easy input message" form — bare `{role, content}`
objects with NO `type` field. Our `ResponsesInputItem` is internally
tagged (`#[serde(tag="type")]`), so every such item failed the untagged
`ResponsesInput` deserialize and axum's `Json` extractor returned 422:

    OpenAIException - Failed to deserialize the JSON body into the target
    type: data did not match any variant of untagged enum ResponsesInput

This was a *total* failure for agent-zero (both the main model on beast
and the utility model on benjy), confirmed by on-wire capture of 15 live
requests: 36/36 input items were bare easy-messages. Other clients
(/v1/chat/completions, /v1/messages) were unaffected — only the
Responses path was exercised this strictly, for the first time.

Make `input`-item parsing match OpenAI's real tolerance, mirroring the
forward-compat `extra: Value` already at the top level of the request:

- New `ResponsesInputElement` wraps the existing typed item enum with
  two more shapes: `EasyMessage { role, content }` (bare, no type;
  `content` optional so an assistant turn with `content: null` parses)
  and `Other(Value)` — a catch-all so a single unmodeled item can never
  again 422 the whole request. The typed enum is unchanged.
- `ResponsesContentPart` gains a `#[serde(other)] Unknown` arm (e.g.
  `refusal`, audio) — dropped in translation, not rejected.
- `FunctionCallOutput.output` is now `Value` (string OR array of content
  parts, per OpenAI) so a structured tool result isn't lost.
- Translator handles all three element shapes; easy-messages translate
  exactly like typed messages, `Other` and unknown parts are dropped.

Tests cover the bare-message, null-content, unknown-item, unknown-part,
and array-tool-output shapes, validated against the 15 captured bodies.

Tools forwarding + native function_call projection on the Responses path
is deliberately a follow-up (Round 2), gated on observing how agent-zero
consumes responses once unblocked (in-band JSON vs native tool calls).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RLyKaJVFDAYnAGiLvVrK8K
2026-06-26 19:45:22 +03:00
7e11a7688c Merge feat/F6-beta-polish: public-beta banner + same-origin nginx + deploy (F6)
All checks were successful
build-prerelease / Resolve version stamps + change detection (push) Successful in 39s
build-prerelease / Lint (fmt + clippy) (push) Successful in 3m17s
build-prerelease / Build neuron-blackwell (push) Successful in 1m43s
build-prerelease / Build neuron-ada (push) Successful in 2m15s
build-prerelease / Build neuron-ampere (push) Successful in 2m17s
build-prerelease / Build cortex binary (push) Successful in 2m50s
build-prerelease / Build helexa-bench binary (push) Successful in 2m25s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 2m1s
build-prerelease / Build helexa-upstream binary (push) Successful in 4m8s
build-prerelease / Test (push) Successful in 6m28s
build-prerelease / Package helexa-bench RPM (push) Successful in 1m18s
build-prerelease / Package helexa-upstream RPM (push) Successful in 1m19s
build-prerelease / Package cortex RPM (push) Successful in 1m24s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 1m39s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 1m44s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 58s
2026-06-23 12:13:35 +03:00
5600575ba2 Merge feat/B7-upstream-packaging: RPM + systemd + CI for helexa-upstream (B7)
Some checks failed
build-prerelease / Resolve version stamps + change detection (push) Has been cancelled
build-prerelease / Lint (fmt + clippy) (push) Has been cancelled
build-prerelease / Test (push) Has been cancelled
build-prerelease / Build cortex binary (push) Has been cancelled
build-prerelease / Build helexa-bench binary (push) Has been cancelled
build-prerelease / Build helexa-upstream binary (push) Has been cancelled
build-prerelease / Build neuron-blackwell (push) Has been cancelled
build-prerelease / Build neuron-ampere (push) Has been cancelled
build-prerelease / Build neuron-ada (push) Has been cancelled
build-prerelease / Package cortex RPM (push) Has been cancelled
build-prerelease / Package helexa-bench RPM (push) Has been cancelled
build-prerelease / Package helexa-upstream RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ada RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-ampere RPM (push) Has been cancelled
build-prerelease / Package helexa-neuron-blackwell RPM (push) Has been cancelled
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Has been cancelled
2026-06-23 12:13:25 +03:00
452d7d9b3d feat(B7): packaging + CI for helexa-upstream
All checks were successful
CI / CUDA type-check (push) Successful in 1m32s
CI / Format (push) Successful in 57s
CI / Clippy (push) Successful in 2m51s
CI / Test (push) Successful in 8m0s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
Ships the mesh authority as an RPM on the same prebuilt-binary pipeline as
helexa-bench.

- rpm/helexa-upstream-prerelease.spec: wraps the CI-built binary; installs
  the systemd unit, sysusers, firewalld service (tcp/8090), and
  /etc/helexa-upstream/helexa-upstream.toml (config noreplace).
- data/helexa-upstream.{service,sysusers.conf,firewalld.xml}: Type=simple
  unit (serve --config), dedicated system user with a StateDirectory home,
  inbound 8090 (/authz/v1 + /web/v1). PostgreSQL is reached outbound; the
  schema migrates on startup.
- build-prerelease.yml: build-upstream + package-upstream jobs with
  change-detection over crates/helexa-upstream/ (UPSTREAM_RE), gated into
  publish. SQLX_OFFLINE=true is set defensively — helexa-upstream uses the
  sqlx runtime query API (no compile-time macros), so it builds with no
  database and no .sqlx cache; DB integration tests stay gated behind
  UPSTREAM_TEST_DATABASE_URL.

Validated: workflow YAML parses, rpmspec expands, and
`SQLX_OFFLINE=true cargo build --release -p helexa-upstream` succeeds.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01F6o3ddqmYNh9kzdwq6eowh
2026-06-23 12:00:41 +03:00
61 changed files with 8830 additions and 299 deletions

View File

@@ -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:

View File

@@ -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

View File

@@ -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
View File

@@ -2972,6 +2972,7 @@ dependencies = [
"axum",
"base64 0.22.1",
"candle-core",
"candle-flash-attn",
"candle-nn",
"candle-transformers",
"clap",

View File

@@ -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.

View File

@@ -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)]

View File

@@ -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`.

View File

@@ -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);

View File

@@ -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();

View File

@@ -581,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,
},
);
}
@@ -638,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,
});
}
}
@@ -694,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,
});
}
}
@@ -728,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,

View File

@@ -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);
}

View File

@@ -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.

View File

@@ -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

View File

@@ -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

View File

@@ -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(())

View File

@@ -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,
}];

View File

@@ -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.

View File

@@ -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"}]"#;

View File

@@ -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,
}

View File

@@ -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()) },
}

View File

@@ -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 {

View File

@@ -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));
}
}

View File

@@ -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,
}
}

View File

@@ -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.

View File

@@ -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 {

View File

@@ -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");
}
}

View File

@@ -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(

View File

@@ -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

View File

@@ -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);
}
}

View File

@@ -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,

View File

@@ -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

View 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
);
}
}
}

View File

@@ -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

View File

@@ -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.

View File

@@ -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 {
@@ -472,7 +480,7 @@ pub struct TpLoadedModel {
/// so this Mutex no longer covers the leader's KV cache; it just
/// serialises subprocess RPC traffic on the pool's
/// `Vec<Worker>` channels.
pub pool: tokio::sync::Mutex<super::tp::WorkerPool>,
pub pool: Arc<tokio::sync::Mutex<super::tp::WorkerPool>>,
/// Bounded admission scheduler (#53), mirroring the single-GPU path.
/// Gated before the pool lock so an overloaded TP model returns fast
/// backpressure instead of an unbounded, untimed wait.
@@ -537,6 +545,13 @@ pub struct TpLoadedModel {
/// Mint for pool-wide snapshot ids. Plain counter; uniqueness only
/// needs to hold per model lifetime (snapshots die with the model).
pub next_snapshot_id: std::sync::atomic::AtomicU64,
/// Lockstep batched decode engine (#98) — the TP mirror of
/// `LoadedModel::engine`. Set once at load (after the `Arc` is
/// built, so the engine can hold a `Weak` back to this model);
/// unset when `max_in_flight` is 1 or batching is killed. Text
/// chat streams route through it; the engine holds the pool mutex
/// while it has active slots.
pub engine: std::sync::OnceLock<super::engine::EngineHandle>,
/// Cached tightest free VRAM (MiB) for the control plane (#53) — see
/// [`LoadedModel::last_free_mb`]. Read by `derived_limit` so `GET /models`
/// never fans a VRAM query out to the (inference-saturated) TP workers.
@@ -679,6 +694,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 +834,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 +953,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 +1009,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 +1080,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 +1121,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 +1181,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 +1555,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 +1614,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],
@@ -1572,7 +1664,7 @@ async fn chunked_prefill_via_worker(
/// `start_offset` skips a restored cached prefix, as in
/// [`chunked_prefill_local`].
#[cfg(feature = "cuda")]
async fn chunked_prefill_tp(
pub(crate) async fn chunked_prefill_tp(
pool: &mut super::tp::WorkerPool,
model_id: &str,
leader_handle: super::device_worker::TpHandle,
@@ -1974,9 +2066,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 +2118,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 +2368,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 +2445,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 +2497,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 +2559,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 +2907,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 +2946,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 +3011,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 +3455,43 @@ 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(),
tokenizer: tokenizer.clone(),
reasoning_tokens: reasoning_tokens.clone(),
tool_call_tokens: tool_call_tokens.clone(),
max_slots: self.admission_cfg.max_in_flight,
backend: super::engine::BackendConfig::Single {
worker: Arc::clone(w),
handle: h,
prefix_cache: prefix_cache.clone(),
prefill_rate: Arc::clone(&prefill_rate),
poisoned: Arc::clone(&poisoned),
inference_lock: Arc::clone(&inference_lock),
},
},
))
}
_ => None,
};
let loaded = Arc::new(LoadedModel {
model_id: spec.model_id.clone(),
arch: arch_local,
@@ -3316,10 +3499,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 +3511,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!(
@@ -3413,12 +3597,35 @@ impl Harness for CandleHarness {
"TP unload: DropTp RPC failed (leader model may leak in worker slab)"
);
}
let mut pool = tp.pool.into_inner();
if let Err(e) = pool.unload_model(model_id).await {
tracing::warn!(model = %model_id, error = %e, "TP unload RPC failed");
}
if let Err(e) = pool.shutdown().await {
tracing::warn!(model = %model_id, error = %e, "TP pool shutdown failed");
// The pool mutex is Arc-shared with the batch engine's
// active-phase guard (#98). `Arc::try_unwrap(tp)`
// succeeding above means the engine is idle (it holds
// `Arc<TpLoadedModel>` whenever it holds the pool
// guard), so sole ownership is the expected case; the
// fallback covers the narrow race where the engine's
// guard is mid-release.
match Arc::try_unwrap(tp.pool) {
Ok(pool_mutex) => {
let mut pool = pool_mutex.into_inner();
if let Err(e) = pool.unload_model(model_id).await {
tracing::warn!(model = %model_id, error = %e, "TP unload RPC failed");
}
if let Err(e) = pool.shutdown().await {
tracing::warn!(model = %model_id, error = %e, "TP pool shutdown failed");
}
}
Err(pool_arc) => {
tracing::warn!(
model = %model_id,
"TP unload: pool mutex still referenced (engine guard \
mid-release); unloading without explicit pool shutdown — \
worker children reap when the last reference drops"
);
let mut pool = pool_arc.lock().await;
if let Err(e) = pool.unload_model(model_id).await {
tracing::warn!(model = %model_id, error = %e, "TP unload RPC failed");
}
}
}
}
}
@@ -3573,7 +3780,7 @@ impl CandleHarness {
model_id: spec.model_id.clone(),
tokenizer,
devices: devices.clone(),
pool: TMutex::new(pool),
pool: StdArc::new(TMutex::new(pool)),
admission: super::admission::AdmissionController::new(&self.admission_cfg),
leader_handle,
leader_device: leader_device.clone(),
@@ -3603,7 +3810,33 @@ impl CandleHarness {
derived_input_cap: AtomicUsize::new(0),
last_free_mb: AtomicU64::new(0),
next_snapshot_id: std::sync::atomic::AtomicU64::new(1),
engine: std::sync::OnceLock::new(),
});
// Batched decode engine (#98): spawned when the operator raised
// max_in_flight above 1 on a snapshot-capable TP model. The
// engine holds only a Weak — an idle engine never keeps an
// unloaded model alive.
if tp_loaded.prefix_cache.is_some()
&& 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 for TP model (#98)"
);
let handle = super::engine::EngineHandle::spawn(super::engine::EngineConfig {
model_id: spec.model_id.clone(),
tokenizer: tp_loaded.tokenizer.clone(),
reasoning_tokens: tp_loaded.reasoning_tokens.clone(),
tool_call_tokens: tp_loaded.tool_call_tokens.clone(),
max_slots: self.admission_cfg.max_in_flight,
backend: super::engine::BackendConfig::Tp {
tp: StdArc::downgrade(&tp_loaded),
},
});
let _ = tp_loaded.engine.set(handle);
}
if tp_loaded.prefix_cache.is_some() {
tracing::info!(
model = %spec.model_id,
@@ -3923,6 +4156,38 @@ impl CandleHarness {
.map_err(InferenceError::from)?;
let tool_schemas = build_tool_schemas(&request);
// Batched decode engine (#98): text streams multiplex through
// the per-model engine instead of serializing on the pool
// mutex per request. Vision requests keep the direct path and
// serialize against the engine via the pool lock it holds
// while active.
if vision_route.is_none()
&& let Some(engine) = tp.engine.get()
{
engine
.submit(super::engine::EngineRequest {
prompt_tokens,
max_new,
temperature,
top_p,
seed,
eos_id,
tool_schemas,
tx,
admit,
span,
})
.await
.map_err(InferenceError::Other)?;
let reasoning_markers = tp.reasoning_tokens.clone();
return Ok(InferenceStream {
events: event_rx,
id: projector_id,
created,
model_id: projector_model_id,
reasoning_markers,
});
}
let tp_for_task = Arc::clone(&tp);
tokio::spawn(
async move {
@@ -3957,6 +4222,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 +4354,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 +4388,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 +4667,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 +5003,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 +5072,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 +5112,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 +5137,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 +5155,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 +5172,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 +5210,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 +5255,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 +5862,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 +5889,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 +5938,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>,
@@ -5697,7 +6001,7 @@ fn restore_or_clear_local(
/// (some restored, some not) — the clear fallback resets every rank,
/// restoring consistency.
#[cfg(feature = "cuda")]
async fn restore_or_clear_tp(
pub(crate) async fn restore_or_clear_tp(
pool: &mut super::tp::WorkerPool,
tp: &TpLoadedModel,
prompt_tokens: &[u32],
@@ -5754,7 +6058,7 @@ async fn restore_or_clear_tp(
/// On any rank failing, drops the id everywhere (idempotent) so no
/// rank leaks a half-stored snapshot.
#[cfg(feature = "cuda")]
async fn store_prefix_snapshot_tp(
pub(crate) async fn store_prefix_snapshot_tp(
pool: &mut super::tp::WorkerPool,
tp: &TpLoadedModel,
prompt_tokens: Vec<u32>,
@@ -6064,7 +6368,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 +6382,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 +6528,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 +6667,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 +6689,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 +6799,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 +6828,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 +6968,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 }"#;

View File

@@ -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,
@@ -436,6 +476,50 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
let _ = reply.send(());
}
#[cfg(feature = "cuda")]
Job::TpForwardLogitsBatch {
handle,
tokens,
prefix_lens,
padded_len,
step,
reply,
} => {
let result = tp_forward_logits_batch(
&mut state,
handle,
&tokens,
&prefix_lens,
padded_len,
step,
);
let _ = reply.send(result);
}
#[cfg(feature = "cuda")]
Job::TpAssembleKvBatch {
handle,
seqs,
reply,
} => {
let result = tp_assemble_kv_batch(&mut state, handle, &seqs);
if result.is_ok() {
trim_device_pool(&state);
}
let _ = reply.send(result);
}
#[cfg(feature = "cuda")]
Job::TpExtractKvRows {
handle,
rows,
padded_len,
steps,
snapshot_ids,
reply,
} => {
let result =
tp_extract_kv_rows(&mut state, handle, &rows, padded_len, steps, &snapshot_ids);
let _ = reply.send(result);
}
#[cfg(feature = "cuda")]
Job::TpForwardLogits {
handle,
tokens,
@@ -760,9 +844,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 +893,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 +964,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 +979,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)"
),
};
@@ -943,6 +1035,114 @@ fn tp_forward_logits(
Ok(values)
}
/// TP-equivalent of [`forward_logits_batch`] on the leader's shard
/// (#98). The caller has already fanned the matching
/// `GenerateStepBatch` out to the subprocess ranks.
#[cfg(feature = "cuda")]
fn tp_forward_logits_batch(
state: &mut DeviceWorkerState,
handle: TpHandle,
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, "TpForwardLogitsBatch: empty batch");
anyhow::ensure!(
prefix_lens.len() == b,
"TpForwardLogitsBatch: {} prefix_lens for batch of {b}",
prefix_lens.len()
);
let input = Tensor::from_vec(tokens.to_vec(), (b, 1), &state.device)?;
let model = state
.tp_models
.get_mut(&handle)
.ok_or_else(|| anyhow::anyhow!("TpForwardLogitsBatch: 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 = model.batch_decode_mask(prefix_lens, padded_len, total_len)?;
let logits = model.forward_batch_decode(&input, &positions, mask.as_ref())?;
let logits = logits.to_dtype(DType::F32)?;
let (rows, _, vocab) = logits.dims3()?;
anyhow::ensure!(
rows == b,
"TpForwardLogitsBatch: 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())
}
/// TP-equivalent of [`assemble_kv_batch`] against the TP slab (#98),
/// keyed by pool-minted snapshot ids.
#[cfg(feature = "cuda")]
fn tp_assemble_kv_batch(
state: &mut DeviceWorkerState,
handle: TpHandle,
seqs: &[(u64, usize)],
) -> anyhow::Result<usize> {
let DeviceWorkerState {
tp_models,
tp_kv_snapshots,
..
} = state;
let mut pairs = Vec::with_capacity(seqs.len());
for (id, len) in seqs {
let snap = tp_kv_snapshots.get(&(handle, *id)).ok_or_else(|| {
anyhow::anyhow!(
"TpAssembleKvBatch: no snapshot {id} for handle {}",
handle.0
)
})?;
pairs.push((snap, *len));
}
let batch = crate::harness::arch::qwen3_5::snapshot::assemble_batch(&pairs)?;
let model = tp_models
.get_mut(&handle)
.ok_or_else(|| anyhow::anyhow!("TpAssembleKvBatch: no model for handle {}", handle.0))?;
model.restore_kv_cache(&batch.snapshot)?;
Ok(batch.padded_len)
}
/// TP-equivalent of [`extract_kv_rows`] against the TP slab (#98),
/// storing each extracted row under the caller's pre-minted pool id.
/// Returns total snapshot bytes (budget accounting).
#[cfg(feature = "cuda")]
fn tp_extract_kv_rows(
state: &mut DeviceWorkerState,
handle: TpHandle,
rows: &[(usize, usize)],
padded_len: usize,
steps: usize,
snapshot_ids: &[u64],
) -> anyhow::Result<u64> {
anyhow::ensure!(
snapshot_ids.len() == rows.len(),
"TpExtractKvRows: {} rows vs {} snapshot_ids",
rows.len(),
snapshot_ids.len()
);
let live = state
.tp_models
.get(&handle)
.ok_or_else(|| anyhow::anyhow!("TpExtractKvRows: no model for handle {}", handle.0))?
.snapshot_kv_cache()?;
let mut total = 0u64;
for (&(row, prefix_len), &id) in rows.iter().zip(snapshot_ids) {
let snap = crate::harness::arch::qwen3_5::snapshot::extract_row(
&live, row, prefix_len, padded_len, steps,
)?;
total += snap.size_bytes();
state.tp_kv_snapshots.insert((handle, id), snap);
}
Ok(total)
}
/// Image-bearing leader forward (rank 0). Preprocesses each source
/// `image_data_uris` entry through the same deterministic
/// `preprocess_data_uri` every rank runs, uploads to the leader's
@@ -1034,6 +1234,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 +1500,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()));
}
@@ -1230,6 +1552,18 @@ fn drain_poisoned(job: Job, device_index: u32) {
let _ = reply.send(Err(err()));
}
#[cfg(feature = "cuda")]
Job::TpForwardLogitsBatch { reply, .. } => {
let _ = reply.send(Err(err()));
}
#[cfg(feature = "cuda")]
Job::TpAssembleKvBatch { reply, .. } => {
let _ = reply.send(Err(err()));
}
#[cfg(feature = "cuda")]
Job::TpExtractKvRows { reply, .. } => {
let _ = reply.send(Err(err()));
}
#[cfg(feature = "cuda")]
Job::TpDropKvSnapshot { reply, .. } => {
// Bookkeeping-only — unit reply so eviction never wedges
// on a poisoned worker (same shape as DropKvSnapshot).

View File

@@ -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.
///
@@ -304,6 +348,41 @@ pub enum Job {
offset: usize,
reply: oneshot::Sender<Result<Vec<f32>>>,
},
/// Leader half of a TP batched decode step (#98) — mirrors the
/// single-GPU `ForwardLogitsBatch` against the TP slab. The caller
/// (`WorkerPool::generate_step_batch`) fans out the matching
/// `GenerateStepBatch` RPC to subprocess ranks first so the
/// row-parallel collectives pair up.
#[cfg(feature = "cuda")]
TpForwardLogitsBatch {
handle: TpHandle,
tokens: Vec<u32>,
prefix_lens: Vec<usize>,
padded_len: usize,
step: usize,
reply: oneshot::Sender<Result<Vec<Vec<f32>>>>,
},
/// Leader half of a TP batch assembly (#98) — mirrors
/// `AssembleKvBatch` against the TP slab, with **pool-minted**
/// snapshot ids (the same ids every subprocess rank stores under).
#[cfg(feature = "cuda")]
TpAssembleKvBatch {
handle: TpHandle,
seqs: Vec<(u64, usize)>,
reply: oneshot::Sender<Result<usize>>,
},
/// Leader half of a TP row extraction (#98) — mirrors
/// `ExtractKvRows` against the TP slab, storing each extracted row
/// under the **pre-minted** id in `snapshot_ids` (one per row).
#[cfg(feature = "cuda")]
TpExtractKvRows {
handle: TpHandle,
rows: Vec<(usize, usize)>,
padded_len: usize,
steps: usize,
snapshot_ids: Vec<u64>,
reply: oneshot::Sender<Result<u64>>,
},
/// Image-bearing leader (rank 0) forward for the single-shot vision
/// prefill. The handler preprocesses each `image_data_uris` entry
/// (the same deterministic path every rank runs), encodes through

View File

@@ -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
@@ -769,6 +874,112 @@ impl DeviceWorkerHandle {
}
}
/// Leader half of a TP batched decode step (#98) — per-row
/// `[vocab]` logits back for per-slot sampling. The caller fans
/// the matching `GenerateStepBatch` out to the subprocess ranks.
#[cfg(feature = "cuda")]
pub async fn tp_forward_logits_batch(
&self,
handle: TpHandle,
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::TpForwardLogitsBatch {
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,
}),
}
}
/// Leader half of a TP batch assembly (#98). Returns the padded
/// uniform KV length. Snapshot ids are pool-minted by the caller.
#[cfg(feature = "cuda")]
pub async fn tp_assemble_kv_batch(
&self,
handle: TpHandle,
seqs: Vec<(u64, 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::TpAssembleKvBatch {
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,
}),
}
}
/// Leader half of a TP row extraction (#98). Stores each extracted
/// row under the caller's pre-minted pool id; returns total bytes.
#[cfg(feature = "cuda")]
pub async fn tp_extract_kv_rows(
&self,
handle: TpHandle,
rows: Vec<(usize, usize)>,
padded_len: usize,
steps: usize,
snapshot_ids: Vec<u64>,
) -> Result<u64, WorkerError> {
if self.poisoned.load(Ordering::Acquire) {
return Err(WorkerError::Poisoned {
device_index: self.device_index,
});
}
let (reply_tx, reply_rx) = oneshot::channel();
self.tx
.send(Job::TpExtractKvRows {
handle,
rows,
padded_len,
steps,
snapshot_ids,
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,
}),
}
}
/// Image-bearing TP leader forward (single-shot vision prefill).
/// Routes `Job::TpForwardLogitsWithImages` onto the worker thread;
/// the handler preprocesses + encodes + splices + forwards and
@@ -925,6 +1136,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");

File diff suppressed because it is too large Load Diff

View File

@@ -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;

View File

@@ -94,6 +94,37 @@ impl TpLeaderModel {
}
}
/// Lockstep batched decode step on rank 0 (#98). Only qwen3_5
/// batches (same gate as snapshots).
pub fn forward_batch_decode(
&mut self,
input: &candle_core::Tensor,
positions: &[usize],
attn_mask: Option<&candle_core::Tensor>,
) -> candle_core::Result<candle_core::Tensor> {
match self {
TpLeaderModel::Qwen3(_) => {
candle_core::bail!("forward_batch_decode: qwen3 (dense) has no batched decode")
}
TpLeaderModel::Qwen3_5(m) => m.forward_batch_decode(input, positions, attn_mask),
}
}
/// Padding mask for a batched decode step (#98).
pub fn batch_decode_mask(
&self,
prefix_lens: &[usize],
padded_len: usize,
total_len: usize,
) -> candle_core::Result<Option<candle_core::Tensor>> {
match self {
TpLeaderModel::Qwen3(_) => {
candle_core::bail!("batch_decode_mask: qwen3 (dense) has no batched decode")
}
TpLeaderModel::Qwen3_5(m) => m.batch_decode_mask(prefix_lens, padded_len, total_len),
}
}
/// Whether this arch supports prefix snapshots (#11). Gates the
/// pool fan-out so unsupported archs never even ask the ranks.
pub fn supports_kv_snapshot(&self) -> bool {
@@ -837,6 +868,179 @@ impl WorkerPool {
}
}
/// One lockstep batched decode step across every rank (#98).
/// Same fan-out / leader-forward / always-drain shape as
/// [`Self::generate_step`]; every rank derives positions + mask
/// locally from the broadcast geometry. Returns one `[vocab]`
/// logits row per batch row from the leader's rank-0 shard.
#[cfg(feature = "cuda")]
pub async fn generate_step_batch(
&mut self,
model_id: &str,
leader_handle: super::device_worker::TpHandle,
tokens: Vec<u32>,
prefix_lens: Vec<usize>,
padded_len: usize,
step: usize,
) -> Result<Vec<Vec<f32>>> {
for w in &mut self.workers {
w.send_only(&WorkerRequest::GenerateStepBatch {
model_id: model_id.to_string(),
tokens: tokens.clone(),
prefix_lens: prefix_lens.clone(),
padded_len,
step,
})
.await?;
}
let timeout = tp_step_timeout();
let leader_fut = self.leader_worker.tp_forward_logits_batch(
leader_handle,
tokens,
prefix_lens,
padded_len,
step,
);
let leader_result = match tokio::time::timeout(timeout, leader_fut).await {
Ok(r) => r,
Err(_elapsed) => {
// Watchdog (#17 Stage 2) — same rationale as
// `generate_step`: abort the wedged comm, fail without
// draining, let auto-recovery restart the pool.
self.watchdog_abort_leader_comm(model_id, timeout.as_secs());
anyhow::bail!(
"tp watchdog: leader batched forward exceeded {}s deadline; aborted wedged \
NCCL comm — model will auto-recover",
timeout.as_secs()
);
}
};
let worker_errors = drain_workers(&mut self.workers, |r| match r {
WorkerResponse::GenerateStepOk => Ok(()),
WorkerResponse::Error { kind, message } => Err(format!("[{kind}]: {message}")),
other => Err(format!("expected GenerateStepOk, got {other:?}")),
})
.await;
match leader_result {
Ok(rows) => {
if worker_errors.is_empty() {
Ok(rows)
} else {
anyhow::bail!(
"GenerateStepBatch: leader succeeded but workers failed: {}",
worker_errors.join("; ")
)
}
}
Err(e) => Err(anyhow::Error::new(e).context(if worker_errors.is_empty() {
"GenerateStepBatch: leader forward failed".to_string()
} else {
format!(
"GenerateStepBatch: leader forward failed and workers also failed: {}",
worker_errors.join("; ")
)
})),
}
}
/// Assemble stored per-sequence snapshots into every rank's live
/// batched state (#98). Snapshot ids in `seqs` are pool-minted;
/// every rank (leader + subprocesses) assembles the same geometry
/// and the returned padded length is asserted identical.
#[cfg(feature = "cuda")]
pub async fn assemble_kv_batch(
&mut self,
model_id: &str,
leader_handle: super::device_worker::TpHandle,
seqs: Vec<(u64, usize)>,
) -> Result<usize> {
for w in &mut self.workers {
w.send_only(&WorkerRequest::AssembleKvBatch {
model_id: model_id.to_string(),
seqs: seqs.clone(),
})
.await?;
}
let leader_result = self
.leader_worker
.tp_assemble_kv_batch(leader_handle, seqs)
.await;
let leader_padded = leader_result.as_ref().ok().copied();
let worker_errors = drain_workers(&mut self.workers, |r| match r {
WorkerResponse::KvBatchAssembled { padded_len } => {
if leader_padded.is_some_and(|lp| lp as u64 != padded_len) {
Err(format!(
"rank assembled padded_len {padded_len} != leader {}",
leader_padded.unwrap_or(0)
))
} else {
Ok(())
}
}
WorkerResponse::Error { kind, message } => Err(format!("[{kind}]: {message}")),
other => Err(format!("expected KvBatchAssembled, got {other:?}")),
})
.await;
let padded = leader_result.map_err(|e| {
anyhow::Error::new(e).context("AssembleKvBatch: leader assembly failed")
})?;
if !worker_errors.is_empty() {
anyhow::bail!(
"AssembleKvBatch: leader succeeded but workers failed: {}",
worker_errors.join("; ")
);
}
Ok(padded)
}
/// Extract live batched-state rows into per-sequence snapshots on
/// every rank (#98), stored under the pre-minted `snapshot_ids`
/// (one per row, minted from the pool's snapshot counter).
#[cfg(feature = "cuda")]
pub async fn extract_kv_rows(
&mut self,
model_id: &str,
leader_handle: super::device_worker::TpHandle,
rows: Vec<(usize, usize)>,
padded_len: usize,
steps: usize,
snapshot_ids: Vec<u64>,
) -> Result<()> {
for w in &mut self.workers {
w.send_only(&WorkerRequest::ExtractKvRows {
model_id: model_id.to_string(),
rows: rows.clone(),
padded_len,
steps,
snapshot_ids: snapshot_ids.clone(),
})
.await?;
}
let leader_result = self
.leader_worker
.tp_extract_kv_rows(leader_handle, rows, padded_len, steps, snapshot_ids)
.await;
let worker_errors = drain_workers(&mut self.workers, |r| match r {
WorkerResponse::KvRowsExtracted { .. } => Ok(()),
WorkerResponse::Error { kind, message } => Err(format!("[{kind}]: {message}")),
other => Err(format!("expected KvRowsExtracted, got {other:?}")),
})
.await;
leader_result.map_err(|e| {
anyhow::Error::new(e).context("ExtractKvRows: leader extraction failed")
})?;
if !worker_errors.is_empty() {
anyhow::bail!(
"ExtractKvRows: leader succeeded but workers failed: {}",
worker_errors.join("; ")
);
}
Ok(())
}
/// Image-bearing variant of [`Self::generate_step`] for the
/// single-shot vision prefill. Identical fan-out / leader-forward /
/// drain shape, but every rank runs the encode + splice path:

View File

@@ -115,6 +115,48 @@ pub enum WorkerRequest {
chunk_size: usize,
},
/// One lockstep batched decode step (#98): `tokens[i]` is batch
/// row i's next token at position `prefix_lens[i] + step`.
/// Identical on every rank (the leader mirrors it through
/// `Job::TpForwardLogitsBatch`); each rank derives the per-row
/// positions and padding mask locally and discards its logits,
/// same as `GenerateStep`.
GenerateStepBatch {
model_id: String,
tokens: Vec<u32>,
prefix_lens: Vec<usize>,
/// Uniform padded KV length the batch was assembled to.
padded_len: usize,
/// Decode steps since the last rebatch.
step: usize,
},
/// Assemble stored per-sequence snapshots into one batched cache
/// state and install it as this rank's live state (#98). `seqs`
/// pairs pool-minted snapshot ids with true token lengths; every
/// rank holds symmetric shard snapshots under the same ids, so
/// the assembled geometry (padded length) is identical across
/// ranks. Source snapshots remain stored. Replies
/// `KvBatchAssembled { padded_len }`.
AssembleKvBatch {
model_id: String,
seqs: Vec<(u64, usize)>,
},
/// Extract rows of this rank's live batched state back into
/// contiguous single-sequence snapshots (#98) — the first half of
/// a rebatch. `rows` pairs batch-row indexes with prefix lengths;
/// `snapshot_ids` carries the pool-minted id to store each
/// extracted row under (one per entry of `rows`, so every rank
/// keys identically). Replies `KvRowsExtracted`.
ExtractKvRows {
model_id: String,
rows: Vec<(usize, usize)>,
padded_len: usize,
steps: usize,
snapshot_ids: Vec<u64>,
},
/// Reset the KV cache for this model on this rank. Sent at the
/// start of every inference so a fresh request doesn't accidentally
/// attend over the previous one's tokens.
@@ -192,6 +234,14 @@ pub enum WorkerResponse {
/// Reply to `ClearKvCache`. Empty payload.
KvCacheCleared,
/// Reply to `AssembleKvBatch`. The uniform padded KV length the
/// batch was assembled to — the leader asserts every rank agrees.
KvBatchAssembled { padded_len: u64 },
/// Reply to `ExtractKvRows`. Total bytes of the extracted
/// snapshots on this rank (budget accounting only).
KvRowsExtracted { bytes: u64 },
/// Reply to `QueryVram`. This rank's device VRAM in MiB.
VramInfo { free_mb: u64, total_mb: u64 },

View File

@@ -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")`).

File diff suppressed because it is too large Load Diff

View File

@@ -82,6 +82,37 @@ impl WorkerModel {
}
}
/// Lockstep batched decode step (#98). Only qwen3_5 batches — the
/// leader gates on its own `TpLeaderModel` support first.
fn forward_batch_decode(
&mut self,
input: &candle_core::Tensor,
positions: &[usize],
attn_mask: Option<&candle_core::Tensor>,
) -> candle_core::Result<candle_core::Tensor> {
match self {
WorkerModel::Qwen3(_) => {
candle_core::bail!("forward_batch_decode: qwen3 (dense) has no batched decode")
}
WorkerModel::Qwen3_5(m) => m.forward_batch_decode(input, positions, attn_mask),
}
}
/// Padding mask for a batched decode step (#98).
fn batch_decode_mask(
&self,
prefix_lens: &[usize],
padded_len: usize,
total_len: usize,
) -> candle_core::Result<Option<candle_core::Tensor>> {
match self {
WorkerModel::Qwen3(_) => {
candle_core::bail!("batch_decode_mask: qwen3 (dense) has no batched decode")
}
WorkerModel::Qwen3_5(m) => m.batch_decode_mask(prefix_lens, padded_len, total_len),
}
}
/// Capture this rank's cache state for a prefix snapshot (#11).
/// Only qwen3_5 exposes its state; the dense qwen3 arch errors —
/// the leader never asks, because it gates on its own
@@ -248,6 +279,23 @@ impl WorkerState {
image_data_uris,
chunk_size,
),
WorkerRequest::GenerateStepBatch {
model_id,
tokens,
prefix_lens,
padded_len,
step,
} => self.handle_generate_step_batch(&model_id, tokens, prefix_lens, padded_len, step),
WorkerRequest::AssembleKvBatch { model_id, seqs } => {
self.handle_assemble_kv_batch(&model_id, &seqs)
}
WorkerRequest::ExtractKvRows {
model_id,
rows,
padded_len,
steps,
snapshot_ids,
} => self.handle_extract_kv_rows(&model_id, &rows, padded_len, steps, &snapshot_ids),
WorkerRequest::ClearKvCache { model_id } => self.handle_clear_kv_cache(&model_id),
WorkerRequest::SnapshotKvCache {
model_id,
@@ -405,8 +453,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 +487,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)"
),
};
}
@@ -538,6 +589,206 @@ impl WorkerState {
}
}
/// One lockstep batched decode step on this rank (#98). Mirrors
/// the leader's `forward_logits_batch`: derives per-row positions
/// and the padding mask locally from the broadcast geometry, runs
/// the batched forward, and discards the logits (the leader
/// samples from rank 0; the NCCL collectives are the point).
#[cfg(feature = "cuda")]
fn handle_generate_step_batch(
&mut self,
model_id: &str,
tokens: Vec<u32>,
prefix_lens: Vec<usize>,
padded_len: usize,
step: usize,
) -> WorkerResponse {
use candle_core::Tensor;
let b = tokens.len();
if b == 0 || prefix_lens.len() != b {
return WorkerResponse::Error {
kind: "bad_request".into(),
message: format!(
"GenerateStepBatch: {b} tokens vs {} prefix_lens",
prefix_lens.len()
),
};
}
let Some(model) = self.models.get_mut(model_id) else {
return WorkerResponse::Error {
kind: "model_not_loaded".into(),
message: format!("model '{model_id}' not loaded on rank {}", self.config.rank),
};
};
let device = model.device().clone();
let result = (|| -> candle_core::Result<()> {
let input = Tensor::from_vec(tokens, (b, 1), &device)?;
let positions: Vec<usize> = prefix_lens.iter().map(|&len| len + step).collect();
let total_len = padded_len + step + 1;
let mask = model.batch_decode_mask(&prefix_lens, padded_len, total_len)?;
model.forward_batch_decode(&input, &positions, mask.as_ref())?;
Ok(())
})();
match result {
Ok(()) => WorkerResponse::GenerateStepOk,
Err(e) => {
tracing::warn!(
rank = self.config.rank,
model = %model_id,
error = %e,
"worker GenerateStepBatch: forward failed"
);
WorkerResponse::Error {
kind: "forward_failed".into(),
message: format!("TP batched forward: {e}"),
}
}
}
}
#[cfg(not(feature = "cuda"))]
fn handle_generate_step_batch(
&mut self,
_model_id: &str,
_tokens: Vec<u32>,
_prefix_lens: Vec<usize>,
_padded_len: usize,
_step: usize,
) -> WorkerResponse {
WorkerResponse::Error {
kind: "cuda_feature_not_enabled".into(),
message: "GenerateStepBatch requires --features cuda".into(),
}
}
/// Assemble stored per-sequence snapshots into this rank's live
/// batched state (#98). Mirrors the leader's `assemble_kv_batch`.
#[cfg(feature = "cuda")]
fn handle_assemble_kv_batch(
&mut self,
model_id: &str,
seqs: &[(u64, usize)],
) -> WorkerResponse {
let Some(model) = self.models.get_mut(model_id) else {
return WorkerResponse::Error {
kind: "model_not_loaded".into(),
message: format!("model '{model_id}' not loaded on rank {}", self.config.rank),
};
};
let mut pairs = Vec::with_capacity(seqs.len());
for (id, len) in seqs {
let Some(snap) = self.kv_snapshots.get(&(model_id.to_string(), *id)) else {
return WorkerResponse::Error {
kind: "snapshot_not_found".into(),
message: format!(
"AssembleKvBatch: no snapshot {id} for '{model_id}' on rank {}",
self.config.rank
),
};
};
pairs.push((snap, *len));
}
let result =
crate::harness::arch::qwen3_5::snapshot::assemble_batch(&pairs).and_then(|batch| {
model.restore_kv_cache(&batch.snapshot)?;
Ok(batch.padded_len)
});
match result {
Ok(padded_len) => WorkerResponse::KvBatchAssembled {
padded_len: padded_len as u64,
},
Err(e) => WorkerResponse::Error {
kind: "assemble_failed".into(),
message: format!("AssembleKvBatch: {e}"),
},
}
}
#[cfg(not(feature = "cuda"))]
fn handle_assemble_kv_batch(
&mut self,
_model_id: &str,
_seqs: &[(u64, usize)],
) -> WorkerResponse {
WorkerResponse::Error {
kind: "cuda_feature_not_enabled".into(),
message: "AssembleKvBatch requires --features cuda".into(),
}
}
/// Extract live batched-state rows into stored per-sequence
/// snapshots under pool-minted ids (#98). Mirrors the leader's
/// `extract_kv_rows`.
#[cfg(feature = "cuda")]
fn handle_extract_kv_rows(
&mut self,
model_id: &str,
rows: &[(usize, usize)],
padded_len: usize,
steps: usize,
snapshot_ids: &[u64],
) -> WorkerResponse {
if snapshot_ids.len() != rows.len() {
return WorkerResponse::Error {
kind: "bad_request".into(),
message: format!(
"ExtractKvRows: {} rows vs {} snapshot_ids",
rows.len(),
snapshot_ids.len()
),
};
}
let Some(model) = self.models.get(model_id) else {
return WorkerResponse::Error {
kind: "model_not_loaded".into(),
message: format!("model '{model_id}' not loaded on rank {}", self.config.rank),
};
};
let live = match model.snapshot_kv_cache() {
Ok(s) => s,
Err(e) => {
return WorkerResponse::Error {
kind: "extract_failed".into(),
message: format!("ExtractKvRows: snapshot live state: {e}"),
};
}
};
let mut total_bytes = 0u64;
for (&(row, prefix_len), &id) in rows.iter().zip(snapshot_ids) {
match crate::harness::arch::qwen3_5::snapshot::extract_row(
&live, row, prefix_len, padded_len, steps,
) {
Ok(snap) => {
total_bytes += snap.size_bytes();
self.kv_snapshots.insert((model_id.to_string(), id), snap);
}
Err(e) => {
return WorkerResponse::Error {
kind: "extract_failed".into(),
message: format!("ExtractKvRows: row {row}: {e}"),
};
}
}
}
WorkerResponse::KvRowsExtracted { bytes: total_bytes }
}
#[cfg(not(feature = "cuda"))]
fn handle_extract_kv_rows(
&mut self,
_model_id: &str,
_rows: &[(usize, usize)],
_padded_len: usize,
_steps: usize,
_snapshot_ids: &[u64],
) -> WorkerResponse {
WorkerResponse::Error {
kind: "cuda_feature_not_enabled".into(),
message: "ExtractKvRows requires --features cuda".into(),
}
}
/// Image-bearing prefill on this rank. Preprocesses each source data
/// URI through the same deterministic `preprocess_data_uri` the
/// leader runs, encodes through this rank's replicated tower, and

View File

@@ -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`].
///

View File

@@ -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,
};

View File

@@ -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();

View File

@@ -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();

View 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
}

View File

@@ -0,0 +1,7 @@
{
"_from_model_config": true,
"output_attentions": false,
"output_hidden_states": false,
"transformers_version": "5.9.0",
"use_cache": true
}

Binary file not shown.

View 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"
}
}

View 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");
}

View 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>

View 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

View 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

View File

@@ -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`).

View File

@@ -0,0 +1,109 @@
# Prebuilt-binary spec for helexa-upstream.
#
# Wraps a pre-built `helexa-upstream` binary produced by an upstream CI
# job and packages it for rpm.lair.cafe. The %build phase is a no-op.
# helexa-upstream is a pure-Rust, non-CUDA daemon: the mesh account +
# budget authority. It serves an inbound HTTP API on tcp/8090 (/authz/v1
# for cortex, /web/v1 for the frontend) — hence the firewalld service —
# and connects out to PostgreSQL (the system of record), running schema
# migrations on startup.
#
# Required defines at rpmbuild time:
# upstream_version e.g. "0.1.16"
# upstream_prerelease e.g. "0.1.20260518140530.gitabcdef0"
# ^^^^^^^^^^^^^^^^^^ ^^^^^^^^
# commit time (sec) commit sha
# (used as Release; the timestamp prefix
# keeps same-day builds strictly ordered.)
%global _build_id_links none
%global debug_package %{nil}
%global __strip /usr/bin/true
%{!?upstream_version: %global upstream_version 0.0.0}
%if 0%{?upstream_prerelease:1}
%global upstream_release %{upstream_prerelease}
%else
%global upstream_release 1
%endif
Name: helexa-upstream
Version: %{upstream_version}
Release: %{upstream_release}%{?dist}
Summary: Mesh account + budget authority for helexa (prebuilt)
License: GPL-3.0-or-later
URL: https://git.lair.cafe/helexa/helexa
Source0: helexa-upstream
Source1: helexa-upstream.service
Source2: helexa-upstream-sysusers.conf
Source3: helexa-upstream.example.toml
Source4: LICENSE
Source5: helexa-upstream-firewalld.xml
Requires: firewalld-filesystem
ExclusiveArch: x86_64
Requires(pre): shadow-utils
Requires: systemd
Provides: user(helexa-upstream)
%description
helexa-upstream is the mesh-level authority: it issues user accounts and
API keys, holds the allocation ledger (free grant + redeemable top-up
codes), enforces per-key budgets via a transactional reserve/settle
contract, and reconciles served usage for operator compensation. cortex
validates locally-unrecognised keys against its /authz/v1 surface
(fail-closed); the helexa.ai frontend drives account self-service via
/web/v1. PostgreSQL is the system of record; the schema is migrated on
startup.
%prep
cp %{SOURCE0} ./helexa-upstream
cp %{SOURCE1} .
cp %{SOURCE2} .
cp %{SOURCE3} .
cp %{SOURCE4} .
cp %{SOURCE5} .
%build
# Already built in the upstream CI build job.
%install
install -Dm755 helexa-upstream %{buildroot}%{_bindir}/helexa-upstream
install -Dm644 helexa-upstream.service %{buildroot}%{_unitdir}/helexa-upstream.service
install -Dm644 helexa-upstream-sysusers.conf %{buildroot}%{_sysusersdir}/helexa-upstream.conf
install -Dm644 helexa-upstream-firewalld.xml %{buildroot}%{_prefix}/lib/firewalld/services/helexa-upstream.xml
install -dm755 %{buildroot}%{_sysconfdir}/helexa-upstream
install -Dm644 helexa-upstream.example.toml %{buildroot}%{_sysconfdir}/helexa-upstream/helexa-upstream.toml
%pre
getent group helexa-upstream >/dev/null || groupadd -r helexa-upstream
getent passwd helexa-upstream >/dev/null || \
useradd -r -g helexa-upstream -d /var/lib/helexa-upstream -s /sbin/nologin \
-c "helexa-upstream authority" helexa-upstream
%post
%systemd_post helexa-upstream.service
%preun
%systemd_preun helexa-upstream.service
%postun
%systemd_postun_with_restart helexa-upstream.service
%files
%license LICENSE
%{_bindir}/helexa-upstream
%{_unitdir}/helexa-upstream.service
%{_sysusersdir}/helexa-upstream.conf
%{_prefix}/lib/firewalld/services/helexa-upstream.xml
%dir %{_sysconfdir}/helexa-upstream
%config(noreplace) %{_sysconfdir}/helexa-upstream/helexa-upstream.toml
%changelog
* Mon Jun 23 2026 Gitea Actions <actions@git.lair.cafe> - %{upstream_version}-%{upstream_release}
- Prerelease build from upstream CI binary.

View File

@@ -0,0 +1,156 @@
#!/usr/bin/env python3
"""Synthesize a tiny qwen3_next parity fixture (#92).
Unlike script/dump_reference.py (which replays a real HF snapshot and
therefore needs the weights on disk), this builds a TINY random-weight
`Qwen3NextForCausalLM` from scratch, saves the checkpoint INTO the
fixture directory, runs the reference forward on fixed token ids, and
dumps the final-position logits. The whole fixture (weights included)
is a few hundred KB, so it is committed and the companion Rust test
(crates/neuron/tests/qwen3_next_parity.rs) runs everywhere — no env
var, no snapshot, CI included.
What this pins, exactly: neuron's qwen3_next wiring against upstream —
flat config normalisation, the `model.*` weight prefix, the
per-key-head-group de-interleave of the fused `in_proj_qkvz` /
`in_proj_ba` projections, hybrid layer interleaving, and the MoE block
(softmax→top-k→renorm routing, per-expert SwiGLU, shared expert +
sigmoid gate). The full-size 80B checkpoint differs only in dimensions.
The config mirrors the real 80B's *shape decisions* at doll-house
scale: interval-4 hybrid, 8 layers (so two full-attention layers),
every layer MoE (decoder_sparse_step 1), 16 experts / top-4 + shared
expert, partial rotary 0.25, 2 KV heads.
Usage (host with torch + transformers>=4.57, e.g. beast):
python3 script/dump_qwen3_next_tiny.py \
--out crates/neuron/tests/fixtures/numerical/qwen3_next-tiny
Regenerate whenever the transformers reference implementation changes;
record the transformers version from the manifest in the commit
message.
"""
import argparse
import json
import os
import struct
# ---------------------------------------------------------------------------
# Compat shim (same as dump_reference.py): transformers 5.9 constructs
# kernels-hub repository objects at import time without the
# revision/version that kernels 0.15 requires. The hub kernels are
# never used here; the constructors just must not throw.
os.environ.setdefault("USE_HUB_KERNELS", "NO")
try:
import kernels.layer.func as _kf
import kernels.layer.layer as _kl
def _patch(cls):
orig = cls.__init__
def patched(self, *a, **kw):
if "revision" not in kw and "version" not in kw:
kw["revision"] = "main"
orig(self, *a, **kw)
cls.__init__ = patched
_patch(_kl.LayerRepository)
_patch(_kf.FuncRepository)
except Exception: # noqa: BLE001 — older/newer kernels may not need it
pass
# ---------------------------------------------------------------------------
import torch # noqa: E402
# Fixed input: 96 token ids (> 64 so neuron's chunked delta-rule
# prefill path is exercised), deterministic arithmetic sequence folded
# into the tiny vocab.
TOKEN_IDS = [(7 * i + 3) % 512 for i in range(96)]
SEED = 92
def tiny_config():
from transformers import Qwen3NextConfig
return Qwen3NextConfig(
vocab_size=512,
hidden_size=64,
intermediate_size=128,
num_hidden_layers=8,
num_attention_heads=4,
num_key_value_heads=2,
head_dim=32,
max_position_embeddings=512,
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=16,
linear_num_key_heads=2,
linear_num_value_heads=4,
linear_value_head_dim=16,
decoder_sparse_step=1,
mlp_only_layers=[],
moe_intermediate_size=32,
norm_topk_prob=True,
num_experts=16,
num_experts_per_tok=4,
shared_expert_intermediate_size=32,
)
def write_f32(path, tensor):
data = tensor.detach().to(torch.float32).cpu().contiguous().reshape(-1)
with open(path, "wb") as f:
f.write(struct.pack(f"<{data.numel()}f", *data.tolist()))
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--out", required=True, help="fixture directory to write")
args = ap.parse_args()
import transformers
from transformers import Qwen3NextForCausalLM
os.makedirs(args.out, exist_ok=True)
torch.manual_seed(SEED)
cfg = tiny_config()
model = Qwen3NextForCausalLM(cfg).to(torch.float32).eval()
# Save the checkpoint into the fixture itself (config.json +
# model.safetensors) — the Rust test loads neuron's implementation
# from exactly these files.
model.save_pretrained(args.out, safe_serialization=True)
ids = torch.tensor([TOKEN_IDS], dtype=torch.long)
with torch.no_grad():
out = model(input_ids=ids)
logits = out.logits[0, -1] # final position, (vocab,)
write_f32(f"{args.out}/logits.f32", logits)
manifest = {
"case": "qwen3_next-tiny",
"seed": SEED,
"token_ids": TOKEN_IDS,
"files": {"logits": {"file": "logits.f32", "shape": [cfg.vocab_size]}},
"versions": {
"transformers": transformers.__version__,
"torch": torch.__version__,
},
}
with open(f"{args.out}/manifest.json", "w") as f:
json.dump(manifest, f, indent=2)
print(f"fixture written to {args.out}")
print(f" transformers {transformers.__version__}, torch {torch.__version__}")
print(f" logits[:4] = {logits[:4].tolist()}")
if __name__ == "__main__":
main()