Compare commits
56 Commits
4cb52e3144
...
fix/126-li
| Author | SHA1 | Date | |
|---|---|---|---|
|
98d095099a
|
|||
| 1773df34ad | |||
|
687324745f
|
|||
|
66632db578
|
|||
| abac5cb303 | |||
|
daeab6be12
|
|||
|
2a366a1c17
|
|||
|
4c3a2735b0
|
|||
| 7454ec7744 | |||
|
d3ef9513f8
|
|||
|
ec7a5b750a
|
|||
|
7f2fcc3527
|
|||
|
9ac1a1586c
|
|||
|
05b6c70a3a
|
|||
| 9c037c5a0b | |||
|
cb6ddbbd98
|
|||
| e824ea3274 | |||
|
9139b07d71
|
|||
|
4762537225
|
|||
| 9b635bfe5c | |||
|
28aa06b768
|
|||
| fc95faaed6 | |||
|
2df829ed73
|
|||
| f611e3479b | |||
|
41ec05e40c
|
|||
|
3242cce39b
|
|||
| 35041f11de | |||
|
40fecb2cf7
|
|||
| 8d50918a1a | |||
|
29a7054c23
|
|||
|
7489a49d24
|
|||
| 341f97f705 | |||
|
fc56eafbf7
|
|||
|
a157c8fa27
|
|||
|
201cc54a7b
|
|||
| 781d92d9a3 | |||
|
9bf13f09dd
|
|||
|
63fb8c6e76
|
|||
|
a1426f177c
|
|||
| a9f5ac7ab4 | |||
|
0126144a93
|
|||
|
869033d08e
|
|||
|
7c5011e5d3
|
|||
|
d1366614f9
|
|||
| 1884535542 | |||
|
140f9124ee
|
|||
| 3b77dbfa32 | |||
|
842caf5d4d
|
|||
| a488ade675 | |||
|
d472b6428a
|
|||
| 0488768afc | |||
|
209bf58207
|
|||
| 1ea7425585 | |||
|
a415454962
|
|||
| 06a36566d1 | |||
|
afc1f7a706
|
@@ -87,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"
|
||||
@@ -377,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
|
||||
|
||||
@@ -82,7 +82,9 @@ jobs:
|
||||
# see commit history).
|
||||
cuda-check:
|
||||
name: CUDA type-check
|
||||
timeout-minutes: 35
|
||||
# flash-attn kernel compilation dominates the first uncached run;
|
||||
# sccache + the cargo target cache absorb it afterwards.
|
||||
timeout-minutes: 70
|
||||
runs-on: cuda-13.0
|
||||
# The workflow-level env sets `RUSTC_WRAPPER: sccache`
|
||||
# unconditionally, which hard-fails cargo if the CUDA image
|
||||
@@ -115,7 +117,7 @@ jobs:
|
||||
export PATH="/usr/local/cuda-13.0/bin:${PATH}"
|
||||
export LD_LIBRARY_PATH="/usr/local/cuda-13.0/targets/x86_64-linux/lib:/usr/local/cuda-13.0/lib64:${LD_LIBRARY_PATH:-}"
|
||||
export LIBRARY_PATH="/usr/local/cuda-13.0/targets/x86_64-linux/lib:/usr/local/cuda-13.0/lib64:${LIBRARY_PATH:-}"
|
||||
script/ci-cargo-escalate.sh cargo check -p neuron --features cuda --all-targets
|
||||
script/ci-cargo-escalate.sh cargo check -p neuron --features cuda,flash-attn --all-targets
|
||||
|
||||
srpm-cortex:
|
||||
name: Build cortex SRPM
|
||||
|
||||
@@ -274,19 +274,52 @@ jobs:
|
||||
|
||||
echo "LLM probe against ${model}"
|
||||
probe_body=$(printf '{"model":"%s","messages":[{"role":"user","content":"Reply with exactly one word: pineapple"}],"max_tokens":512,"temperature":0}' "${model}")
|
||||
resp=$(curl -fsS --max-time 180 -H "content-type: application/json" \
|
||||
-d "${probe_body}" http://localhost:13131/v1/chat/completions) || {
|
||||
echo "FAIL: probe request errored"
|
||||
exit 1
|
||||
}
|
||||
if printf %s "${resp}" | grep -qi pineapple; then
|
||||
echo "LLM probe passed"
|
||||
else
|
||||
echo "FAIL: probe response missing expected token"
|
||||
printf %s "${resp}" | head -c 2000
|
||||
echo
|
||||
exit 1
|
||||
fi
|
||||
# The probe races real traffic: a deploy publishes a new build
|
||||
# SHA, which is exactly what triggers helexa-bench to re-sweep
|
||||
# every scenario against this neuron — long capability-probe
|
||||
# generations can hold the batch-1 model for minutes. Admission
|
||||
# control answers concurrent requests with 429/503 +
|
||||
# Retry-After (#53/#63); treat those as "busy, try again", not
|
||||
# deploy failure. Overall budget ~6 min.
|
||||
probe_deadline=$(( $(date +%s) + 360 ))
|
||||
attempt=0
|
||||
while :; do
|
||||
attempt=$(( attempt + 1 ))
|
||||
hdrs=$(mktemp)
|
||||
resp=$(curl -sS --max-time 180 -D "${hdrs}" \
|
||||
-H "content-type: application/json" \
|
||||
-d "${probe_body}" http://localhost:13131/v1/chat/completions) || resp=""
|
||||
status=$(awk 'toupper($1) ~ /^HTTP/ {code=$2} END {print code}' "${hdrs}")
|
||||
retry_after=$(awk 'tolower($1) == "retry-after:" {print $2+0; exit}' "${hdrs}")
|
||||
rm -f "${hdrs}"
|
||||
if [ "${status}" = "200" ]; then
|
||||
if printf %s "${resp}" | grep -qi pineapple; then
|
||||
echo "LLM probe passed (attempt ${attempt})"
|
||||
break
|
||||
fi
|
||||
echo "FAIL: probe response missing expected token"
|
||||
printf %s "${resp}" | head -c 2000
|
||||
echo
|
||||
exit 1
|
||||
fi
|
||||
case "${status}" in
|
||||
429|503) ;; # busy/queue-full — retryable
|
||||
*)
|
||||
echo "FAIL: probe request errored (HTTP ${status:-none})"
|
||||
printf %s "${resp}" | head -c 2000
|
||||
echo
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
if [ "$(date +%s)" -ge "${probe_deadline}" ]; then
|
||||
echo "FAIL: model still busy (HTTP ${status}) after probe budget"
|
||||
exit 1
|
||||
fi
|
||||
wait_s=${retry_after:-15}
|
||||
[ "${wait_s}" -ge 5 ] 2>/dev/null || wait_s=15
|
||||
echo "model busy (HTTP ${status}), retrying in ${wait_s}s (attempt ${attempt})"
|
||||
sleep "${wait_s}"
|
||||
done
|
||||
DEPLOY
|
||||
|
||||
- name: Ensure firewalld allows helexa-neuron
|
||||
|
||||
1
Cargo.lock
generated
1
Cargo.lock
generated
@@ -2972,6 +2972,7 @@ dependencies = [
|
||||
"axum",
|
||||
"base64 0.22.1",
|
||||
"candle-core",
|
||||
"candle-flash-attn",
|
||||
"candle-nn",
|
||||
"candle-transformers",
|
||||
"clap",
|
||||
|
||||
@@ -18,6 +18,35 @@ db_path = "/var/lib/helexa-bench/bench.sqlite"
|
||||
[scenarios]
|
||||
prompt_sizes = [128, 4096]
|
||||
max_tokens = 256
|
||||
# Concurrency / agentic-load scenarios (#89), enabled so the 27B baseline
|
||||
# carries p95-under-concurrency data before the F3 A/B gate (#94) needs
|
||||
# it for comparison. Levels mirror the real a0/hermes/opencode fan-out.
|
||||
concurrency_levels = [2, 4, 8]
|
||||
concurrency_prompt_tokens = 512
|
||||
|
||||
# Capability probes (#91) — the reasoning/planning axis the speed
|
||||
# scenarios miss; scored manually via `helexa-bench score` (O7). Enabled
|
||||
# for the same reason: the F3 gate compares 80B-A3B variants against the
|
||||
# 27B on planning quality, so the 27B needs scored artifacts first.
|
||||
[[scenarios.capability_probes]]
|
||||
name = "rust-plan"
|
||||
max_tokens = 4096
|
||||
prompt = """
|
||||
Write an implementation plan for adding rate limiting to an Axum service.
|
||||
Honor existing conventions, call out trade-offs, and sequence the work.
|
||||
"""
|
||||
|
||||
[[scenarios.capability_probes]]
|
||||
name = "debug-reason"
|
||||
max_tokens = 4096
|
||||
prompt = """
|
||||
A Rust axum server streams SSE responses through a reverse proxy. Clients
|
||||
report that streams stall for exactly 60 seconds and then resume, but only
|
||||
when response chunks are small and infrequent. Curling the backend directly
|
||||
never stalls. List the most likely causes in order of probability, explain
|
||||
the mechanism behind each, and describe the smallest experiment that would
|
||||
confirm or eliminate each cause.
|
||||
"""
|
||||
|
||||
# Read-only JSON API consumed by the bench UI (hosted separately) and for
|
||||
# programmatic access. Served alongside the sweep loop.
|
||||
|
||||
@@ -16,6 +16,14 @@ name = "candle"
|
||||
|
||||
[harness.candle]
|
||||
|
||||
# Batched decode engine (#98): up to 8 concurrent text streams
|
||||
# multiplex through one lockstep (B,1) forward per decode step.
|
||||
# max_in_flight is both the admission bound and the engine's slot
|
||||
# count. NEURON_BATCHING=0 (systemd drop-in) is the kill switch;
|
||||
# removing this section reverts to batch-1 serialization.
|
||||
[harness.candle.admission]
|
||||
max_in_flight = 8
|
||||
|
||||
[[default_models]]
|
||||
model_id = "Qwen/Qwen3.6-27B"
|
||||
harness = "candle"
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -35,6 +35,9 @@ pub fn api_routes(state: ApiState) -> Router {
|
||||
.route("/api/health", get(health))
|
||||
.route("/api/dimensions", get(dimensions))
|
||||
.route("/api/summary", get(summary))
|
||||
.route("/api/scaling", get(scaling))
|
||||
.route("/api/swap", get(swap))
|
||||
.route("/api/capability", get(capability))
|
||||
.route("/api/series", get(series))
|
||||
.route("/api/runs", get(runs))
|
||||
.layer(CorsLayer::permissive())
|
||||
@@ -73,6 +76,30 @@ async fn summary(
|
||||
store.summary().map(Json).map_err(err500)
|
||||
}
|
||||
|
||||
/// Context-length scaling curves per (target, model) — prefill & decode
|
||||
/// tok/s vs context, with decode-flatness (#88).
|
||||
async fn scaling(
|
||||
State(s): State<ApiState>,
|
||||
) -> Result<Json<Vec<crate::store::ScalingCurve>>, ApiError> {
|
||||
let store = s.lock().await;
|
||||
store.scaling().map(Json).map_err(err500)
|
||||
}
|
||||
|
||||
/// Cold-load / model-swap costs per (target, model) — reload latency + cold
|
||||
/// first-request (#90).
|
||||
async fn swap(State(s): State<ApiState>) -> Result<Json<Vec<crate::store::SwapCost>>, ApiError> {
|
||||
let store = s.lock().await;
|
||||
store.swap_costs().map(Json).map_err(err500)
|
||||
}
|
||||
|
||||
/// Capability-probe runs — stored artifacts + quality scores (#91).
|
||||
async fn capability(
|
||||
State(s): State<ApiState>,
|
||||
) -> Result<Json<Vec<crate::store::CapabilityRun>>, ApiError> {
|
||||
let store = s.lock().await;
|
||||
store.capability_runs(false).map(Json).map_err(err500)
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct SeriesQuery {
|
||||
/// Optional — when omitted the store resolves the host serving this model.
|
||||
|
||||
@@ -3,10 +3,10 @@
|
||||
//! `openai` targets use the OpenAI-compatible surface (preliminary).
|
||||
|
||||
use crate::config::{TargetConfig, TargetKind};
|
||||
use anyhow::{Context, Result};
|
||||
use anyhow::{Context, Result, anyhow};
|
||||
use cortex_core::build_info::BuildInfo;
|
||||
use cortex_core::discovery::DiscoveryResponse;
|
||||
use cortex_core::harness::ModelInfo;
|
||||
use cortex_core::discovery::{DiscoveryResponse, HealthResponse};
|
||||
use cortex_core::harness::{ModelInfo, ModelSpec};
|
||||
use cortex_core::openai::ModelsResponse;
|
||||
use std::time::Duration;
|
||||
|
||||
@@ -94,6 +94,84 @@ impl TargetClient {
|
||||
Ok(Some(disco))
|
||||
}
|
||||
|
||||
/// Runtime device health (neuron only): per-GPU VRAM used/free,
|
||||
/// utilization, and temperature from `GET /health`. Bench samples this
|
||||
/// around each measured run to record VRAM high-water + GPU telemetry
|
||||
/// (#87). Returns `Ok(None)` for non-neuron targets; a soft `Ok(None)`
|
||||
/// (not an error) on transport failure so a flaky `/health` never fails
|
||||
/// a measurement.
|
||||
pub async fn fetch_health(&self, target: &TargetConfig) -> Result<Option<HealthResponse>> {
|
||||
if target.kind != TargetKind::Neuron {
|
||||
return Ok(None);
|
||||
}
|
||||
let base = target.endpoint.trim_end_matches('/');
|
||||
let health = self
|
||||
.http
|
||||
.get(format!("{base}/health"))
|
||||
.timeout(META_TIMEOUT)
|
||||
.send()
|
||||
.await
|
||||
.context("GET /health")?
|
||||
.error_for_status()
|
||||
.context("GET /health status")?
|
||||
.json::<HealthResponse>()
|
||||
.await
|
||||
.context("decoding /health")?;
|
||||
Ok(Some(health))
|
||||
}
|
||||
|
||||
/// Unload a model (neuron only): `POST /models/unload {model_id}`.
|
||||
/// Used by the deliberate swap-cost measurement (#90), never the sweep.
|
||||
pub async fn unload_model(&self, target: &TargetConfig, model_id: &str) -> Result<()> {
|
||||
let base = target.endpoint.trim_end_matches('/');
|
||||
self.http
|
||||
.post(format!("{base}/models/unload"))
|
||||
.json(&serde_json::json!({ "model_id": model_id }))
|
||||
.send()
|
||||
.await
|
||||
.context("POST /models/unload")?
|
||||
.error_for_status()
|
||||
.context("POST /models/unload status")?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Load a model from a spec (neuron only): `POST /models/load`. neuron
|
||||
/// returns synchronously once loaded, so the call duration is the reload
|
||||
/// cost the swap-cost measurement records (#90).
|
||||
pub async fn load_model(&self, target: &TargetConfig, spec: &ModelSpec) -> Result<()> {
|
||||
let base = target.endpoint.trim_end_matches('/');
|
||||
self.http
|
||||
.post(format!("{base}/models/load"))
|
||||
.json(spec)
|
||||
// A cold load can take tens of seconds; use the full request
|
||||
// timeout rather than the short metadata one.
|
||||
.send()
|
||||
.await
|
||||
.context("POST /models/load")?
|
||||
.error_for_status()
|
||||
.context("POST /models/load status")?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Reconstruct a reload [`ModelSpec`] from a model's `/models` entry.
|
||||
/// Tensor-parallel is inferred from the device count; `quant` is left
|
||||
/// `None` for neuron to resolve from the catalogue / its prior load.
|
||||
pub fn spec_from_info(info: &ModelInfo) -> Result<ModelSpec> {
|
||||
if info.devices.is_empty() {
|
||||
return Err(anyhow!(
|
||||
"model '{}' reports no devices; cannot reconstruct a load spec",
|
||||
info.id
|
||||
));
|
||||
}
|
||||
Ok(ModelSpec {
|
||||
model_id: info.id.clone(),
|
||||
harness: info.harness.clone(),
|
||||
quant: None,
|
||||
tensor_parallel: (info.devices.len() > 1).then_some(info.devices.len() as u32),
|
||||
devices: Some(info.devices.clone()),
|
||||
})
|
||||
}
|
||||
|
||||
/// Warm models — those ready to serve without a cold load.
|
||||
///
|
||||
/// Neuron: `GET /models` filtered to `status == "loaded"` (skips
|
||||
|
||||
@@ -104,6 +104,37 @@ pub struct ScenarioConfig {
|
||||
/// Max generated tokens per request.
|
||||
#[serde(default = "default_max_tokens")]
|
||||
pub max_tokens: u64,
|
||||
/// Concurrency levels (#89): one `concurrency:<n>` scenario per entry,
|
||||
/// each firing N simultaneous streams. Defaults to empty (opt-in) — a
|
||||
/// burst puts real load on a serving fleet, so operators enable it
|
||||
/// deliberately, e.g. `concurrency_levels = [2, 4, 8]`.
|
||||
#[serde(default)]
|
||||
pub concurrency_levels: Vec<u32>,
|
||||
/// Approximate prompt size (tokens) used by the concurrency scenarios.
|
||||
#[serde(default = "default_concurrency_prompt_tokens")]
|
||||
pub concurrency_prompt_tokens: u32,
|
||||
/// Capability probes (#91): one `capability:<name>` scenario per entry,
|
||||
/// each running a fixed prompt and storing the full output artifact for
|
||||
/// quality scoring (manual now, LLM-judge later). Defaults to empty
|
||||
/// (opt-in) — these generate long outputs and exist to compare reasoning
|
||||
/// quality across models, not to run on every sweep by default.
|
||||
#[serde(default)]
|
||||
pub capability_probes: Vec<CapabilityProbe>,
|
||||
}
|
||||
|
||||
/// One capability probe: a named prompt whose output is stored and scored
|
||||
/// for quality (#91). The probe is deterministic (temperature 0) so the
|
||||
/// same model+build produces a stable artifact to score.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct CapabilityProbe {
|
||||
/// Stable id fragment — the scenario becomes `capability:<name>`.
|
||||
pub name: String,
|
||||
/// The prompt sent verbatim as the user message.
|
||||
pub prompt: String,
|
||||
/// Generation budget for the probe (planning answers are long; the
|
||||
/// default 256 is too small). Defaults to 2048.
|
||||
#[serde(default = "default_capability_max_tokens")]
|
||||
pub max_tokens: u64,
|
||||
}
|
||||
|
||||
impl Default for ScenarioConfig {
|
||||
@@ -111,6 +142,9 @@ impl Default for ScenarioConfig {
|
||||
ScenarioConfig {
|
||||
prompt_sizes: default_prompt_sizes(),
|
||||
max_tokens: default_max_tokens(),
|
||||
concurrency_levels: Vec::new(),
|
||||
concurrency_prompt_tokens: default_concurrency_prompt_tokens(),
|
||||
capability_probes: Vec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -187,6 +221,12 @@ fn default_prompt_sizes() -> Vec<u32> {
|
||||
fn default_max_tokens() -> u64 {
|
||||
256
|
||||
}
|
||||
fn default_concurrency_prompt_tokens() -> u32 {
|
||||
512
|
||||
}
|
||||
fn default_capability_max_tokens() -> u64 {
|
||||
2048
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
// Jail's closure must return figment::Result; the large-Err type is
|
||||
|
||||
@@ -43,6 +43,33 @@ enum Command {
|
||||
#[arg(short, long, default_value = "helexa-bench.toml")]
|
||||
config: String,
|
||||
},
|
||||
/// Measure cold-load / model-swap cost (#90): for each neuron target's
|
||||
/// warm models, unload → time reload → time a cold first request, recorded
|
||||
/// under scenario "swap". DELIBERATE — takes each model offline for its
|
||||
/// reload, so run it in a maintenance window, not against live traffic.
|
||||
SwapCost {
|
||||
#[arg(short, long, default_value = "helexa-bench.toml")]
|
||||
config: String,
|
||||
},
|
||||
/// Attach a quality score to a capability-probe run (#91). Find run ids
|
||||
/// with `report --capability`. `--scorer` records who scored it
|
||||
/// (defaults to "manual"); a future LLM-judge would set e.g. "llm:…".
|
||||
Score {
|
||||
#[arg(short, long, default_value = "helexa-bench.toml")]
|
||||
config: String,
|
||||
/// Override the SQLite path (skips reading the config file).
|
||||
#[arg(long)]
|
||||
db: Option<String>,
|
||||
/// The run id to score.
|
||||
#[arg(long)]
|
||||
id: i64,
|
||||
/// The quality score to attach (scale is the operator's rubric).
|
||||
#[arg(long)]
|
||||
score: f64,
|
||||
/// Who/what produced the score.
|
||||
#[arg(long, default_value = "manual")]
|
||||
scorer: String,
|
||||
},
|
||||
/// Render recorded results. Uses `--db` if given, else the db_path
|
||||
/// from `--config`.
|
||||
Report {
|
||||
@@ -54,6 +81,19 @@ enum Command {
|
||||
/// Output format.
|
||||
#[arg(long, default_value = "md")]
|
||||
format: Format,
|
||||
/// Render the context-length scaling view (prefill & decode tok/s
|
||||
/// vs context per model, with decode-flatness) instead of the flat
|
||||
/// results table (#88).
|
||||
#[arg(long)]
|
||||
scaling: bool,
|
||||
/// Render the cold-load / model-swap cost view (#90) instead of the
|
||||
/// flat results table.
|
||||
#[arg(long)]
|
||||
swap: bool,
|
||||
/// Render the capability-probe view (#91): stored artifacts + quality
|
||||
/// scores, with per-model median.
|
||||
#[arg(long)]
|
||||
capability: bool,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -122,16 +162,79 @@ async fn run(cli: Cli) -> Result<()> {
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Command::Report { config, db, format } => {
|
||||
Command::SwapCost { config } => {
|
||||
let cfg = load_config(&config)?;
|
||||
require_targets(&cfg)?;
|
||||
let sweeper = Sweeper::new(cfg)?;
|
||||
tracing::warn!(
|
||||
"swap-cost: cycling each warm model (unload → reload → cold request); models go offline during reload"
|
||||
);
|
||||
let summary = sweeper.swap_cost_once().await?;
|
||||
tracing::info!(
|
||||
measured = summary.measured,
|
||||
failed = summary.failed,
|
||||
unreachable = summary.targets_unreachable,
|
||||
"swap-cost measurement complete"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Command::Score {
|
||||
config,
|
||||
db,
|
||||
id,
|
||||
score,
|
||||
scorer,
|
||||
} => {
|
||||
let db_path = match db {
|
||||
Some(p) => p,
|
||||
None => load_config(&config)?.bench.db_path,
|
||||
};
|
||||
let store = Store::open(&db_path)?;
|
||||
let rows = store.report_rows()?;
|
||||
let rendered = match format {
|
||||
Format::Md => report::render_markdown(&rows),
|
||||
Format::Json => report::render_json(&rows)?,
|
||||
match store.set_score(id, score, &scorer)? {
|
||||
0 => anyhow::bail!("no run with id {id}"),
|
||||
_ => {
|
||||
println!("scored run {id}: {score} ({scorer})");
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
Command::Report {
|
||||
config,
|
||||
db,
|
||||
format,
|
||||
scaling,
|
||||
swap,
|
||||
capability,
|
||||
} => {
|
||||
let db_path = match db {
|
||||
Some(p) => p,
|
||||
None => load_config(&config)?.bench.db_path,
|
||||
};
|
||||
let store = Store::open(&db_path)?;
|
||||
let rendered = if capability {
|
||||
let runs = store.capability_runs(false)?;
|
||||
match format {
|
||||
Format::Md => report::render_capability_markdown(&runs),
|
||||
Format::Json => report::render_capability_json(&runs)?,
|
||||
}
|
||||
} else if swap {
|
||||
let costs = store.swap_costs()?;
|
||||
match format {
|
||||
Format::Md => report::render_swap_markdown(&costs),
|
||||
Format::Json => report::render_swap_json(&costs)?,
|
||||
}
|
||||
} else if scaling {
|
||||
let curves = store.scaling()?;
|
||||
match format {
|
||||
Format::Md => report::render_scaling_markdown(&curves),
|
||||
Format::Json => report::render_scaling_json(&curves)?,
|
||||
}
|
||||
} else {
|
||||
let rows = store.report_rows()?;
|
||||
match format {
|
||||
Format::Md => report::render_markdown(&rows),
|
||||
Format::Json => report::render_json(&rows)?,
|
||||
}
|
||||
};
|
||||
println!("{rendered}");
|
||||
Ok(())
|
||||
|
||||
@@ -3,28 +3,36 @@
|
||||
//! doc: engine, model, prompt tok, TTFT (s), decode tok/s, total (s),
|
||||
//! plus the build SHA each cell was measured against.
|
||||
|
||||
use crate::store::ReportRow;
|
||||
use crate::store::{CapabilityRun, ReportRow, ScalingCurve, SwapCost};
|
||||
use anyhow::Result;
|
||||
|
||||
pub fn render_markdown(rows: &[ReportRow]) -> String {
|
||||
let mut out = String::new();
|
||||
out.push_str(
|
||||
"| engine | model | prompt tok | TTFT (s) | decode tok/s | total (s) | build | n |\n",
|
||||
"| engine | model | prompt tok | prefill tok/s | TTFT (s) | TTFT p95 | \
|
||||
decode tok/s | total (s) | total p95 | VRAM (GB) | conc | queue ms | rej | build | n |\n",
|
||||
);
|
||||
out.push_str("|---|---|---:|---:|---:|---:|---|---:|\n");
|
||||
out.push_str("|---|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---|---:|\n");
|
||||
for r in rows {
|
||||
let ptok = r
|
||||
.prompt_tokens
|
||||
.map(|t| t.to_string())
|
||||
.unwrap_or_else(|| format!("~{}", r.prompt_size_approx));
|
||||
out.push_str(&format!(
|
||||
"| {} | {} | {} | {} | {} | {} | `{}` | {} |\n",
|
||||
"| {} | {} | {} | {} | {} | {} | {} | {} | {} | {} | {} | {} | {} | `{}` | {} |\n",
|
||||
r.target_name,
|
||||
r.model_id,
|
||||
ptok,
|
||||
fmt_opt(r.prefill_tps_median, 1),
|
||||
fmt_opt(r.ttft_s_median, 3),
|
||||
fmt_opt(r.ttft_s_p95, 3),
|
||||
fmt_opt(r.decode_tps_median, 1),
|
||||
fmt_opt(r.total_s_median, 3),
|
||||
fmt_opt(r.total_s_p95, 3),
|
||||
fmt_vram(r.vram_used_mb_median, r.vram_total_mb),
|
||||
fmt_u64(r.concurrency),
|
||||
fmt_opt(r.queue_wait_ms_median, 0),
|
||||
fmt_opt(r.rejected_median, 0),
|
||||
r.git_sha,
|
||||
r.samples,
|
||||
));
|
||||
@@ -43,8 +51,23 @@ pub fn render_json(rows: &[ReportRow]) -> Result<String> {
|
||||
"prompt_size_approx": r.prompt_size_approx,
|
||||
"prompt_tokens": r.prompt_tokens,
|
||||
"ttft_s_median": r.ttft_s_median,
|
||||
"ttft_s_p95": r.ttft_s_p95,
|
||||
"ttft_s_p99": r.ttft_s_p99,
|
||||
"decode_tps_median": r.decode_tps_median,
|
||||
"total_s_median": r.total_s_median,
|
||||
"total_s_p95": r.total_s_p95,
|
||||
"total_s_p99": r.total_s_p99,
|
||||
"prefill_ms_median": r.prefill_ms_median,
|
||||
"decode_ms_median": r.decode_ms_median,
|
||||
"prefill_tps_median": r.prefill_tps_median,
|
||||
"vram_used_mb_median": r.vram_used_mb_median,
|
||||
"vram_total_mb": r.vram_total_mb,
|
||||
"gpu_util_pct_median": r.gpu_util_pct_median,
|
||||
"gpu_temp_c_median": r.gpu_temp_c_median,
|
||||
"concurrency": r.concurrency,
|
||||
"ttft_p95_load_s": r.ttft_p95_load_s,
|
||||
"queue_wait_ms_median": r.queue_wait_ms_median,
|
||||
"rejected_median": r.rejected_median,
|
||||
"git_sha": r.git_sha,
|
||||
"samples": r.samples,
|
||||
"gpu": r.gpu,
|
||||
@@ -54,6 +77,179 @@ pub fn render_json(rows: &[ReportRow]) -> Result<String> {
|
||||
Ok(serde_json::to_string_pretty(&arr)?)
|
||||
}
|
||||
|
||||
/// Context-length scaling view (#88): one block per (target, model) with
|
||||
/// prefill & decode tok/s vs context, then the decode-flatness verdict.
|
||||
pub fn render_scaling_markdown(curves: &[ScalingCurve]) -> String {
|
||||
let mut out = String::new();
|
||||
for c in curves {
|
||||
let gpu = c.gpu.as_deref().unwrap_or("");
|
||||
out.push_str(&format!(
|
||||
"### {} · {} (`{}`{})\n\n",
|
||||
c.target_name,
|
||||
c.model_id,
|
||||
c.git_sha,
|
||||
if gpu.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
format!(", {gpu}")
|
||||
},
|
||||
));
|
||||
out.push_str("| ctx tok | prefill tok/s | decode tok/s | n |\n");
|
||||
out.push_str("|---:|---:|---:|---:|\n");
|
||||
for p in &c.points {
|
||||
let ctx = p
|
||||
.prompt_tokens
|
||||
.map(|t| t.to_string())
|
||||
.unwrap_or_else(|| format!("~{}", p.prompt_size));
|
||||
out.push_str(&format!(
|
||||
"| {} | {} | {} | {} |\n",
|
||||
ctx,
|
||||
fmt_opt(p.prefill_tps, 1),
|
||||
fmt_opt(p.decode_tps, 1),
|
||||
p.samples,
|
||||
));
|
||||
}
|
||||
match c.decode_flatness {
|
||||
Some(f) => out.push_str(&format!(
|
||||
"\ndecode flatness: {f:.2} — decode tok/s {} across the context range \
|
||||
({})\n\n",
|
||||
if f >= 0.9 {
|
||||
"holds"
|
||||
} else if f >= 0.7 {
|
||||
"softens"
|
||||
} else {
|
||||
"drops sharply"
|
||||
},
|
||||
if f >= 0.9 {
|
||||
"Gated-DeltaNet O(1) decode confirmed"
|
||||
} else {
|
||||
"investigate where it breaks"
|
||||
},
|
||||
)),
|
||||
None => out.push_str("\ndecode flatness: — (need ≥2 context points)\n\n"),
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
pub fn render_scaling_json(curves: &[ScalingCurve]) -> Result<String> {
|
||||
Ok(serde_json::to_string_pretty(curves)?)
|
||||
}
|
||||
|
||||
/// Cold-load / model-swap cost view (#90): reload latency + cold
|
||||
/// first-request per model.
|
||||
pub fn render_swap_markdown(costs: &[SwapCost]) -> String {
|
||||
let mut out = String::new();
|
||||
out.push_str(
|
||||
"| engine | model | unload (s) | reload (s) | cold TTFT (s) | cold total (s) | build | n |\n",
|
||||
);
|
||||
out.push_str("|---|---|---:|---:|---:|---:|---|---:|\n");
|
||||
for c in costs {
|
||||
out.push_str(&format!(
|
||||
"| {} | {} | {} | {} | {} | {} | `{}` | {} |\n",
|
||||
c.target_name,
|
||||
c.model_id,
|
||||
fmt_ms_as_s(c.unload_ms_median),
|
||||
fmt_ms_as_s(c.load_ms_median),
|
||||
fmt_opt(c.cold_ttft_s_median, 3),
|
||||
fmt_opt(c.cold_total_s_median, 3),
|
||||
c.git_sha,
|
||||
c.samples,
|
||||
));
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
pub fn render_swap_json(costs: &[SwapCost]) -> Result<String> {
|
||||
Ok(serde_json::to_string_pretty(costs)?)
|
||||
}
|
||||
|
||||
/// Capability-probe view (#91): per (model, probe) the median quality score
|
||||
/// (the A/B number), then each run's id, score, and an artifact snippet so
|
||||
/// unscored runs can be located and scored (`helexa-bench score --id …`).
|
||||
pub fn render_capability_markdown(runs: &[CapabilityRun]) -> String {
|
||||
use std::collections::BTreeMap;
|
||||
let mut groups: BTreeMap<(String, String, String), Vec<&CapabilityRun>> = BTreeMap::new();
|
||||
for r in runs {
|
||||
groups
|
||||
.entry((
|
||||
r.target_name.clone(),
|
||||
r.model_id.clone(),
|
||||
r.scenario_id.clone(),
|
||||
))
|
||||
.or_default()
|
||||
.push(r);
|
||||
}
|
||||
let mut out = String::new();
|
||||
for ((target, model, scenario), rs) in groups {
|
||||
let scores: Vec<f64> = rs.iter().filter_map(|r| r.quality_score).collect();
|
||||
let median = median_slice(&scores);
|
||||
out.push_str(&format!(
|
||||
"### {target} · {model} · {scenario} — median score {} ({}/{} scored)\n\n",
|
||||
median
|
||||
.map(|m| format!("{m:.1}"))
|
||||
.unwrap_or_else(|| "—".into()),
|
||||
scores.len(),
|
||||
rs.len(),
|
||||
));
|
||||
out.push_str("| run | score | scorer | build | artifact (snippet) |\n");
|
||||
out.push_str("|---:|---:|---|---|---|\n");
|
||||
for r in rs {
|
||||
out.push_str(&format!(
|
||||
"| {} | {} | {} | `{}` | {} |\n",
|
||||
r.id,
|
||||
r.quality_score
|
||||
.map(|s| format!("{s:.1}"))
|
||||
.unwrap_or_else(|| "—".into()),
|
||||
r.scorer.as_deref().unwrap_or("—"),
|
||||
r.git_sha,
|
||||
snippet(r.artifact.as_deref()),
|
||||
));
|
||||
}
|
||||
out.push('\n');
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
pub fn render_capability_json(runs: &[CapabilityRun]) -> Result<String> {
|
||||
Ok(serde_json::to_string_pretty(runs)?)
|
||||
}
|
||||
|
||||
/// First ~80 chars of an artifact on one line, for the table cell.
|
||||
fn snippet(artifact: Option<&str>) -> String {
|
||||
match artifact {
|
||||
Some(a) => {
|
||||
let one_line: String = a.split_whitespace().collect::<Vec<_>>().join(" ");
|
||||
let trimmed: String = one_line.chars().take(80).collect();
|
||||
if one_line.chars().count() > 80 {
|
||||
format!("{trimmed}…")
|
||||
} else {
|
||||
trimmed
|
||||
}
|
||||
}
|
||||
None => "—".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn median_slice(v: &[f64]) -> Option<f64> {
|
||||
if v.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let mut s = v.to_vec();
|
||||
s.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||
let lo = (s.len() - 1) / 2;
|
||||
let hi = s.len() / 2;
|
||||
Some((s[lo] + s[hi]) / 2.0)
|
||||
}
|
||||
|
||||
/// Milliseconds rendered as seconds (reload costs read naturally in s).
|
||||
fn fmt_ms_as_s(ms: Option<f64>) -> String {
|
||||
match ms {
|
||||
Some(x) => format!("{:.2}", x / 1000.0),
|
||||
None => "—".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn fmt_opt(v: Option<f64>, places: usize) -> String {
|
||||
match v {
|
||||
Some(x) => format!("{x:.places$}"),
|
||||
@@ -61,9 +257,113 @@ fn fmt_opt(v: Option<f64>, places: usize) -> String {
|
||||
}
|
||||
}
|
||||
|
||||
/// Integer cell (concurrency width); `—` when unset (non-concurrency rows).
|
||||
fn fmt_u64(v: Option<u64>) -> String {
|
||||
match v {
|
||||
Some(x) => x.to_string(),
|
||||
None => "—".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// `used/total` in GB (e.g. `42.0/64.0`) — the headroom-at-a-glance cell.
|
||||
/// `used` alone if the node total is unknown; `—` if no telemetry.
|
||||
fn fmt_vram(used_mb: Option<f64>, total_mb: Option<u64>) -> String {
|
||||
match (used_mb, total_mb) {
|
||||
(Some(u), Some(t)) => format!("{:.1}/{:.1}", u / 1024.0, t as f64 / 1024.0),
|
||||
(Some(u), None) => format!("{:.1}", u / 1024.0),
|
||||
_ => "—".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::store::{ScalingCurve, ScalingPoint};
|
||||
|
||||
#[test]
|
||||
fn capability_markdown_groups_with_median_and_snippet() {
|
||||
let runs = vec![
|
||||
CapabilityRun {
|
||||
id: 7,
|
||||
ts: "t".into(),
|
||||
target_name: "beast".into(),
|
||||
model_id: "m".into(),
|
||||
scenario_id: "capability:plan".into(),
|
||||
git_sha: "abc".into(),
|
||||
quality_score: Some(8.0),
|
||||
scorer: Some("manual".into()),
|
||||
artifact: Some("A detailed plan with trade-offs and sequencing.".into()),
|
||||
},
|
||||
CapabilityRun {
|
||||
id: 8,
|
||||
ts: "t".into(),
|
||||
target_name: "beast".into(),
|
||||
model_id: "m".into(),
|
||||
scenario_id: "capability:plan".into(),
|
||||
git_sha: "abc".into(),
|
||||
quality_score: Some(6.0),
|
||||
scorer: Some("manual".into()),
|
||||
artifact: Some("Shorter plan.".into()),
|
||||
},
|
||||
];
|
||||
let md = render_capability_markdown(&runs);
|
||||
assert!(md.contains("capability:plan"));
|
||||
assert!(md.contains("median score 7.0")); // median(8,6)
|
||||
assert!(md.contains("trade-offs"));
|
||||
assert!(md.contains("| 7 |"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn swap_markdown_renders_reload_and_cold_costs() {
|
||||
let costs = vec![SwapCost {
|
||||
target_name: "beast".into(),
|
||||
model_id: "Qwen/Qwen3.6-27B".into(),
|
||||
git_sha: "abc1234".into(),
|
||||
gpu: Some("2× RTX 5090".into()),
|
||||
unload_ms_median: Some(320.0),
|
||||
load_ms_median: Some(25000.0),
|
||||
cold_ttft_s_median: Some(2.5),
|
||||
cold_total_s_median: Some(5.0),
|
||||
samples: 3,
|
||||
}];
|
||||
let md = render_swap_markdown(&costs);
|
||||
assert!(md.contains("reload (s)"));
|
||||
assert!(md.contains("beast"));
|
||||
assert!(md.contains("25.00")); // 25000 ms → 25.00 s
|
||||
assert!(md.contains("2.500"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scaling_markdown_renders_curve_and_flatness() {
|
||||
let curves = vec![ScalingCurve {
|
||||
target_name: "beast".into(),
|
||||
model_id: "Qwen/Qwen3.6-27B".into(),
|
||||
git_sha: "abc1234".into(),
|
||||
gpu: Some("2× RTX 5090".into()),
|
||||
points: vec![
|
||||
ScalingPoint {
|
||||
prompt_size: 128,
|
||||
prompt_tokens: Some(130),
|
||||
prefill_tps: Some(900.0),
|
||||
decode_tps: Some(50.0),
|
||||
samples: 5,
|
||||
},
|
||||
ScalingPoint {
|
||||
prompt_size: 4096,
|
||||
prompt_tokens: Some(4100),
|
||||
prefill_tps: Some(2800.0),
|
||||
decode_tps: Some(48.0),
|
||||
samples: 5,
|
||||
},
|
||||
],
|
||||
decode_flatness: Some(0.96),
|
||||
}];
|
||||
let md = render_scaling_markdown(&curves);
|
||||
assert!(md.contains("### beast · Qwen/Qwen3.6-27B"));
|
||||
assert!(md.contains("ctx tok"));
|
||||
assert!(md.contains("decode flatness: 0.96"));
|
||||
assert!(md.contains("holds"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn markdown_has_header_and_row() {
|
||||
@@ -77,14 +377,36 @@ mod tests {
|
||||
ttft_s_median: Some(0.123),
|
||||
decode_tps_median: Some(45.6),
|
||||
total_s_median: Some(1.234),
|
||||
ttft_s_p95: Some(0.222),
|
||||
ttft_s_p99: Some(0.250),
|
||||
total_s_p95: Some(1.5),
|
||||
total_s_p99: Some(1.6),
|
||||
prefill_ms_median: Some(120.0),
|
||||
decode_ms_median: Some(1100.0),
|
||||
prefill_tps_median: Some(1066.7),
|
||||
vram_used_mb_median: Some(43008.0),
|
||||
vram_total_mb: Some(65536),
|
||||
gpu_util_pct_median: Some(89.0),
|
||||
gpu_temp_c_median: Some(64.0),
|
||||
concurrency: None,
|
||||
ttft_p95_load_s: None,
|
||||
queue_wait_ms_median: None,
|
||||
rejected_median: None,
|
||||
samples: 5,
|
||||
gpu: Some("2× RTX 5090".into()),
|
||||
}];
|
||||
let md = render_markdown(&rows);
|
||||
assert!(md.contains("| engine |"));
|
||||
assert!(md.contains("prefill tok/s"));
|
||||
assert!(md.contains("VRAM (GB)"));
|
||||
assert!(md.contains("conc"));
|
||||
assert!(md.contains("beast"));
|
||||
assert!(md.contains("`30d50d6`"));
|
||||
assert!(md.contains("0.123"));
|
||||
// p95 column rendered.
|
||||
assert!(md.contains("0.222"));
|
||||
// VRAM used/total in GB (43008/65536 MiB → 42.0/64.0).
|
||||
assert!(md.contains("42.0/64.0"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -99,6 +421,21 @@ mod tests {
|
||||
ttft_s_median: Some(0.1),
|
||||
decode_tps_median: None,
|
||||
total_s_median: Some(0.5),
|
||||
ttft_s_p95: Some(0.1),
|
||||
ttft_s_p99: Some(0.1),
|
||||
total_s_p95: Some(0.5),
|
||||
total_s_p99: Some(0.5),
|
||||
prefill_ms_median: None,
|
||||
decode_ms_median: None,
|
||||
prefill_tps_median: None,
|
||||
vram_used_mb_median: None,
|
||||
vram_total_mb: None,
|
||||
gpu_util_pct_median: None,
|
||||
gpu_temp_c_median: None,
|
||||
concurrency: None,
|
||||
ttft_p95_load_s: None,
|
||||
queue_wait_ms_median: None,
|
||||
rejected_median: None,
|
||||
samples: 1,
|
||||
gpu: None,
|
||||
}];
|
||||
|
||||
@@ -62,6 +62,41 @@ pub struct ScenarioMetrics {
|
||||
pub prompt_tokens: Option<u64>,
|
||||
/// Completion tokens: from `usage` when present, else content-chunk count.
|
||||
pub completion_tokens: u64,
|
||||
/// Server-measured prefill duration (ms), from the `usage.helexa_timing`
|
||||
/// extension (#85). `None` when the server didn't emit it (external
|
||||
/// engines, non-instrumented paths). The honest prefill-phase number,
|
||||
/// distinct from client-observed `ttft_s` which also includes request
|
||||
/// setup + first-byte network latency.
|
||||
pub prefill_ms: Option<u64>,
|
||||
/// Server-measured decode duration (ms), from `usage.helexa_timing`.
|
||||
pub decode_ms: Option<u64>,
|
||||
/// Tokens submitted to prefill — the denominator for prefill tok/s.
|
||||
pub prefill_tokens: Option<u64>,
|
||||
// ── Concurrency / agentic-load fields (#89) ──────────────────────────
|
||||
// Set only by the concurrency scenario, which fans out N simultaneous
|
||||
// streams to characterize the real a0/hermes/opencode workload that
|
||||
// batch-1 single-request measurement can't see. `None` for single
|
||||
// requests. For a concurrency burst, the inherited fields carry the
|
||||
// aggregate: `ttft_s` = median TTFT across streams, `decode_tps` = node
|
||||
// throughput (total tokens / burst window), `total_s` = burst wall-clock,
|
||||
// `completion_tokens` = total across streams.
|
||||
/// Number of simultaneous streams in the burst (the cell dimension).
|
||||
pub concurrency: Option<u32>,
|
||||
/// p95 of per-stream TTFT within the burst — the tail under simultaneous
|
||||
/// load, where batch-1 serialization actually hurts.
|
||||
pub ttft_p95_s: Option<f64>,
|
||||
/// Median per-stream admission queue-wait (ms), approximated as
|
||||
/// `ttft − prefill_ms` (#85): on a batch-1 server, later streams wait for
|
||||
/// earlier ones, so TTFT inflates while server prefill stays constant —
|
||||
/// the gap is the wait. `None` if streams didn't report `helexa_timing`.
|
||||
pub queue_wait_ms_median: Option<f64>,
|
||||
/// Streams shed by admission control (HTTP 429/503) during the burst —
|
||||
/// honest backpressure, not silent failures.
|
||||
pub rejected: Option<u32>,
|
||||
/// Full generated text, captured only by the capability probe (#91) so
|
||||
/// the output can be quality-scored later (manual or LLM-judge). `None`
|
||||
/// for latency/throughput scenarios, which discard the text.
|
||||
pub artifact: Option<String>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -85,10 +120,13 @@ pub trait Scenario: Send + Sync {
|
||||
async fn run(&self, ctx: &RunCtx) -> Result<ScenarioMetrics>;
|
||||
}
|
||||
|
||||
/// Build the active scenario set from config. One chat-latency scenario
|
||||
/// per configured prompt size.
|
||||
/// Build the active scenario set from config: one chat-latency scenario per
|
||||
/// prompt size, plus one concurrency scenario per configured level (#89).
|
||||
/// Concurrency levels default to empty (opt-in), since a burst puts real
|
||||
/// simultaneous load on a serving fleet — operators enable it deliberately.
|
||||
pub fn build_scenarios(cfg: &ScenarioConfig) -> Vec<Box<dyn Scenario>> {
|
||||
cfg.prompt_sizes
|
||||
let mut scenarios: Vec<Box<dyn Scenario>> = cfg
|
||||
.prompt_sizes
|
||||
.iter()
|
||||
.map(|&size| {
|
||||
Box::new(ChatLatencyScenario {
|
||||
@@ -96,7 +134,46 @@ pub fn build_scenarios(cfg: &ScenarioConfig) -> Vec<Box<dyn Scenario>> {
|
||||
approx_prompt_tokens: size,
|
||||
}) as Box<dyn Scenario>
|
||||
})
|
||||
.collect()
|
||||
.collect();
|
||||
for &n in &cfg.concurrency_levels {
|
||||
scenarios.push(Box::new(ConcurrencyScenario {
|
||||
id: format!("concurrency:{n}"),
|
||||
concurrency: n,
|
||||
approx_prompt_tokens: cfg.concurrency_prompt_tokens,
|
||||
}) as Box<dyn Scenario>);
|
||||
}
|
||||
for probe in &cfg.capability_probes {
|
||||
scenarios.push(Box::new(CapabilityScenario {
|
||||
id: format!("capability:{}", probe.name),
|
||||
prompt: probe.prompt.clone(),
|
||||
max_tokens: probe.max_tokens,
|
||||
}) as Box<dyn Scenario>);
|
||||
}
|
||||
scenarios
|
||||
}
|
||||
|
||||
/// A single small streamed request, timed like a chat-latency run. Used by
|
||||
/// the swap-cost measurement (#90) to capture the cold first-request latency
|
||||
/// straight after a reload. Reuses the shared SSE-timing core.
|
||||
pub async fn cold_probe(ctx: &RunCtx<'_>) -> Result<ScenarioMetrics> {
|
||||
let prompt = build_prompt(128);
|
||||
let payload = chat_payload(ctx, &prompt);
|
||||
tokio::time::timeout(ctx.timeout, stream_and_measure(ctx, &payload))
|
||||
.await
|
||||
.map_err(|_| anyhow!("cold probe timed out after {:?}", ctx.timeout))?
|
||||
}
|
||||
|
||||
/// The chat-completions request body shared by the latency and concurrency
|
||||
/// scenarios — streamed, deterministic (temperature 0), usage included.
|
||||
fn chat_payload(ctx: &RunCtx, prompt: &str) -> serde_json::Value {
|
||||
json!({
|
||||
"model": ctx.model_id,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"max_tokens": ctx.max_tokens,
|
||||
"temperature": 0,
|
||||
"stream": true,
|
||||
"stream_options": {"include_usage": true},
|
||||
})
|
||||
}
|
||||
|
||||
/// Streamed single-request chat-completions latency probe — the batch-1
|
||||
@@ -118,15 +195,7 @@ impl Scenario for ChatLatencyScenario {
|
||||
|
||||
async fn run(&self, ctx: &RunCtx) -> Result<ScenarioMetrics> {
|
||||
let prompt = build_prompt(self.approx_prompt_tokens);
|
||||
let payload = json!({
|
||||
"model": ctx.model_id,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"max_tokens": ctx.max_tokens,
|
||||
"temperature": 0,
|
||||
"stream": true,
|
||||
"stream_options": {"include_usage": true},
|
||||
});
|
||||
|
||||
let payload = chat_payload(ctx, &prompt);
|
||||
let fut = stream_and_measure(ctx, &payload);
|
||||
tokio::time::timeout(ctx.timeout, fut)
|
||||
.await
|
||||
@@ -134,11 +203,182 @@ impl Scenario for ChatLatencyScenario {
|
||||
}
|
||||
}
|
||||
|
||||
/// Fan-out load probe: fire `concurrency` identical streams at once and
|
||||
/// measure how the fleet behaves under simultaneous pressure (#89). This is
|
||||
/// the only scenario that exercises the real a0/hermes/opencode pattern —
|
||||
/// many agentic requests per user turn — which batch-1 single-request
|
||||
/// timing cannot characterize. On a batch-1 serialized server, aggregate
|
||||
/// throughput stays ~flat while TTFT/queue-wait inflate with `concurrency`;
|
||||
/// that gap is the evidence for/against continuous batching.
|
||||
pub struct ConcurrencyScenario {
|
||||
id: String,
|
||||
concurrency: u32,
|
||||
approx_prompt_tokens: u32,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Scenario for ConcurrencyScenario {
|
||||
fn id(&self) -> &str {
|
||||
&self.id
|
||||
}
|
||||
|
||||
fn prompt_size(&self) -> u32 {
|
||||
self.approx_prompt_tokens
|
||||
}
|
||||
|
||||
async fn run(&self, ctx: &RunCtx) -> Result<ScenarioMetrics> {
|
||||
let prompt = build_prompt(self.approx_prompt_tokens);
|
||||
let payload = chat_payload(ctx, &prompt);
|
||||
|
||||
// Fire all streams at once; each is independently timed and capped by
|
||||
// the per-request timeout so one hung stream can't stall the burst.
|
||||
let burst_start = Instant::now();
|
||||
let futs = (0..self.concurrency).map(|_| async {
|
||||
tokio::time::timeout(ctx.timeout, stream_and_measure(ctx, &payload)).await
|
||||
});
|
||||
let results = futures::future::join_all(futs).await;
|
||||
let burst_window = burst_start.elapsed().as_secs_f64();
|
||||
|
||||
let mut streams: Vec<ScenarioMetrics> = Vec::new();
|
||||
let mut rejected: u32 = 0;
|
||||
for r in results {
|
||||
match r {
|
||||
Ok(Ok(m)) => streams.push(m),
|
||||
// Admission backpressure (429/503) is shed load, counted
|
||||
// separately from genuine failures/timeouts.
|
||||
Ok(Err(e)) if is_admission_reject(&e) => rejected += 1,
|
||||
Ok(Err(_)) | Err(_) => {}
|
||||
}
|
||||
}
|
||||
if streams.is_empty() {
|
||||
return Err(anyhow!(
|
||||
"all {} concurrent streams failed ({rejected} shed by admission)",
|
||||
self.concurrency
|
||||
));
|
||||
}
|
||||
|
||||
let total_tokens: u64 = streams.iter().map(|m| m.completion_tokens).sum();
|
||||
let ttfts: Vec<f64> = streams.iter().map(|m| m.ttft_s).collect();
|
||||
// queue-wait ≈ TTFT − server prefill (#85); only for streams that
|
||||
// reported helexa_timing.
|
||||
let queue_waits: Vec<f64> = streams
|
||||
.iter()
|
||||
.filter_map(|m| {
|
||||
m.prefill_ms
|
||||
.map(|p| (m.ttft_s * 1000.0 - p as f64).max(0.0))
|
||||
})
|
||||
.collect();
|
||||
// Aggregate decode throughput across the whole node for the burst.
|
||||
let aggregate_tps = if burst_window > 0.0 {
|
||||
Some(total_tokens as f64 / burst_window)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
Ok(ScenarioMetrics {
|
||||
ttft_s: median(&ttfts).unwrap_or(0.0),
|
||||
decode_tps: aggregate_tps,
|
||||
total_s: burst_window,
|
||||
prompt_tokens: streams.iter().find_map(|m| m.prompt_tokens),
|
||||
completion_tokens: total_tokens,
|
||||
prefill_ms: None,
|
||||
decode_ms: None,
|
||||
prefill_tokens: None,
|
||||
concurrency: Some(self.concurrency),
|
||||
ttft_p95_s: percentile(&ttfts, 95.0),
|
||||
queue_wait_ms_median: median(&queue_waits),
|
||||
rejected: Some(rejected),
|
||||
artifact: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Quality probe (#91): runs a fixed prompt and stores the full generated
|
||||
/// text as an artifact for later scoring (manual now, LLM-judge later). The
|
||||
/// point is to compare reasoning/planning quality across models — the axis
|
||||
/// speed-only scenarios miss — so the frontier A/B (F3) picks on capability,
|
||||
/// not just throughput.
|
||||
pub struct CapabilityScenario {
|
||||
id: String,
|
||||
prompt: String,
|
||||
max_tokens: u64,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Scenario for CapabilityScenario {
|
||||
fn id(&self) -> &str {
|
||||
&self.id
|
||||
}
|
||||
|
||||
/// Capability probes have no synthetic prompt-token target; the cell is
|
||||
/// keyed by the scenario id alone.
|
||||
fn prompt_size(&self) -> u32 {
|
||||
0
|
||||
}
|
||||
|
||||
async fn run(&self, ctx: &RunCtx) -> Result<ScenarioMetrics> {
|
||||
let payload = json!({
|
||||
"model": ctx.model_id,
|
||||
"messages": [{"role": "user", "content": self.prompt}],
|
||||
"max_tokens": self.max_tokens,
|
||||
"temperature": 0,
|
||||
"stream": true,
|
||||
"stream_options": {"include_usage": true},
|
||||
});
|
||||
let fut = stream_and_measure_inner(ctx, &payload, true);
|
||||
tokio::time::timeout(ctx.timeout, fut)
|
||||
.await
|
||||
.map_err(|_| anyhow!("capability probe timed out after {:?}", ctx.timeout))?
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether a stream error was admission backpressure (HTTP 429/503) rather
|
||||
/// than a genuine failure. `stream_and_measure` renders the upstream status
|
||||
/// into the error string, so a substring check is sufficient.
|
||||
fn is_admission_reject(e: &anyhow::Error) -> bool {
|
||||
let s = e.to_string();
|
||||
s.contains("429") || s.contains("503")
|
||||
}
|
||||
|
||||
/// Median of a slice (sorted copy). `None` if empty.
|
||||
fn median(values: &[f64]) -> Option<f64> {
|
||||
if values.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let mut v = values.to_vec();
|
||||
v.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||
let lo = (v.len() - 1) / 2;
|
||||
let hi = v.len() / 2;
|
||||
Some((v[lo] + v[hi]) / 2.0)
|
||||
}
|
||||
|
||||
/// Nearest-rank percentile of a slice (`p` in 0..=100). `None` if empty.
|
||||
fn percentile(values: &[f64], p: f64) -> Option<f64> {
|
||||
if values.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let mut v = values.to_vec();
|
||||
v.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||
let rank = (p / 100.0 * v.len() as f64).ceil() as usize;
|
||||
Some(v[rank.clamp(1, v.len()) - 1])
|
||||
}
|
||||
|
||||
/// The SSE-timing core, ported from `bench.py::one_run`. Kept free of the
|
||||
/// `Scenario` trait so it's unit-testable against a mock byte stream.
|
||||
async fn stream_and_measure(
|
||||
ctx: &RunCtx<'_>,
|
||||
payload: &serde_json::Value,
|
||||
) -> Result<ScenarioMetrics> {
|
||||
stream_and_measure_inner(ctx, payload, false).await
|
||||
}
|
||||
|
||||
/// As [`stream_and_measure`] but accumulates the full visible text when
|
||||
/// `capture_text` is set — used by the capability probe (#91) to store the
|
||||
/// generated artifact for later quality scoring.
|
||||
async fn stream_and_measure_inner(
|
||||
ctx: &RunCtx<'_>,
|
||||
payload: &serde_json::Value,
|
||||
capture_text: bool,
|
||||
) -> Result<ScenarioMetrics> {
|
||||
let start = Instant::now();
|
||||
let resp = ctx
|
||||
@@ -160,6 +400,10 @@ async fn stream_and_measure(
|
||||
let mut chunk_count: u64 = 0;
|
||||
let mut prompt_tokens: Option<u64> = None;
|
||||
let mut completion_tokens: Option<u64> = None;
|
||||
let mut prefill_ms: Option<u64> = None;
|
||||
let mut decode_ms: Option<u64> = None;
|
||||
let mut prefill_tokens: Option<u64> = None;
|
||||
let mut captured = String::new();
|
||||
|
||||
while let Some(event) = stream.next().await {
|
||||
let event = event.context("reading SSE stream")?;
|
||||
@@ -172,46 +416,86 @@ async fn stream_and_measure(
|
||||
Ok(c) => c,
|
||||
Err(_) => continue, // tolerate non-JSON keepalive frames
|
||||
};
|
||||
if let Some(choice) = chunk.choices.first()
|
||||
&& choice
|
||||
if let Some(choice) = chunk.choices.first() {
|
||||
// Liveness counts ANY generated delta (#117). Thinking
|
||||
// models (Qwen3-Next-Thinking, Qwen3 with thinking on)
|
||||
// stream `reasoning_content` first — sometimes for their
|
||||
// entire budget — and a content-only view misread that as
|
||||
// a dead stream ("no content chunks received") while also
|
||||
// producing impossible client-side rates (reasoning-
|
||||
// inclusive token counts over a visible-content-only
|
||||
// window; observed: "244 tok/s" on a 3060). For
|
||||
// non-thinking models the first delta IS content, so
|
||||
// `ttft_s` semantics are unchanged for them.
|
||||
let content = choice
|
||||
.delta
|
||||
.get("content")
|
||||
.and_then(|c| c.as_str())
|
||||
.is_some_and(|s| !s.is_empty())
|
||||
{
|
||||
if first.is_none() {
|
||||
first = Some(now);
|
||||
.filter(|c| !c.is_empty());
|
||||
let reasoning = choice
|
||||
.delta
|
||||
.get("reasoning_content")
|
||||
.and_then(|c| c.as_str())
|
||||
.filter(|c| !c.is_empty());
|
||||
if content.is_some() || reasoning.is_some() {
|
||||
if first.is_none() {
|
||||
first = Some(now);
|
||||
}
|
||||
last = Some(now);
|
||||
chunk_count += 1;
|
||||
}
|
||||
if capture_text && let Some(text) = content {
|
||||
captured.push_str(text);
|
||||
}
|
||||
last = Some(now);
|
||||
chunk_count += 1;
|
||||
}
|
||||
if let Some(usage) = chunk.usage {
|
||||
prompt_tokens = Some(usage.prompt_tokens);
|
||||
completion_tokens = Some(usage.completion_tokens);
|
||||
if let Some(t) = usage.helexa_timing {
|
||||
prefill_ms = Some(t.prefill_ms);
|
||||
decode_ms = Some(t.decode_ms);
|
||||
prefill_tokens = Some(t.prefill_tokens);
|
||||
}
|
||||
}
|
||||
}
|
||||
let end = Instant::now();
|
||||
|
||||
let first = first.ok_or_else(|| anyhow!("no content chunks received"))?;
|
||||
let first = first.ok_or_else(|| anyhow!("no generated chunks received"))?;
|
||||
|
||||
// neuron emits one SSE chunk per visible token, so chunk_count is an
|
||||
// engine-truth count when no usage frame is sent.
|
||||
// neuron emits one SSE chunk per generated token, so chunk_count is
|
||||
// an engine-truth count when no usage frame is sent.
|
||||
let tokens = completion_tokens.filter(|&t| t > 0).unwrap_or(chunk_count);
|
||||
// decode rate is only meaningful over a real inter-chunk window.
|
||||
// Decode rate: prefer the server-measured split (#85) — it counts
|
||||
// every generated token over the actual decode window, immune to
|
||||
// reasoning-suppression frame mismatches. Fall back to the client
|
||||
// inter-chunk window with the CHUNK count (same frame) — never
|
||||
// usage.completion_tokens over the chunk window, which mixes a
|
||||
// reasoning-inclusive numerator with a visible-only denominator.
|
||||
let window = last
|
||||
.filter(|&l| l > first)
|
||||
.map(|l| (l - first).as_secs_f64())
|
||||
.unwrap_or(0.0);
|
||||
let decode_tps = match decode_ms {
|
||||
Some(ms) if ms > 200 && tokens > 0 => Some(tokens as f64 / (ms as f64 / 1000.0)),
|
||||
_ if window > 0.2 => Some(chunk_count as f64 / window),
|
||||
_ => None,
|
||||
};
|
||||
Ok(ScenarioMetrics {
|
||||
ttft_s: (first - start).as_secs_f64(),
|
||||
decode_tps: if window > 0.2 {
|
||||
Some(tokens as f64 / window)
|
||||
} else {
|
||||
None
|
||||
},
|
||||
decode_tps,
|
||||
total_s: (end - start).as_secs_f64(),
|
||||
prompt_tokens,
|
||||
completion_tokens: tokens,
|
||||
prefill_ms,
|
||||
decode_ms,
|
||||
prefill_tokens,
|
||||
// Concurrency fields unset on the single-request path; the
|
||||
// concurrency scenario builds its own aggregate (#89).
|
||||
concurrency: None,
|
||||
ttft_p95_s: None,
|
||||
queue_wait_ms_median: None,
|
||||
rejected: None,
|
||||
artifact: if capture_text { Some(captured) } else { None },
|
||||
})
|
||||
}
|
||||
|
||||
@@ -229,6 +513,54 @@ mod tests {
|
||||
assert!(small.ends_with("/no_think"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn median_and_percentile_basics() {
|
||||
assert_eq!(median(&[3.0, 1.0, 2.0]), Some(2.0));
|
||||
assert_eq!(median(&[]), None);
|
||||
let v = [1.0, 2.0, 3.0, 4.0, 5.0];
|
||||
assert_eq!(percentile(&v, 50.0), Some(3.0));
|
||||
assert_eq!(percentile(&v, 95.0), Some(5.0)); // nearest-rank → max with n=5
|
||||
assert_eq!(percentile(&[], 95.0), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn admission_rejects_detected_by_status() {
|
||||
assert!(is_admission_reject(&anyhow!(
|
||||
"upstream returned 429 Too Many Requests"
|
||||
)));
|
||||
assert!(is_admission_reject(&anyhow!(
|
||||
"upstream returned 503 Service Unavailable"
|
||||
)));
|
||||
assert!(!is_admission_reject(&anyhow!(
|
||||
"upstream returned 500 Internal"
|
||||
)));
|
||||
assert!(!is_admission_reject(&anyhow!("connection refused")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn concurrency_scenarios_built_from_config() {
|
||||
use crate::config::{CapabilityProbe, ScenarioConfig};
|
||||
let cfg = ScenarioConfig {
|
||||
prompt_sizes: vec![128],
|
||||
max_tokens: 64,
|
||||
concurrency_levels: vec![2, 8],
|
||||
concurrency_prompt_tokens: 512,
|
||||
capability_probes: vec![CapabilityProbe {
|
||||
name: "plan".into(),
|
||||
prompt: "Write a plan.".into(),
|
||||
max_tokens: 2048,
|
||||
}],
|
||||
};
|
||||
let ids: Vec<String> = build_scenarios(&cfg)
|
||||
.iter()
|
||||
.map(|s| s.id().to_string())
|
||||
.collect();
|
||||
assert!(ids.contains(&"chat:128".to_string()));
|
||||
assert!(ids.contains(&"concurrency:2".to_string()));
|
||||
assert!(ids.contains(&"concurrency:8".to_string()));
|
||||
assert!(ids.contains(&"capability:plan".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prompt_floor_for_tiny_targets() {
|
||||
// max(approx,16) floor means even 0 yields a non-trivial prompt.
|
||||
|
||||
@@ -51,6 +51,35 @@ pub struct RunRecord {
|
||||
pub decode_tps: Option<f64>,
|
||||
pub total_s: Option<f64>,
|
||||
pub completion_tokens: Option<u64>,
|
||||
// server-measured prefill/decode split (#85), null on engines/paths
|
||||
// that don't emit `usage.helexa_timing`.
|
||||
pub prefill_ms: Option<u64>,
|
||||
pub decode_ms: Option<u64>,
|
||||
pub prefill_tokens: Option<u64>,
|
||||
// GPU telemetry sampled from /health around the run (#87), null for
|
||||
// non-neuron targets or when /health was unreachable. vram_used_mb is
|
||||
// the node sum; util/temp are the hottest single device.
|
||||
pub vram_used_mb: Option<u64>,
|
||||
pub gpu_util_pct: Option<u32>,
|
||||
pub gpu_temp_c: Option<u32>,
|
||||
// concurrency / agentic-load burst metrics (#89), null for single-request
|
||||
// scenarios. For a burst, ttft_s/decode_tps/total_s carry the aggregate.
|
||||
pub concurrency: Option<u32>,
|
||||
pub ttft_p95_s: Option<f64>,
|
||||
pub queue_wait_ms: Option<f64>,
|
||||
pub rejected: Option<u32>,
|
||||
// cold-load / model-swap cost (#90), set only by the deliberate
|
||||
// `swap-cost` measurement (scenario_id = "swap"). The other metric
|
||||
// fields carry the cold first-request after reload.
|
||||
pub swap_unload_ms: Option<u64>,
|
||||
pub swap_load_ms: Option<u64>,
|
||||
// capability probe (#91): the full generated text (scenario_id =
|
||||
// "capability:<name>"), scored for quality later. `quality_score` /
|
||||
// `scorer` are null at insert — set by `score` (manual) or a future
|
||||
// LLM-judge.
|
||||
pub artifact: Option<String>,
|
||||
pub quality_score: Option<f64>,
|
||||
pub scorer: Option<String>,
|
||||
// outcome
|
||||
pub ok: bool,
|
||||
pub error: Option<String>,
|
||||
@@ -123,6 +152,21 @@ impl Store {
|
||||
decode_tps REAL,
|
||||
total_s REAL,
|
||||
completion_tokens INTEGER,
|
||||
prefill_ms INTEGER,
|
||||
decode_ms INTEGER,
|
||||
prefill_tokens INTEGER,
|
||||
vram_used_mb INTEGER,
|
||||
gpu_util_pct INTEGER,
|
||||
gpu_temp_c INTEGER,
|
||||
concurrency INTEGER,
|
||||
ttft_p95_s REAL,
|
||||
queue_wait_ms REAL,
|
||||
rejected INTEGER,
|
||||
swap_unload_ms INTEGER,
|
||||
swap_load_ms INTEGER,
|
||||
artifact TEXT,
|
||||
quality_score REAL,
|
||||
scorer TEXT,
|
||||
ok INTEGER NOT NULL,
|
||||
error TEXT
|
||||
);
|
||||
@@ -133,6 +177,51 @@ impl Store {
|
||||
"#,
|
||||
)
|
||||
.context("initialising sqlite schema")?;
|
||||
// Additive migrations for DBs created before a column existed.
|
||||
// `CREATE TABLE IF NOT EXISTS` above only seeds fresh DBs; existing
|
||||
// ones need the columns backfilled (as NULL) so older rows coexist
|
||||
// with new metrics. There is no migration framework — each entry is
|
||||
// an idempotent "add if missing".
|
||||
Self::ensure_columns(
|
||||
conn,
|
||||
"runs",
|
||||
&[
|
||||
("prefill_ms", "INTEGER"),
|
||||
("decode_ms", "INTEGER"),
|
||||
("prefill_tokens", "INTEGER"),
|
||||
("vram_used_mb", "INTEGER"),
|
||||
("gpu_util_pct", "INTEGER"),
|
||||
("gpu_temp_c", "INTEGER"),
|
||||
("concurrency", "INTEGER"),
|
||||
("ttft_p95_s", "REAL"),
|
||||
("queue_wait_ms", "REAL"),
|
||||
("rejected", "INTEGER"),
|
||||
("swap_unload_ms", "INTEGER"),
|
||||
("swap_load_ms", "INTEGER"),
|
||||
("artifact", "TEXT"),
|
||||
("quality_score", "REAL"),
|
||||
("scorer", "TEXT"),
|
||||
],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Add any of `columns` that the table is missing (`ALTER TABLE ADD
|
||||
/// COLUMN`). Idempotent: existing columns are read from
|
||||
/// `PRAGMA table_info` and skipped, so this is safe to run on every open.
|
||||
fn ensure_columns(conn: &Connection, table: &str, columns: &[(&str, &str)]) -> Result<()> {
|
||||
let mut existing = std::collections::HashSet::new();
|
||||
let mut stmt = conn.prepare(&format!("PRAGMA table_info({table})"))?;
|
||||
let names = stmt.query_map([], |row| row.get::<_, String>(1))?;
|
||||
for name in names {
|
||||
existing.insert(name?);
|
||||
}
|
||||
for (name, ty) in columns {
|
||||
if !existing.contains(*name) {
|
||||
conn.execute_batch(&format!("ALTER TABLE {table} ADD COLUMN {name} {ty};"))
|
||||
.with_context(|| format!("adding column {table}.{name}"))?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -166,6 +255,11 @@ impl Store {
|
||||
model_id, harness, capabilities_json, devices_json,
|
||||
scenario_id, prompt_size_approx, prompt_tokens_actual, max_tokens,
|
||||
ttft_s, decode_tps, total_s, completion_tokens,
|
||||
prefill_ms, decode_ms, prefill_tokens,
|
||||
vram_used_mb, gpu_util_pct, gpu_temp_c,
|
||||
concurrency, ttft_p95_s, queue_wait_ms, rejected,
|
||||
swap_unload_ms, swap_load_ms,
|
||||
artifact, quality_score, scorer,
|
||||
ok, error
|
||||
) VALUES (
|
||||
?1, ?2, ?3, ?4,
|
||||
@@ -176,7 +270,12 @@ impl Store {
|
||||
?20, ?21, ?22, ?23,
|
||||
?24, ?25, ?26, ?27,
|
||||
?28, ?29, ?30, ?31,
|
||||
?32, ?33
|
||||
?32, ?33, ?34,
|
||||
?35, ?36, ?37,
|
||||
?38, ?39, ?40, ?41,
|
||||
?42, ?43,
|
||||
?44, ?45, ?46,
|
||||
?47, ?48
|
||||
)",
|
||||
params![
|
||||
r.ts,
|
||||
@@ -210,6 +309,21 @@ impl Store {
|
||||
r.decode_tps,
|
||||
r.total_s,
|
||||
r.completion_tokens,
|
||||
r.prefill_ms,
|
||||
r.decode_ms,
|
||||
r.prefill_tokens,
|
||||
r.vram_used_mb,
|
||||
r.gpu_util_pct,
|
||||
r.gpu_temp_c,
|
||||
r.concurrency,
|
||||
r.ttft_p95_s,
|
||||
r.queue_wait_ms,
|
||||
r.rejected,
|
||||
r.swap_unload_ms,
|
||||
r.swap_load_ms,
|
||||
r.artifact,
|
||||
r.quality_score,
|
||||
r.scorer,
|
||||
r.ok as i64,
|
||||
r.error,
|
||||
],
|
||||
@@ -224,7 +338,10 @@ impl Store {
|
||||
// successful run, then median that SHA's samples.
|
||||
let mut stmt = self.conn.prepare(
|
||||
"SELECT target_name, model_id, scenario_id, prompt_size_approx, git_sha,
|
||||
ttft_s, decode_tps, total_s, prompt_tokens_actual, gpus_json
|
||||
ttft_s, decode_tps, total_s, prompt_tokens_actual, gpus_json,
|
||||
prefill_ms, decode_ms, prefill_tokens,
|
||||
vram_used_mb, gpu_util_pct, gpu_temp_c,
|
||||
concurrency, ttft_p95_s, queue_wait_ms, rejected
|
||||
FROM runs
|
||||
WHERE ok=1
|
||||
ORDER BY target_name, model_id, scenario_id, id",
|
||||
@@ -241,12 +358,188 @@ impl Store {
|
||||
total_s: row.get(7)?,
|
||||
prompt_tokens_actual: row.get(8)?,
|
||||
gpus_json: row.get(9)?,
|
||||
prefill_ms: row.get(10)?,
|
||||
decode_ms: row.get(11)?,
|
||||
prefill_tokens: row.get(12)?,
|
||||
vram_used_mb: row.get(13)?,
|
||||
gpu_util_pct: row.get(14)?,
|
||||
gpu_temp_c: row.get(15)?,
|
||||
concurrency: row.get(16)?,
|
||||
ttft_p95_s: row.get(17)?,
|
||||
queue_wait_ms: row.get(18)?,
|
||||
rejected: row.get(19)?,
|
||||
})
|
||||
})?;
|
||||
let raws: Vec<RawRow> = rows.collect::<rusqlite::Result<_>>()?;
|
||||
Ok(aggregate(raws))
|
||||
}
|
||||
|
||||
/// Context-length scaling curves (#88): per (target, model), the latest
|
||||
/// build's `chat:<n>` cells pivoted by prompt size into prefill & decode
|
||||
/// tok/s vs context. The headline is `decode_flatness` — decode tok/s at
|
||||
/// the largest context divided by the smallest. Near 1.0 confirms the
|
||||
/// Gated-DeltaNet O(1)-in-sequence-length decode; a sharp drop locates
|
||||
/// where the model stops scaling for free.
|
||||
pub fn scaling(&self) -> Result<Vec<ScalingCurve>> {
|
||||
use std::collections::BTreeMap;
|
||||
// Reuse the already-aggregated report cells; the chat:<n> rows are
|
||||
// exactly the per-context measurement points.
|
||||
let mut by_model: BTreeMap<(String, String), Vec<ReportRow>> = BTreeMap::new();
|
||||
for r in self.report_rows()? {
|
||||
if r.scenario_id.starts_with("chat:") {
|
||||
by_model
|
||||
.entry((r.target_name.clone(), r.model_id.clone()))
|
||||
.or_default()
|
||||
.push(r);
|
||||
}
|
||||
}
|
||||
let mut out = Vec::new();
|
||||
for ((target_name, model_id), mut rows) in by_model {
|
||||
rows.sort_by_key(|r| r.prompt_size_approx);
|
||||
let points: Vec<ScalingPoint> = rows
|
||||
.iter()
|
||||
.map(|r| ScalingPoint {
|
||||
prompt_size: r.prompt_size_approx,
|
||||
prompt_tokens: r.prompt_tokens,
|
||||
prefill_tps: r.prefill_tps_median,
|
||||
decode_tps: r.decode_tps_median,
|
||||
samples: r.samples,
|
||||
})
|
||||
.collect();
|
||||
// Flatness across the smallest→largest points that both have a
|
||||
// decode rate (skips cells where the decode window was too short).
|
||||
let with_decode: Vec<&ScalingPoint> =
|
||||
points.iter().filter(|p| p.decode_tps.is_some()).collect();
|
||||
let decode_flatness = match (with_decode.first(), with_decode.last()) {
|
||||
(Some(lo), Some(hi)) if with_decode.len() >= 2 => {
|
||||
match (lo.decode_tps, hi.decode_tps) {
|
||||
(Some(a), Some(b)) if a > 0.0 => Some(b / a),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
out.push(ScalingCurve {
|
||||
target_name,
|
||||
model_id,
|
||||
git_sha: rows.first().map(|r| r.git_sha.clone()).unwrap_or_default(),
|
||||
gpu: rows.iter().find_map(|r| r.gpu.clone()),
|
||||
points,
|
||||
decode_flatness,
|
||||
});
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Cold-load / model-swap costs (#90): per (target, model) at the latest
|
||||
/// build, the median unload→reload latency and the cold first-request
|
||||
/// latency after reload (the `scenario_id = "swap"` rows).
|
||||
pub fn swap_costs(&self) -> Result<Vec<SwapCost>> {
|
||||
use std::collections::BTreeMap;
|
||||
let mut stmt = self.conn.prepare(
|
||||
"SELECT target_name, model_id, git_sha, gpus_json,
|
||||
swap_unload_ms, swap_load_ms, ttft_s, total_s
|
||||
FROM runs
|
||||
WHERE ok=1 AND scenario_id='swap'
|
||||
ORDER BY target_name, model_id, id",
|
||||
)?;
|
||||
struct Raw {
|
||||
target: String,
|
||||
model: String,
|
||||
sha: String,
|
||||
gpus_json: Option<String>,
|
||||
unload_ms: Option<f64>,
|
||||
load_ms: Option<f64>,
|
||||
ttft_s: Option<f64>,
|
||||
total_s: Option<f64>,
|
||||
}
|
||||
let raws: Vec<Raw> = stmt
|
||||
.query_map([], |r| {
|
||||
Ok(Raw {
|
||||
target: r.get(0)?,
|
||||
model: r.get(1)?,
|
||||
sha: r.get(2)?,
|
||||
gpus_json: r.get(3)?,
|
||||
unload_ms: r.get::<_, Option<i64>>(4)?.map(|v| v as f64),
|
||||
load_ms: r.get::<_, Option<i64>>(5)?.map(|v| v as f64),
|
||||
ttft_s: r.get(6)?,
|
||||
total_s: r.get(7)?,
|
||||
})
|
||||
})?
|
||||
.collect::<rusqlite::Result<_>>()?;
|
||||
let mut by: BTreeMap<(String, String), Vec<Raw>> = BTreeMap::new();
|
||||
for r in raws {
|
||||
by.entry((r.target.clone(), r.model.clone()))
|
||||
.or_default()
|
||||
.push(r);
|
||||
}
|
||||
let mut out = Vec::new();
|
||||
for ((target, model), rows) in by {
|
||||
let latest = rows.last().map(|r| r.sha.clone()).unwrap_or_default();
|
||||
let cell: Vec<&Raw> = rows.iter().filter(|r| r.sha == latest).collect();
|
||||
out.push(SwapCost {
|
||||
target_name: target,
|
||||
model_id: model,
|
||||
git_sha: latest,
|
||||
gpu: cell
|
||||
.iter()
|
||||
.find_map(|r| r.gpus_json.as_deref().and_then(gpu_label)),
|
||||
unload_ms_median: median(cell.iter().filter_map(|r| r.unload_ms)),
|
||||
load_ms_median: median(cell.iter().filter_map(|r| r.load_ms)),
|
||||
cold_ttft_s_median: median(cell.iter().filter_map(|r| r.ttft_s)),
|
||||
cold_total_s_median: median(cell.iter().filter_map(|r| r.total_s)),
|
||||
samples: cell.len(),
|
||||
});
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Capability-probe runs (#91): the stored output artifacts and their
|
||||
/// quality scores. `unscored_only` filters to runs awaiting a score —
|
||||
/// the worklist for manual scoring or a future LLM-judge.
|
||||
pub fn capability_runs(&self, unscored_only: bool) -> Result<Vec<CapabilityRun>> {
|
||||
let sql = format!(
|
||||
"SELECT id, ts, target_name, model_id, scenario_id, git_sha,
|
||||
quality_score, scorer, artifact
|
||||
FROM runs
|
||||
WHERE ok=1 AND scenario_id LIKE 'capability:%'{}
|
||||
ORDER BY id DESC",
|
||||
if unscored_only {
|
||||
" AND quality_score IS NULL"
|
||||
} else {
|
||||
""
|
||||
},
|
||||
);
|
||||
let mut stmt = self.conn.prepare(&sql)?;
|
||||
let rows = stmt
|
||||
.query_map([], |r| {
|
||||
Ok(CapabilityRun {
|
||||
id: r.get(0)?,
|
||||
ts: r.get(1)?,
|
||||
target_name: r.get(2)?,
|
||||
model_id: r.get(3)?,
|
||||
scenario_id: r.get(4)?,
|
||||
git_sha: r.get(5)?,
|
||||
quality_score: r.get(6)?,
|
||||
scorer: r.get(7)?,
|
||||
artifact: r.get(8)?,
|
||||
})
|
||||
})?
|
||||
.collect::<rusqlite::Result<_>>()?;
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
/// Attach a quality score to one capability run (#91). `scorer` records
|
||||
/// who/what scored it ("manual", "llm:claude-…"). Returns the row count
|
||||
/// updated (0 if the id doesn't exist).
|
||||
pub fn set_score(&self, run_id: i64, score: f64, scorer: &str) -> Result<usize> {
|
||||
let n = self.conn.execute(
|
||||
"UPDATE runs SET quality_score=?1, scorer=?2 WHERE id=?3",
|
||||
params![score, scorer, run_id],
|
||||
)?;
|
||||
Ok(n)
|
||||
}
|
||||
|
||||
// ── Read API surface (consumed by api.rs) ─────────────────────────
|
||||
|
||||
/// Total recorded runs (for `/api/health`).
|
||||
@@ -379,7 +672,10 @@ impl Store {
|
||||
"SELECT id, ts, target_name, hostname, git_sha, build_timestamp, package_version,
|
||||
model_id, harness, scenario_id, prompt_size_approx, prompt_tokens_actual,
|
||||
max_tokens, ttft_s, decode_tps, total_s, completion_tokens, ok, error,
|
||||
gpus_json
|
||||
gpus_json, prefill_ms, decode_ms, prefill_tokens,
|
||||
vram_used_mb, gpu_util_pct, gpu_temp_c,
|
||||
concurrency, ttft_p95_s, queue_wait_ms, rejected,
|
||||
swap_unload_ms, swap_load_ms
|
||||
FROM runs",
|
||||
);
|
||||
let mut conds: Vec<String> = Vec::new();
|
||||
@@ -435,6 +731,18 @@ impl Store {
|
||||
completion_tokens: r.get(16)?,
|
||||
ok: r.get::<_, i64>(17)? != 0,
|
||||
error: r.get(18)?,
|
||||
prefill_ms: r.get(20)?,
|
||||
decode_ms: r.get(21)?,
|
||||
prefill_tokens: r.get(22)?,
|
||||
vram_used_mb: r.get(23)?,
|
||||
gpu_util_pct: r.get(24)?,
|
||||
gpu_temp_c: r.get(25)?,
|
||||
concurrency: r.get(26)?,
|
||||
ttft_p95_s: r.get(27)?,
|
||||
queue_wait_ms: r.get(28)?,
|
||||
rejected: r.get(29)?,
|
||||
swap_unload_ms: r.get(30)?,
|
||||
swap_load_ms: r.get(31)?,
|
||||
})
|
||||
})?
|
||||
.collect::<rusqlite::Result<_>>()?;
|
||||
@@ -554,6 +862,18 @@ pub struct RunRow {
|
||||
pub decode_tps: Option<f64>,
|
||||
pub total_s: Option<f64>,
|
||||
pub completion_tokens: Option<u64>,
|
||||
pub prefill_ms: Option<u64>,
|
||||
pub decode_ms: Option<u64>,
|
||||
pub prefill_tokens: Option<u64>,
|
||||
pub vram_used_mb: Option<u64>,
|
||||
pub gpu_util_pct: Option<u64>,
|
||||
pub gpu_temp_c: Option<u64>,
|
||||
pub concurrency: Option<u64>,
|
||||
pub ttft_p95_s: Option<f64>,
|
||||
pub queue_wait_ms: Option<f64>,
|
||||
pub rejected: Option<u64>,
|
||||
pub swap_unload_ms: Option<u64>,
|
||||
pub swap_load_ms: Option<u64>,
|
||||
pub ok: bool,
|
||||
pub error: Option<String>,
|
||||
}
|
||||
@@ -569,6 +889,16 @@ struct RawRow {
|
||||
total_s: Option<f64>,
|
||||
prompt_tokens_actual: Option<u64>,
|
||||
gpus_json: Option<String>,
|
||||
prefill_ms: Option<u64>,
|
||||
decode_ms: Option<u64>,
|
||||
prefill_tokens: Option<u64>,
|
||||
vram_used_mb: Option<u64>,
|
||||
gpu_util_pct: Option<u64>,
|
||||
gpu_temp_c: Option<u64>,
|
||||
concurrency: Option<u64>,
|
||||
ttft_p95_s: Option<f64>,
|
||||
queue_wait_ms: Option<f64>,
|
||||
rejected: Option<u64>,
|
||||
}
|
||||
|
||||
/// An aggregated cell ready for the report table.
|
||||
@@ -583,11 +913,100 @@ pub struct ReportRow {
|
||||
pub ttft_s_median: Option<f64>,
|
||||
pub decode_tps_median: Option<f64>,
|
||||
pub total_s_median: Option<f64>,
|
||||
/// Latency tail percentiles — where batch-1 pain actually shows up, and
|
||||
/// invisible behind a bare median. p95/p99 nearest-rank; with few
|
||||
/// samples they collapse toward the max (honest, not interpolated).
|
||||
pub ttft_s_p95: Option<f64>,
|
||||
pub ttft_s_p99: Option<f64>,
|
||||
pub total_s_p95: Option<f64>,
|
||||
pub total_s_p99: Option<f64>,
|
||||
/// Server-measured prefill/decode split (#85). `prefill_tps_median` is
|
||||
/// the true prompt-encoding rate (prefill_tokens / prefill_ms),
|
||||
/// complementing `decode_tps_median` (the generation rate).
|
||||
pub prefill_ms_median: Option<f64>,
|
||||
pub decode_ms_median: Option<f64>,
|
||||
pub prefill_tps_median: Option<f64>,
|
||||
/// GPU telemetry sampled from /health around the run (#87). `vram_used_mb`
|
||||
/// is the node sum; `vram_total_mb` (from discovery) lets the report show
|
||||
/// real headroom — the "2/3 used" hunch as a number. util/temp are the
|
||||
/// hottest device. All `None` for non-neuron targets.
|
||||
pub vram_used_mb_median: Option<f64>,
|
||||
pub vram_total_mb: Option<u64>,
|
||||
pub gpu_util_pct_median: Option<f64>,
|
||||
pub gpu_temp_c_median: Option<f64>,
|
||||
/// Concurrency / agentic-load burst metrics (#89). `concurrency` is the
|
||||
/// burst width (constant per cell). `ttft_p95_load_s` is the within-burst
|
||||
/// TTFT tail; `queue_wait_ms_median` the admission wait; `rejected_median`
|
||||
/// the per-burst shed count. All `None` for non-concurrency scenarios.
|
||||
pub concurrency: Option<u64>,
|
||||
pub ttft_p95_load_s: Option<f64>,
|
||||
pub queue_wait_ms_median: Option<f64>,
|
||||
pub rejected_median: Option<f64>,
|
||||
pub samples: usize,
|
||||
/// Public-facing resource name (the host's GPU(s)), e.g. "2× RTX 5090".
|
||||
pub gpu: Option<String>,
|
||||
}
|
||||
|
||||
/// One context-length scaling curve for a (target, model) at its latest
|
||||
/// build — the points ordered by prompt size, plus the decode-flatness
|
||||
/// summary (#88).
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct ScalingCurve {
|
||||
pub target_name: String,
|
||||
pub model_id: String,
|
||||
pub git_sha: String,
|
||||
pub gpu: Option<String>,
|
||||
pub points: Vec<ScalingPoint>,
|
||||
/// decode tok/s at the largest context ÷ at the smallest. ~1.0 = flat
|
||||
/// (GDN O(1) decode); <1 quantifies the drop-off. `None` with <2 points.
|
||||
pub decode_flatness: Option<f64>,
|
||||
}
|
||||
|
||||
/// Cold-load / model-swap cost for a (target, model) at its latest build
|
||||
/// (#90): the reload latency and the cold first-request after it.
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct SwapCost {
|
||||
pub target_name: String,
|
||||
pub model_id: String,
|
||||
pub git_sha: String,
|
||||
pub gpu: Option<String>,
|
||||
/// Time to free the model (`POST /models/unload`), ms.
|
||||
pub unload_ms_median: Option<f64>,
|
||||
/// Time to reload it (`POST /models/load`, synchronous), ms — the
|
||||
/// headline swap cost feeding the vision cold-swap policy (F4e).
|
||||
pub load_ms_median: Option<f64>,
|
||||
/// TTFT of the first request after reload (cold caches), seconds.
|
||||
pub cold_ttft_s_median: Option<f64>,
|
||||
/// Total wall-clock of that cold first request, seconds.
|
||||
pub cold_total_s_median: Option<f64>,
|
||||
pub samples: usize,
|
||||
}
|
||||
|
||||
/// One capability-probe run (#91): the stored output artifact plus its
|
||||
/// quality score (null until scored manually or by a future LLM-judge).
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct CapabilityRun {
|
||||
pub id: i64,
|
||||
pub ts: String,
|
||||
pub target_name: String,
|
||||
pub model_id: String,
|
||||
pub scenario_id: String,
|
||||
pub git_sha: String,
|
||||
pub quality_score: Option<f64>,
|
||||
pub scorer: Option<String>,
|
||||
pub artifact: Option<String>,
|
||||
}
|
||||
|
||||
/// One point on a [`ScalingCurve`]: the throughput at a given context size.
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct ScalingPoint {
|
||||
pub prompt_size: u32,
|
||||
pub prompt_tokens: Option<u64>,
|
||||
pub prefill_tps: Option<f64>,
|
||||
pub decode_tps: Option<f64>,
|
||||
pub samples: usize,
|
||||
}
|
||||
|
||||
/// Group by (target, model, scenario), keep only the latest SHA's rows
|
||||
/// (latest = the SHA of the last-inserted row, since input is id-ordered),
|
||||
/// and median each metric.
|
||||
@@ -611,6 +1030,11 @@ fn aggregate(raws: Vec<RawRow>) -> Vec<ReportRow> {
|
||||
let latest_sha = rows.last().map(|r| r.git_sha.clone()).unwrap_or_default();
|
||||
let cell: Vec<&RawRow> = rows.iter().filter(|r| r.git_sha == latest_sha).collect();
|
||||
let prompt_size_approx = cell.first().map(|r| r.prompt_size_approx).unwrap_or(0);
|
||||
// Per-row prefill tok/s, derived from the server-measured split.
|
||||
let prefill_tps = |r: &&RawRow| match (r.prefill_tokens, r.prefill_ms) {
|
||||
(Some(tok), Some(ms)) if ms > 0 => Some(tok as f64 * 1000.0 / ms as f64),
|
||||
_ => None,
|
||||
};
|
||||
out.push(ReportRow {
|
||||
target_name,
|
||||
model_id,
|
||||
@@ -621,6 +1045,27 @@ fn aggregate(raws: Vec<RawRow>) -> Vec<ReportRow> {
|
||||
ttft_s_median: median(cell.iter().filter_map(|r| r.ttft_s)),
|
||||
decode_tps_median: median(cell.iter().filter_map(|r| r.decode_tps)),
|
||||
total_s_median: median(cell.iter().filter_map(|r| r.total_s)),
|
||||
ttft_s_p95: percentile(cell.iter().filter_map(|r| r.ttft_s), 95.0),
|
||||
ttft_s_p99: percentile(cell.iter().filter_map(|r| r.ttft_s), 99.0),
|
||||
total_s_p95: percentile(cell.iter().filter_map(|r| r.total_s), 95.0),
|
||||
total_s_p99: percentile(cell.iter().filter_map(|r| r.total_s), 99.0),
|
||||
prefill_ms_median: median(cell.iter().filter_map(|r| r.prefill_ms.map(|m| m as f64))),
|
||||
decode_ms_median: median(cell.iter().filter_map(|r| r.decode_ms.map(|m| m as f64))),
|
||||
prefill_tps_median: median(cell.iter().filter_map(prefill_tps)),
|
||||
vram_used_mb_median: median(
|
||||
cell.iter().filter_map(|r| r.vram_used_mb.map(|v| v as f64)),
|
||||
),
|
||||
vram_total_mb: cell
|
||||
.iter()
|
||||
.find_map(|r| r.gpus_json.as_deref().and_then(gpu_total_vram_mb)),
|
||||
gpu_util_pct_median: median(
|
||||
cell.iter().filter_map(|r| r.gpu_util_pct.map(|v| v as f64)),
|
||||
),
|
||||
gpu_temp_c_median: median(cell.iter().filter_map(|r| r.gpu_temp_c.map(|v| v as f64))),
|
||||
concurrency: cell.iter().find_map(|r| r.concurrency),
|
||||
ttft_p95_load_s: median(cell.iter().filter_map(|r| r.ttft_p95_s)),
|
||||
queue_wait_ms_median: median(cell.iter().filter_map(|r| r.queue_wait_ms)),
|
||||
rejected_median: median(cell.iter().filter_map(|r| r.rejected.map(|v| v as f64))),
|
||||
samples: cell.len(),
|
||||
gpu: cell
|
||||
.iter()
|
||||
@@ -630,6 +1075,19 @@ fn aggregate(raws: Vec<RawRow>) -> Vec<ReportRow> {
|
||||
out
|
||||
}
|
||||
|
||||
/// Node total VRAM in MB, summed across the devices in a run's stored
|
||||
/// `gpus_json` (the discovery `DeviceInfo` list, each with `vram_total_mb`).
|
||||
/// Pairs with the sampled `vram_used_mb` to report real headroom (#87).
|
||||
/// `None` when empty/absent or no device declares a total.
|
||||
fn gpu_total_vram_mb(gpus_json: &str) -> Option<u64> {
|
||||
let devices: Vec<serde_json::Value> = serde_json::from_str(gpus_json).ok()?;
|
||||
let total: u64 = devices
|
||||
.iter()
|
||||
.filter_map(|d| d.get("vram_total_mb").and_then(|v| v.as_u64()))
|
||||
.sum();
|
||||
(total > 0).then_some(total)
|
||||
}
|
||||
|
||||
/// Compact GPU label from a run's stored `gpus_json` (the discovery device
|
||||
/// list) — e.g. "2× RTX 5090", "RTX 4090". `None` when empty/absent. Used
|
||||
/// as the public-facing resource name in place of internal hostnames.
|
||||
@@ -680,6 +1138,22 @@ fn median(values: impl Iterator<Item = f64>) -> Option<f64> {
|
||||
Some((v[lo] + v[hi]) / 2.0)
|
||||
}
|
||||
|
||||
/// Nearest-rank percentile (`p` in 0..=100). Chosen over interpolation
|
||||
/// because bench cells hold only a handful of samples: with n=5, p95/p99
|
||||
/// resolve to the max, which honestly says "this is the worst we saw"
|
||||
/// rather than inventing a value between samples we never observed.
|
||||
fn percentile(values: impl Iterator<Item = f64>, p: f64) -> Option<f64> {
|
||||
let mut v: Vec<f64> = values.collect();
|
||||
if v.is_empty() {
|
||||
return None;
|
||||
}
|
||||
v.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||
// rank = ceil(p/100 * n), clamped to [1, n]; index is rank-1.
|
||||
let rank = (p / 100.0 * v.len() as f64).ceil() as usize;
|
||||
let idx = rank.clamp(1, v.len()) - 1;
|
||||
Some(v[idx])
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -717,6 +1191,21 @@ mod tests {
|
||||
decode_tps: Some(50.0),
|
||||
total_s: Some(1.0),
|
||||
completion_tokens: Some(50),
|
||||
prefill_ms: Some(200),
|
||||
decode_ms: Some(1000),
|
||||
prefill_tokens: Some(130),
|
||||
vram_used_mb: Some(42000),
|
||||
gpu_util_pct: Some(88),
|
||||
gpu_temp_c: Some(64),
|
||||
concurrency: None,
|
||||
ttft_p95_s: None,
|
||||
queue_wait_ms: None,
|
||||
rejected: None,
|
||||
swap_unload_ms: None,
|
||||
swap_load_ms: None,
|
||||
artifact: None,
|
||||
quality_score: None,
|
||||
scorer: None,
|
||||
ok,
|
||||
error: if ok { None } else { Some("boom".into()) },
|
||||
}
|
||||
@@ -755,6 +1244,197 @@ mod tests {
|
||||
assert!((rows[0].ttft_s_median.unwrap() - 0.3).abs() < 1e-9);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn report_surfaces_percentiles_and_prefill_split() {
|
||||
let s = Store::open_in_memory().unwrap();
|
||||
// Five samples on one cell with spread TTFT so percentiles differ
|
||||
// from the median, plus a server-measured prefill/decode split.
|
||||
for (i, ttft) in [0.10, 0.12, 0.14, 0.16, 0.50].iter().enumerate() {
|
||||
let mut r = rec("beast", "sha", "m", "chat:128", true);
|
||||
r.ttft_s = Some(*ttft);
|
||||
r.total_s = Some(ttft + 1.0);
|
||||
r.prefill_ms = Some(200 + i as u64);
|
||||
r.prefill_tokens = Some(400);
|
||||
s.insert_run(&r).unwrap();
|
||||
}
|
||||
let rows = s.report_rows().unwrap();
|
||||
assert_eq!(rows.len(), 1);
|
||||
let row = &rows[0];
|
||||
assert_eq!(row.samples, 5);
|
||||
// p50 is the middle value; p95/p99 (nearest-rank, n=5) hit the max.
|
||||
assert!((row.ttft_s_median.unwrap() - 0.14).abs() < 1e-9);
|
||||
assert!((row.ttft_s_p95.unwrap() - 0.50).abs() < 1e-9);
|
||||
assert!((row.ttft_s_p99.unwrap() - 0.50).abs() < 1e-9);
|
||||
// prefill tok/s = 400 tok / ~0.2 s ≈ 2000 tok/s.
|
||||
assert!(row.prefill_tps_median.unwrap() > 1900.0);
|
||||
assert!(row.prefill_ms_median.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn report_surfaces_vram_and_gpu_telemetry() {
|
||||
let s = Store::open_in_memory().unwrap();
|
||||
let mut r = rec("beast", "sha", "m", "chat:128", true);
|
||||
// Node total VRAM from discovery devices → headroom denominator.
|
||||
r.gpus_json =
|
||||
Some(r#"[{"name":"RTX 5090","vram_total_mb":32000},{"name":"RTX 5090","vram_total_mb":32000}]"#.into());
|
||||
r.vram_used_mb = Some(42000);
|
||||
r.gpu_util_pct = Some(90);
|
||||
r.gpu_temp_c = Some(66);
|
||||
s.insert_run(&r).unwrap();
|
||||
let rows = s.report_rows().unwrap();
|
||||
let row = &rows[0];
|
||||
assert_eq!(row.vram_used_mb_median, Some(42000.0));
|
||||
assert_eq!(row.vram_total_mb, Some(64000)); // 2× 32000
|
||||
assert_eq!(row.gpu_util_pct_median, Some(90.0));
|
||||
assert_eq!(row.gpu_temp_c_median, Some(66.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn report_surfaces_concurrency_burst_metrics() {
|
||||
let s = Store::open_in_memory().unwrap();
|
||||
// Two concurrency:8 burst-runs with shed load and a queue-wait tail.
|
||||
for (qw, rej) in [(120.0, 1u32), (180.0, 3u32)] {
|
||||
let mut r = rec("beast", "sha", "m", "concurrency:8", true);
|
||||
r.concurrency = Some(8);
|
||||
r.ttft_p95_s = Some(0.9);
|
||||
r.queue_wait_ms = Some(qw);
|
||||
r.rejected = Some(rej);
|
||||
s.insert_run(&r).unwrap();
|
||||
}
|
||||
let row = s
|
||||
.report_rows()
|
||||
.unwrap()
|
||||
.into_iter()
|
||||
.find(|r| r.scenario_id == "concurrency:8")
|
||||
.unwrap();
|
||||
assert_eq!(row.concurrency, Some(8));
|
||||
assert_eq!(row.queue_wait_ms_median, Some(150.0)); // median(120,180)
|
||||
assert_eq!(row.rejected_median, Some(2.0)); // median(1,3)
|
||||
assert_eq!(row.ttft_p95_load_s, Some(0.9));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capability_runs_store_artifact_and_accept_scores() {
|
||||
let s = Store::open_in_memory().unwrap();
|
||||
let mut r = rec("beast", "sha", "m", "capability:plan", true);
|
||||
r.artifact = Some("a thorough implementation plan...".into());
|
||||
s.insert_run(&r).unwrap();
|
||||
// A non-capability row must not appear in capability_runs.
|
||||
s.insert_run(&rec("beast", "sha", "m", "chat:128", true))
|
||||
.unwrap();
|
||||
|
||||
let unscored = s.capability_runs(true).unwrap();
|
||||
assert_eq!(unscored.len(), 1);
|
||||
let id = unscored[0].id;
|
||||
assert!(unscored[0].artifact.as_deref().unwrap().contains("plan"));
|
||||
assert!(unscored[0].quality_score.is_none());
|
||||
|
||||
// Score it; it then drops off the unscored worklist.
|
||||
assert_eq!(s.set_score(id, 7.5, "manual").unwrap(), 1);
|
||||
assert!(s.capability_runs(true).unwrap().is_empty());
|
||||
let all = s.capability_runs(false).unwrap();
|
||||
assert_eq!(all[0].quality_score, Some(7.5));
|
||||
assert_eq!(all[0].scorer.as_deref(), Some("manual"));
|
||||
// Unknown id updates nothing.
|
||||
assert_eq!(s.set_score(999_999, 1.0, "manual").unwrap(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn swap_costs_pivots_swap_rows() {
|
||||
let s = Store::open_in_memory().unwrap();
|
||||
for (unload, load) in [(300u64, 24000u64), (340, 26000)] {
|
||||
let mut r = rec("beast", "sha", "m", "swap", true);
|
||||
r.swap_unload_ms = Some(unload);
|
||||
r.swap_load_ms = Some(load);
|
||||
r.ttft_s = Some(2.5); // cold first-request
|
||||
r.total_s = Some(5.0);
|
||||
s.insert_run(&r).unwrap();
|
||||
}
|
||||
// A non-swap row must be ignored by swap_costs.
|
||||
s.insert_run(&rec("beast", "sha", "m", "chat:128", true))
|
||||
.unwrap();
|
||||
let costs = s.swap_costs().unwrap();
|
||||
assert_eq!(costs.len(), 1);
|
||||
let c = &costs[0];
|
||||
assert_eq!(c.unload_ms_median, Some(320.0)); // median(300,340)
|
||||
assert_eq!(c.load_ms_median, Some(25000.0)); // median(24000,26000)
|
||||
assert_eq!(c.cold_ttft_s_median, Some(2.5));
|
||||
assert_eq!(c.samples, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scaling_pivots_chat_cells_and_computes_flatness() {
|
||||
let s = Store::open_in_memory().unwrap();
|
||||
// Two context points for one model: decode tok/s 50 @128, 45 @4096.
|
||||
let mut small = rec("beast", "sha", "m", "chat:128", true);
|
||||
small.prompt_size_approx = 128;
|
||||
small.decode_tps = Some(50.0);
|
||||
small.prefill_ms = Some(100);
|
||||
small.prefill_tokens = Some(128);
|
||||
s.insert_run(&small).unwrap();
|
||||
let mut big = rec("beast", "sha", "m", "chat:4096", true);
|
||||
big.prompt_size_approx = 4096;
|
||||
big.decode_tps = Some(45.0);
|
||||
big.prefill_ms = Some(1000);
|
||||
big.prefill_tokens = Some(4096);
|
||||
s.insert_run(&big).unwrap();
|
||||
// A concurrency cell must NOT leak into the scaling curve.
|
||||
let mut conc = rec("beast", "sha", "m", "concurrency:8", true);
|
||||
conc.concurrency = Some(8);
|
||||
s.insert_run(&conc).unwrap();
|
||||
|
||||
let curves = s.scaling().unwrap();
|
||||
assert_eq!(curves.len(), 1);
|
||||
let c = &curves[0];
|
||||
assert_eq!(c.points.len(), 2); // only the two chat:<n> points
|
||||
assert_eq!(c.points[0].prompt_size, 128); // ordered ascending
|
||||
assert_eq!(c.points[1].prompt_size, 4096);
|
||||
// flatness = decode@largest / decode@smallest = 45/50 = 0.9
|
||||
assert!((c.decode_flatness.unwrap() - 0.9).abs() < 1e-9);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gpu_total_vram_sums_devices() {
|
||||
let j = r#"[{"vram_total_mb":32000},{"vram_total_mb":32000}]"#;
|
||||
assert_eq!(gpu_total_vram_mb(j), Some(64000));
|
||||
assert_eq!(gpu_total_vram_mb("[]"), None);
|
||||
assert_eq!(gpu_total_vram_mb("not json"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn percentile_nearest_rank() {
|
||||
let vals = || [1.0, 2.0, 3.0, 4.0, 5.0].into_iter();
|
||||
assert_eq!(percentile(vals(), 50.0), Some(3.0));
|
||||
assert_eq!(percentile(vals(), 95.0), Some(5.0));
|
||||
assert_eq!(percentile(vals(), 99.0), Some(5.0));
|
||||
assert_eq!(percentile(std::iter::empty(), 95.0), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn migration_is_idempotent_and_backfills() {
|
||||
// A DB whose `runs` table predates the prefill columns: create the
|
||||
// pre-#85 shape, insert a row, then run ensure_columns twice.
|
||||
let conn = Connection::open_in_memory().unwrap();
|
||||
conn.execute_batch(
|
||||
"CREATE TABLE runs (id INTEGER PRIMARY KEY, ttft_s REAL);
|
||||
INSERT INTO runs (ttft_s) VALUES (0.1);",
|
||||
)
|
||||
.unwrap();
|
||||
for _ in 0..2 {
|
||||
Store::ensure_columns(
|
||||
&conn,
|
||||
"runs",
|
||||
&[("prefill_ms", "INTEGER"), ("decode_ms", "INTEGER")],
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
// Columns now exist and the old row reads them back as NULL.
|
||||
let got: Option<i64> = conn
|
||||
.query_row("SELECT prefill_ms FROM runs", [], |r| r.get(0))
|
||||
.unwrap();
|
||||
assert_eq!(got, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gpu_label_formats() {
|
||||
let two = r#"[{"name":"NVIDIA GeForce RTX 5090"},{"name":"NVIDIA GeForce RTX 5090"}]"#;
|
||||
|
||||
@@ -9,11 +9,11 @@
|
||||
|
||||
use crate::client::TargetClient;
|
||||
use crate::config::{BenchConfig, TargetConfig, TargetKind};
|
||||
use crate::scenario::{RunCtx, build_scenarios};
|
||||
use crate::scenario::{RunCtx, ScenarioMetrics, build_scenarios};
|
||||
use crate::store::{RunRecord, Store};
|
||||
use anyhow::Result;
|
||||
use cortex_core::build_info::BuildInfo;
|
||||
use cortex_core::discovery::DiscoveryResponse;
|
||||
use cortex_core::discovery::{DiscoveryResponse, HealthResponse};
|
||||
use cortex_core::harness::ModelInfo;
|
||||
|
||||
/// helexa-bench's own build version.
|
||||
@@ -38,6 +38,38 @@ pub struct SweepSummary {
|
||||
pub targets_unreachable: usize,
|
||||
}
|
||||
|
||||
/// Node-level GPU telemetry folded from one `/health` snapshot (#87):
|
||||
/// VRAM used summed across the node's devices, and the hottest/busiest
|
||||
/// single device for utilization and temperature.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct HealthAgg {
|
||||
vram_used_mb: u64,
|
||||
gpu_util_pct: u32,
|
||||
gpu_temp_c: u32,
|
||||
}
|
||||
|
||||
/// Cold-load / model-swap timing for one measure_swap cycle (#90).
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct SwapTiming {
|
||||
unload_ms: u64,
|
||||
load_ms: u64,
|
||||
}
|
||||
|
||||
impl HealthAgg {
|
||||
fn from_health(h: &HealthResponse) -> Self {
|
||||
HealthAgg {
|
||||
vram_used_mb: h.devices.iter().map(|d| d.vram_used_mb).sum(),
|
||||
gpu_util_pct: h
|
||||
.devices
|
||||
.iter()
|
||||
.map(|d| d.utilization_pct)
|
||||
.max()
|
||||
.unwrap_or(0),
|
||||
gpu_temp_c: h.devices.iter().map(|d| d.temp_c).max().unwrap_or(0),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Sweeper {
|
||||
cfg: BenchConfig,
|
||||
client: TargetClient,
|
||||
@@ -72,6 +104,105 @@ impl Sweeper {
|
||||
}
|
||||
}
|
||||
|
||||
/// Deliberate cold-load / model-swap cost measurement (#90), invoked by
|
||||
/// the `swap-cost` subcommand — **never** the continuous sweep. For each
|
||||
/// neuron target and each currently-warm model: unload it, time the
|
||||
/// reload, then time a cold first request. This takes the model offline
|
||||
/// for the reload, so it is an explicit operator action (maintenance
|
||||
/// window), recorded under `scenario_id = "swap"`.
|
||||
pub async fn swap_cost_once(&self) -> Result<SweepSummary> {
|
||||
let mut summary = SweepSummary::default();
|
||||
for target in &self.cfg.targets {
|
||||
if target.kind != TargetKind::Neuron {
|
||||
continue; // load/unload is a neuron-native operation
|
||||
}
|
||||
let build = match self.client.fetch_version(target).await {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
summary.targets_unreachable += 1;
|
||||
tracing::warn!(target = %target.name, error = %format!("{e:#}"), "swap: target unreachable");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let discovery = self.client.fetch_discovery(target).await.unwrap_or(None);
|
||||
let models = self.client.warm_models(target).await.unwrap_or_default();
|
||||
for model in &models {
|
||||
match self
|
||||
.measure_swap(target, &build, discovery.as_ref(), model)
|
||||
.await
|
||||
{
|
||||
Ok(()) => summary.measured += 1,
|
||||
Err(e) => {
|
||||
summary.failed += 1;
|
||||
tracing::warn!(target = %target.name, model = %model.id, error = %format!("{e:#}"), "swap: measurement failed");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
/// Unload → timed reload → timed cold first request for one model.
|
||||
async fn measure_swap(
|
||||
&self,
|
||||
target: &TargetConfig,
|
||||
build: &BuildInfo,
|
||||
discovery: Option<&DiscoveryResponse>,
|
||||
model: &ModelInfo,
|
||||
) -> Result<()> {
|
||||
let spec = TargetClient::spec_from_info(model)?;
|
||||
tracing::warn!(target = %target.name, model = %model.id, "swap: unloading (model goes offline until reload)");
|
||||
|
||||
let t0 = std::time::Instant::now();
|
||||
self.client.unload_model(target, &model.id).await?;
|
||||
let unload_ms = t0.elapsed().as_millis() as u64;
|
||||
|
||||
let t1 = std::time::Instant::now();
|
||||
self.client.load_model(target, &spec).await?;
|
||||
let load_ms = t1.elapsed().as_millis() as u64;
|
||||
tracing::info!(target = %target.name, model = %model.id, unload_ms, load_ms, "swap: reloaded; measuring cold first request");
|
||||
|
||||
// Cold first request — caches empty straight after the load.
|
||||
let ctx = RunCtx {
|
||||
client: self.client.http(),
|
||||
chat_url: self.client.chat_url(target),
|
||||
model_id: model.id.clone(),
|
||||
max_tokens: self.cfg.scenarios.max_tokens,
|
||||
timeout: self.cfg.bench.request_timeout(),
|
||||
};
|
||||
let cold = crate::scenario::cold_probe(&ctx).await;
|
||||
let swap = SwapTiming { unload_ms, load_ms };
|
||||
let rec = match &cold {
|
||||
Ok(m) => self.build_record(
|
||||
target,
|
||||
build,
|
||||
discovery,
|
||||
model,
|
||||
"swap",
|
||||
0,
|
||||
Ok(m),
|
||||
None,
|
||||
Some(swap),
|
||||
),
|
||||
Err(e) => {
|
||||
let msg = format!("{e:#}");
|
||||
self.build_record(
|
||||
target,
|
||||
build,
|
||||
discovery,
|
||||
model,
|
||||
"swap",
|
||||
0,
|
||||
Err(&msg),
|
||||
None,
|
||||
Some(swap),
|
||||
)
|
||||
}
|
||||
};
|
||||
self.store.insert_run(&rec)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// One full pass over all targets.
|
||||
pub async fn run_once(&self) -> Result<SweepSummary> {
|
||||
let mut summary = SweepSummary::default();
|
||||
@@ -131,7 +262,21 @@ impl Sweeper {
|
||||
}
|
||||
|
||||
for i in 0..need {
|
||||
match scenario.run(&ctx).await {
|
||||
let result = scenario.run(&ctx).await;
|
||||
// Sample GPU telemetry right after the run, while the
|
||||
// model is loaded and decode VRAM is at its recent peak
|
||||
// (#87). neuron's /health is ~5s-cached, so this is a
|
||||
// coarse high-water proxy, not an instantaneous peak — but
|
||||
// it's the headroom signal we can read over the wire. A
|
||||
// flaky /health degrades to None, never a failed run.
|
||||
let health = self
|
||||
.client
|
||||
.fetch_health(target)
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
.map(|h| HealthAgg::from_health(&h));
|
||||
match result {
|
||||
Ok(m) => {
|
||||
let rec = self.build_record(
|
||||
target,
|
||||
@@ -141,6 +286,8 @@ impl Sweeper {
|
||||
scenario.id(),
|
||||
scenario.prompt_size(),
|
||||
Ok(&m),
|
||||
health,
|
||||
None,
|
||||
);
|
||||
self.store.insert_run(&rec)?;
|
||||
summary.measured += 1;
|
||||
@@ -160,6 +307,8 @@ impl Sweeper {
|
||||
scenario.id(),
|
||||
scenario.prompt_size(),
|
||||
Err(&msg),
|
||||
health,
|
||||
None,
|
||||
);
|
||||
self.store.insert_run(&rec)?;
|
||||
summary.failed += 1;
|
||||
@@ -186,19 +335,14 @@ impl Sweeper {
|
||||
scenario_id: &str,
|
||||
prompt_size: u32,
|
||||
result: Result<&crate::scenario::ScenarioMetrics, &str>,
|
||||
health: Option<HealthAgg>,
|
||||
swap: Option<SwapTiming>,
|
||||
) -> RunRecord {
|
||||
let (ok, error, ttft, decode, total, prompt_tokens, completion) = match result {
|
||||
Ok(m) => (
|
||||
true,
|
||||
None,
|
||||
Some(m.ttft_s),
|
||||
m.decode_tps,
|
||||
Some(m.total_s),
|
||||
m.prompt_tokens,
|
||||
Some(m.completion_tokens),
|
||||
),
|
||||
Err(e) => (false, Some(e.to_string()), None, None, None, None, None),
|
||||
let (m, error): (Option<&ScenarioMetrics>, Option<String>) = match result {
|
||||
Ok(m) => (Some(m), None),
|
||||
Err(e) => (None, Some(e.to_string())),
|
||||
};
|
||||
let ok = m.is_some();
|
||||
|
||||
RunRecord {
|
||||
ts: chrono::Utc::now().to_rfc3339(),
|
||||
@@ -230,12 +374,29 @@ impl Sweeper {
|
||||
.unwrap_or_else(|_| "[]".to_string()),
|
||||
scenario_id: scenario_id.to_string(),
|
||||
prompt_size_approx: prompt_size,
|
||||
prompt_tokens_actual: prompt_tokens,
|
||||
prompt_tokens_actual: m.and_then(|m| m.prompt_tokens),
|
||||
max_tokens: self.cfg.scenarios.max_tokens,
|
||||
ttft_s: ttft,
|
||||
decode_tps: decode,
|
||||
total_s: total,
|
||||
completion_tokens: completion,
|
||||
ttft_s: m.map(|m| m.ttft_s),
|
||||
decode_tps: m.and_then(|m| m.decode_tps),
|
||||
total_s: m.map(|m| m.total_s),
|
||||
completion_tokens: m.map(|m| m.completion_tokens),
|
||||
prefill_ms: m.and_then(|m| m.prefill_ms),
|
||||
decode_ms: m.and_then(|m| m.decode_ms),
|
||||
prefill_tokens: m.and_then(|m| m.prefill_tokens),
|
||||
vram_used_mb: health.map(|h| h.vram_used_mb),
|
||||
gpu_util_pct: health.map(|h| h.gpu_util_pct),
|
||||
gpu_temp_c: health.map(|h| h.gpu_temp_c),
|
||||
concurrency: m.and_then(|m| m.concurrency),
|
||||
ttft_p95_s: m.and_then(|m| m.ttft_p95_s),
|
||||
queue_wait_ms: m.and_then(|m| m.queue_wait_ms_median),
|
||||
rejected: m.and_then(|m| m.rejected),
|
||||
swap_unload_ms: swap.map(|s| s.unload_ms),
|
||||
swap_load_ms: swap.map(|s| s.load_ms),
|
||||
// Capability artifact (#91); score/scorer are attached later by
|
||||
// the `score` subcommand or a future LLM-judge.
|
||||
artifact: m.and_then(|m| m.artifact.clone()),
|
||||
quality_score: None,
|
||||
scorer: None,
|
||||
ok,
|
||||
error,
|
||||
}
|
||||
|
||||
@@ -46,6 +46,21 @@ fn rec(
|
||||
decode_tps: if ok { Some(30.0) } else { None },
|
||||
total_s: if ok { Some(2.0) } else { None },
|
||||
completion_tokens: if ok { Some(60) } else { None },
|
||||
prefill_ms: if ok { Some(150) } else { None },
|
||||
decode_ms: if ok { Some(1800) } else { None },
|
||||
prefill_tokens: if ok { Some(130) } else { None },
|
||||
vram_used_mb: if ok { Some(42000) } else { None },
|
||||
gpu_util_pct: if ok { Some(85) } else { None },
|
||||
gpu_temp_c: if ok { Some(63) } else { None },
|
||||
concurrency: None,
|
||||
ttft_p95_s: None,
|
||||
queue_wait_ms: None,
|
||||
rejected: None,
|
||||
swap_unload_ms: None,
|
||||
swap_load_ms: None,
|
||||
artifact: None,
|
||||
quality_score: None,
|
||||
scorer: None,
|
||||
ok,
|
||||
error: if ok { None } else { Some("boom".into()) },
|
||||
}
|
||||
|
||||
@@ -89,6 +89,9 @@ fn config_for(endpoint: String, db_path: String) -> BenchConfig {
|
||||
scenarios: ScenarioConfig {
|
||||
prompt_sizes: vec![128], // single scenario keeps assertions simple
|
||||
max_tokens: 16,
|
||||
concurrency_levels: Vec::new(),
|
||||
concurrency_prompt_tokens: 512,
|
||||
capability_probes: Vec::new(),
|
||||
},
|
||||
api: Default::default(),
|
||||
targets: vec![TargetConfig {
|
||||
|
||||
@@ -55,6 +55,12 @@ pub fn aggregate_models(topology: &HashMap<String, CortexTopology>) -> Vec<Corte
|
||||
|
||||
let mut out: Vec<CortexModelEntry> = merged.into_values().collect();
|
||||
out.sort_by(|a, b| a.id.cmp(&b.id));
|
||||
// Re-derive the flat ecosystem fields (#78) from the merged (tightest)
|
||||
// limit — the values deserialized from each cortex are per-operator and
|
||||
// may not match the federation-wide merge.
|
||||
for e in &mut out {
|
||||
e.sync_flat_limit();
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
@@ -74,6 +80,10 @@ fn router_entry(cortex: &str, e: &CortexModelEntry) -> CortexModelEntry {
|
||||
cost: e.cost.clone(),
|
||||
tool_call: e.tool_call,
|
||||
reasoning: e.reasoning,
|
||||
// Derived from `limit` by the final sync pass in aggregate_models.
|
||||
max_model_len: None,
|
||||
max_input_tokens: None,
|
||||
max_output_tokens: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -151,6 +161,9 @@ mod tests {
|
||||
cost: None,
|
||||
tool_call: false,
|
||||
reasoning: false,
|
||||
max_model_len: None,
|
||||
max_input_tokens: None,
|
||||
max_output_tokens: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -239,5 +252,9 @@ mod tests {
|
||||
assert_eq!(out.len(), 1);
|
||||
assert_eq!(out[0].limit.as_ref().unwrap().context, 16_384);
|
||||
assert_eq!(out[0].cost.as_ref().unwrap().input, 0.20);
|
||||
// Flat #78 fields re-derived from the merged (tightest) limit.
|
||||
assert_eq!(out[0].max_model_len, Some(16_384));
|
||||
assert_eq!(out[0].max_input_tokens, None);
|
||||
assert_eq!(out[0].max_output_tokens, Some(4096));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -39,6 +39,9 @@ fn model_entry(loaded: bool, feasible: bool) -> CortexModelEntry {
|
||||
cost: None,
|
||||
tool_call: false,
|
||||
reasoning: false,
|
||||
max_model_len: None,
|
||||
max_input_tokens: None,
|
||||
max_output_tokens: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -94,13 +94,24 @@ impl AdmissionController {
|
||||
/// overall queue is full or the principal is over its fair-share cap —
|
||||
/// then waits up to `max_wait` for an in-flight slot. The returned permit
|
||||
/// must be held for the request's lifetime; dropping it frees the slots.
|
||||
///
|
||||
/// CANCELLATION SAFETY: the semaphore wait below is where a client
|
||||
/// disconnect lands — axum drops the request future mid-await. The
|
||||
/// reservation therefore lives in a RAII [`PendingReservation`] taken
|
||||
/// BEFORE the await: if this future is dropped while queued, the
|
||||
/// guard's Drop rolls the counts back. (The original version
|
||||
/// incremented raw counters and only decremented on the timeout
|
||||
/// branch — every abandoned wait leaked a `pending` + per-principal
|
||||
/// slot, ratcheting the model into a permanent instant-429 state
|
||||
/// under client retry storms. Observed live 2026-07-02:
|
||||
/// `queue_depth: 1` pinned on an idle model.)
|
||||
pub async fn enter(
|
||||
&self,
|
||||
principal: Option<&str>,
|
||||
) -> Result<AdmissionPermit, AdmissionRejection> {
|
||||
// Decision + reservation under one brief lock so concurrent callers
|
||||
// can't both slip past the thresholds. No await is held here.
|
||||
{
|
||||
let reservation = {
|
||||
let mut st = self.state.lock().expect("admission state poisoned");
|
||||
if st.pending >= self.max_pending {
|
||||
return Err(AdmissionRejection::QueueFull {
|
||||
@@ -119,31 +130,25 @@ impl AdmissionController {
|
||||
if let Some(p) = principal {
|
||||
*st.per_principal.entry(p.to_string()).or_insert(0) += 1;
|
||||
}
|
||||
}
|
||||
PendingReservation {
|
||||
state: Arc::clone(&self.state),
|
||||
principal: principal.map(str::to_string),
|
||||
}
|
||||
};
|
||||
|
||||
match tokio::time::timeout(self.max_wait, Arc::clone(&self.slots).acquire_owned()).await {
|
||||
Ok(Ok(permit)) => Ok(AdmissionPermit {
|
||||
_permit: permit,
|
||||
state: Arc::clone(&self.state),
|
||||
principal: principal.map(str::to_string),
|
||||
_reservation: reservation,
|
||||
}),
|
||||
// Semaphore is never closed; treat a closed/elapsed wait the
|
||||
// same. `reservation` drops here, rolling back the counts.
|
||||
Ok(Err(_)) | Err(_) => Err(AdmissionRejection::Timeout {
|
||||
retry_after_secs: self.retry_hint(self.max_pending),
|
||||
}),
|
||||
// Semaphore is never closed; treat a closed/elapsed wait the same.
|
||||
Ok(Err(_)) | Err(_) => {
|
||||
self.release(principal);
|
||||
Err(AdmissionRejection::Timeout {
|
||||
retry_after_secs: self.retry_hint(self.max_pending),
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Roll back a reserved-but-not-admitted slot (wait timed out).
|
||||
fn release(&self, principal: Option<&str>) {
|
||||
let mut st = self.state.lock().expect("admission state poisoned");
|
||||
st.pending = st.pending.saturating_sub(1);
|
||||
decrement_principal(&mut st.per_principal, principal);
|
||||
}
|
||||
|
||||
/// Requests currently running (holding an in-flight slot).
|
||||
pub fn in_flight(&self) -> usize {
|
||||
self.max_in_flight
|
||||
@@ -177,16 +182,17 @@ fn decrement_principal(map: &mut HashMap<String, usize>, principal: Option<&str>
|
||||
}
|
||||
}
|
||||
|
||||
/// Held for a request's lifetime; frees the in-flight + queue slot (and the
|
||||
/// principal's fair-share slot) on drop.
|
||||
/// RAII accounting for one reserved slot (queued or in-flight): decrements
|
||||
/// `pending` and the principal's fair-share count on drop, whichever way
|
||||
/// the reservation ends — admitted-and-finished, wait timeout, or the
|
||||
/// caller's future being dropped mid-queue (client disconnect).
|
||||
#[derive(Debug)]
|
||||
pub struct AdmissionPermit {
|
||||
_permit: OwnedSemaphorePermit,
|
||||
struct PendingReservation {
|
||||
state: Arc<Mutex<AdmissionState>>,
|
||||
principal: Option<String>,
|
||||
}
|
||||
|
||||
impl Drop for AdmissionPermit {
|
||||
impl Drop for PendingReservation {
|
||||
fn drop(&mut self) {
|
||||
let mut st = self.state.lock().expect("admission state poisoned");
|
||||
st.pending = st.pending.saturating_sub(1);
|
||||
@@ -194,6 +200,14 @@ impl Drop for AdmissionPermit {
|
||||
}
|
||||
}
|
||||
|
||||
/// Held for a request's lifetime; frees the in-flight slot (semaphore
|
||||
/// permit) and the queue + fair-share accounting (reservation) on drop.
|
||||
#[derive(Debug)]
|
||||
pub struct AdmissionPermit {
|
||||
_permit: OwnedSemaphorePermit,
|
||||
_reservation: PendingReservation,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -295,4 +309,65 @@ mod tests {
|
||||
drop(_a1);
|
||||
b.await.unwrap().expect("B is served after A releases");
|
||||
}
|
||||
|
||||
/// Regression for the 2026-07-02 retry-storm incident: a client that
|
||||
/// disconnects while QUEUED drops the `enter()` future mid-await.
|
||||
/// The reservation must roll back — the original implementation
|
||||
/// leaked `pending` + the per-principal count on this path, pinning
|
||||
/// the model in a permanent instant-429 state.
|
||||
#[tokio::test]
|
||||
async fn cancelled_queued_waiter_rolls_back_accounting() {
|
||||
let cfg = AdmissionConfig {
|
||||
max_in_flight: 1,
|
||||
max_queue_depth: 2,
|
||||
max_wait_secs: 30,
|
||||
// Cap 3 lets the runner + both waiters coexist; if the two
|
||||
// cancelled waiters leaked their counts, the principal would
|
||||
// sit at 3 == cap and the post-cancel enter below would hit
|
||||
// PrincipalCap instead of queueing.
|
||||
max_per_principal: 3,
|
||||
};
|
||||
let ctrl = Arc::new(AdmissionController::new(&cfg));
|
||||
let running = ctrl.enter(Some("acct/key")).await.expect("admit running");
|
||||
|
||||
// Two waiters from the same principal park in the queue…
|
||||
let mut waiters = Vec::new();
|
||||
for _ in 0..2 {
|
||||
let c = Arc::clone(&ctrl);
|
||||
waiters.push(tokio::spawn(async move {
|
||||
c.enter(Some("acct/key")).await.map(drop)
|
||||
}));
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
assert_eq!(ctrl.queue_depth(), 2);
|
||||
|
||||
// …and both clients vanish (abort = the dropped request future).
|
||||
for w in &waiters {
|
||||
w.abort();
|
||||
}
|
||||
for w in waiters {
|
||||
let _ = w.await;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
|
||||
assert_eq!(
|
||||
ctrl.queue_depth(),
|
||||
0,
|
||||
"cancelled waiters must not leak queue slots"
|
||||
);
|
||||
|
||||
// The principal's fair-share count must also be clean: with the
|
||||
// runner still holding 1 of its cap of 3, a new request from the
|
||||
// same principal queues instead of hitting PrincipalCap (which a
|
||||
// leak of the two cancelled counts would trigger).
|
||||
let c = Arc::clone(&ctrl);
|
||||
let retry = tokio::spawn(async move { c.enter(Some("acct/key")).await.map(drop) });
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
assert_eq!(ctrl.queue_depth(), 1, "post-cancel request queues normally");
|
||||
drop(running);
|
||||
retry
|
||||
.await
|
||||
.unwrap()
|
||||
.expect("post-cancel request is served — no leaked principal count");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -22,6 +22,7 @@ use super::TextConfig;
|
||||
use super::full_attn::Qwen3_5Attention;
|
||||
use super::linear_attn::GatedDeltaNet;
|
||||
use super::mlp::Qwen3_5MLP;
|
||||
use super::moe::Qwen3_5MoeBlock;
|
||||
use super::rmsnorm::Qwen3_5RmsNorm;
|
||||
use super::rope::RotaryEmbedding;
|
||||
use super::snapshot::LayerKvSnapshot;
|
||||
@@ -35,10 +36,27 @@ enum AttentionKind {
|
||||
Linear(GatedDeltaNet),
|
||||
}
|
||||
|
||||
/// The FFN slot: dense SwiGLU (Qwen3.6) or the high-sparsity MoE block
|
||||
/// (qwen3_next 80B-A3B family, #92), selected per layer by
|
||||
/// [`TextConfig::layer_uses_moe`].
|
||||
enum MlpKind {
|
||||
Dense(Qwen3_5MLP),
|
||||
Moe(Qwen3_5MoeBlock),
|
||||
}
|
||||
|
||||
impl Module for MlpKind {
|
||||
fn forward(&self, x: &Tensor) -> candle_core::Result<Tensor> {
|
||||
match self {
|
||||
MlpKind::Dense(mlp) => mlp.forward(x),
|
||||
MlpKind::Moe(moe) => moe.forward(x),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Qwen3_5DecoderLayer {
|
||||
input_layernorm: Qwen3_5RmsNorm,
|
||||
post_attention_layernorm: Qwen3_5RmsNorm,
|
||||
mlp: Qwen3_5MLP,
|
||||
mlp: MlpKind,
|
||||
attention: AttentionKind,
|
||||
}
|
||||
|
||||
@@ -73,7 +91,11 @@ impl Qwen3_5DecoderLayer {
|
||||
),
|
||||
};
|
||||
|
||||
let mlp = Qwen3_5MLP::load(cfg, &vb.pp("mlp"))?;
|
||||
let mlp = if cfg.layer_uses_moe(layer_idx) {
|
||||
MlpKind::Moe(Qwen3_5MoeBlock::load(cfg, &vb.pp("mlp"))?)
|
||||
} else {
|
||||
MlpKind::Dense(Qwen3_5MLP::load(cfg, &vb.pp("mlp"))?)
|
||||
};
|
||||
let input_layernorm =
|
||||
Qwen3_5RmsNorm::load(&vb.pp("input_layernorm"), cfg.hidden_size, cfg.rms_norm_eps)?;
|
||||
let post_attention_layernorm = Qwen3_5RmsNorm::load(
|
||||
|
||||
@@ -29,6 +29,86 @@ use super::TextConfig;
|
||||
use super::rmsnorm::Qwen3_5RmsNorm;
|
||||
use super::rope::RotaryEmbedding;
|
||||
|
||||
/// Runtime kill-switch for the FlashAttention path (#95):
|
||||
/// `NEURON_FLASH_ATTN=0` (or `false`) forces the eager fallback
|
||||
/// without a rebuild — the A/B lever and the rollback if the kernels
|
||||
/// misbehave on some device. Read once.
|
||||
#[cfg(feature = "flash-attn")]
|
||||
fn flash_attn_enabled() -> bool {
|
||||
static ENABLED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
|
||||
*ENABLED.get_or_init(|| {
|
||||
let on = !std::env::var("NEURON_FLASH_ATTN").is_ok_and(|v| v == "0" || v == "false");
|
||||
tracing::info!(enabled = on, "FlashAttention path (#95)");
|
||||
on
|
||||
})
|
||||
}
|
||||
|
||||
/// Attention core shared by the single-GPU and TP full-attention
|
||||
/// layers (#95): `(B, H, L, D)` query and `(B, H_kv, S, D)` key/value
|
||||
/// (post-KV-cache, NOT GQA-repeated) → `(B, H, L, D)` context.
|
||||
///
|
||||
/// With the `flash-attn` feature on a CUDA device in f16/bf16, this
|
||||
/// dispatches to the FlashAttention kernel: GQA is native (no
|
||||
/// repeated-K/V materialisation) and causality is a kernel flag, so
|
||||
/// the O(L²) mask/score tensors never exist. The kernels align the
|
||||
/// causal mask to the BOTTOM-RIGHT when `seqlen_q != seqlen_k`
|
||||
/// (flash-attention v2.1+ semantics), which is exactly what chunked
|
||||
/// prefill continuation needs: a chunk of L new queries against
|
||||
/// `offset + L` cached keys masks correctly.
|
||||
///
|
||||
/// INVARIANT: `attn_mask` is either `None` (decode / single position)
|
||||
/// or the standard causal mask — the only mask the qwen3_5 forward
|
||||
/// constructs. The flash path encodes it as `causal = attn_mask
|
||||
/// .is_some()`; a future non-causal mask must extend this signature,
|
||||
/// not silently pass through.
|
||||
///
|
||||
/// Falls back to the eager matmul→softmax→matmul everywhere else
|
||||
/// (CPU, f32, feature off, or `NEURON_FLASH_ATTN=0`).
|
||||
pub(crate) fn attention_context(
|
||||
q: &Tensor,
|
||||
k: &Tensor,
|
||||
v: &Tensor,
|
||||
attn_mask: Option<&Tensor>,
|
||||
num_kv_groups: usize,
|
||||
scale: f64,
|
||||
) -> candle_core::Result<Tensor> {
|
||||
#[cfg(feature = "flash-attn")]
|
||||
{
|
||||
use candle_core::DType;
|
||||
let dtype = q.dtype();
|
||||
// Prefill only (q_len > 1): measured on beast (27B, 30k-token
|
||||
// prompt, 2x RTX 5090), flash cuts prefill 24.8s → 22.1s but
|
||||
// REGRESSES decode ~20% (50 → 60 ms/token at 30k KV) — FA2
|
||||
// without flash-decoding is weak at query-length 1 and the
|
||||
// per-step layout transposes add overhead. Greedy outputs are
|
||||
// byte-identical either way, so this split is purely a
|
||||
// performance routing decision.
|
||||
if flash_attn_enabled()
|
||||
&& q.dim(2)? > 1
|
||||
&& q.device().is_cuda()
|
||||
&& (dtype == DType::F16 || dtype == DType::BF16)
|
||||
{
|
||||
// flash_attn wants (B, L, H, D); the callers carry (B, H, L, D).
|
||||
let qf = q.transpose(1, 2)?.contiguous()?;
|
||||
let kf = k.transpose(1, 2)?.contiguous()?;
|
||||
let vf = v.transpose(1, 2)?.contiguous()?;
|
||||
let causal = attn_mask.is_some();
|
||||
let ctx = candle_flash_attn::flash_attn(&qf, &kf, &vf, scale as f32, causal)?;
|
||||
return ctx.transpose(1, 2)?.contiguous();
|
||||
}
|
||||
}
|
||||
|
||||
// Eager fallback: materialise GQA-repeated K/V and the score matrix.
|
||||
let k = repeat_kv(k.clone(), num_kv_groups)?.contiguous()?;
|
||||
let v = repeat_kv(v.clone(), num_kv_groups)?.contiguous()?;
|
||||
let mut scores = (q.matmul(&k.transpose(2, 3)?)? * scale)?;
|
||||
if let Some(m) = attn_mask {
|
||||
scores = scores.broadcast_add(m)?;
|
||||
}
|
||||
let probs = candle_nn::ops::softmax_last_dim(&scores)?;
|
||||
probs.matmul(&v)
|
||||
}
|
||||
|
||||
pub struct Qwen3_5Attention {
|
||||
q_proj: Linear,
|
||||
k_proj: Linear,
|
||||
@@ -139,18 +219,10 @@ impl Qwen3_5Attention {
|
||||
// 4. KV cache.
|
||||
let (k, v) = self.kv_cache.append(&k, &v)?;
|
||||
|
||||
// 5. GQA repeat (cheap shape op).
|
||||
let k = repeat_kv(k, self.num_kv_groups)?.contiguous()?;
|
||||
let v = repeat_kv(v, self.num_kv_groups)?.contiguous()?;
|
||||
|
||||
// 6. Scaled dot-product + causal mask.
|
||||
// 5+6. Attention core — FlashAttention when available, eager
|
||||
// GQA-repeat + masked softmax otherwise (#95).
|
||||
let scale = 1.0_f64 / (self.head_dim as f64).sqrt();
|
||||
let mut scores = (q.matmul(&k.transpose(2, 3)?)? * scale)?;
|
||||
if let Some(m) = attn_mask {
|
||||
scores = scores.broadcast_add(m)?;
|
||||
}
|
||||
let probs = candle_nn::ops::softmax_last_dim(&scores)?;
|
||||
let ctx = probs.matmul(&v)?; // (B, H, L, D)
|
||||
let ctx = attention_context(&q, &k, &v, attn_mask, self.num_kv_groups, scale)?; // (B, H, L, D)
|
||||
|
||||
// 7. Reshape back, apply the output gate, project.
|
||||
let ctx = ctx
|
||||
|
||||
@@ -139,10 +139,42 @@ impl GatedDeltaNet {
|
||||
let conv_dim = key_dim * 2 + value_dim;
|
||||
|
||||
// ----- Linear projections (all `bias=False` in the reference). -----
|
||||
let in_proj_qkv = load_linear_no_bias(vb, "in_proj_qkv", cfg.hidden_size, conv_dim)?;
|
||||
let in_proj_z = load_linear_no_bias(vb, "in_proj_z", cfg.hidden_size, value_dim)?;
|
||||
let in_proj_b = load_linear_no_bias(vb, "in_proj_b", cfg.hidden_size, num_v_heads)?;
|
||||
let in_proj_a = load_linear_no_bias(vb, "in_proj_a", cfg.hidden_size, num_v_heads)?;
|
||||
// Two checkpoint layouts exist for the input projections:
|
||||
// - Qwen3.6 (qwen3_5): separate `in_proj_qkv` / `in_proj_z` /
|
||||
// `in_proj_b` / `in_proj_a`, with qkv stored as contiguous
|
||||
// [Q | K | V] blocks — loads directly.
|
||||
// - Qwen3-Next 80B-A3B (qwen3_next, #92): fused `in_proj_qkvz`
|
||||
// + `in_proj_ba`, **interleaved per key-head group** (see
|
||||
// `split_fused_qkvz`/`split_fused_ba`) — de-interleaved once
|
||||
// at load into the same contiguous layout, so the forward
|
||||
// path (incl. the conv over [Q|K|V] channels) is unchanged.
|
||||
let (in_proj_qkv, in_proj_z, in_proj_b, in_proj_a) =
|
||||
if vb.contains_tensor("in_proj_qkvz.weight") {
|
||||
let qkvz = vb
|
||||
.pp("in_proj_qkvz")
|
||||
.get((2 * key_dim + 2 * value_dim, cfg.hidden_size), "weight")
|
||||
.with_context(|| format!("load '{}/in_proj_qkvz/weight'", vb.prefix()))?;
|
||||
let ba = vb
|
||||
.pp("in_proj_ba")
|
||||
.get((2 * num_v_heads, cfg.hidden_size), "weight")
|
||||
.with_context(|| format!("load '{}/in_proj_ba/weight'", vb.prefix()))?;
|
||||
let (qkv_w, z_w) =
|
||||
split_fused_qkvz(&qkvz, num_k_heads, num_v_heads, head_k_dim, head_v_dim)?;
|
||||
let (b_w, a_w) = split_fused_ba(&ba, num_k_heads, num_v_heads)?;
|
||||
(
|
||||
Linear::new(qkv_w, None),
|
||||
Linear::new(z_w, None),
|
||||
Linear::new(b_w, None),
|
||||
Linear::new(a_w, None),
|
||||
)
|
||||
} else {
|
||||
(
|
||||
load_linear_no_bias(vb, "in_proj_qkv", cfg.hidden_size, conv_dim)?,
|
||||
load_linear_no_bias(vb, "in_proj_z", cfg.hidden_size, value_dim)?,
|
||||
load_linear_no_bias(vb, "in_proj_b", cfg.hidden_size, num_v_heads)?,
|
||||
load_linear_no_bias(vb, "in_proj_a", cfg.hidden_size, num_v_heads)?,
|
||||
)
|
||||
};
|
||||
let out_proj = load_linear_no_bias(vb, "out_proj", value_dim, cfg.hidden_size)?;
|
||||
|
||||
// ----- Conv1d weight (depthwise, bias=False). -----
|
||||
@@ -889,6 +921,61 @@ fn load_linear_no_bias(
|
||||
Ok(Linear::new(weight, None))
|
||||
}
|
||||
|
||||
/// De-interleave a fused `in_proj_qkvz.weight` (qwen3_next layout, #92)
|
||||
/// into a contiguous `[Q | K | V]` qkv weight plus a `Z` weight.
|
||||
///
|
||||
/// The fused rows are grouped **per key head**: for each of the
|
||||
/// `num_k_heads` groups (`r = num_v_heads / num_k_heads`, group stride
|
||||
/// `s = 2*head_k + 2*head_v*r`), the group holds
|
||||
/// `[q (head_k) | k (head_k) | v (head_v*r) | z (head_v*r)]` — the
|
||||
/// reshape in upstream `fix_query_key_value_ordering`
|
||||
/// `(num_k_heads, 2*head_k + 2*head_v*num_v/num_k)`. Concatenating the
|
||||
/// per-group regions restores the global-contiguous layout the rest of
|
||||
/// this module (incl. the conv over `[Q|K|V]` channels) expects.
|
||||
pub(crate) fn split_fused_qkvz(
|
||||
qkvz: &Tensor,
|
||||
num_k_heads: usize,
|
||||
num_v_heads: usize,
|
||||
head_k_dim: usize,
|
||||
head_v_dim: usize,
|
||||
) -> Result<(Tensor, Tensor)> {
|
||||
let r = num_v_heads / num_k_heads;
|
||||
let stride = 2 * head_k_dim + 2 * head_v_dim * r;
|
||||
let (mut qs, mut ks, mut vs, mut zs) = (Vec::new(), Vec::new(), Vec::new(), Vec::new());
|
||||
for g in 0..num_k_heads {
|
||||
let base = g * stride;
|
||||
qs.push(qkvz.narrow(0, base, head_k_dim)?);
|
||||
ks.push(qkvz.narrow(0, base + head_k_dim, head_k_dim)?);
|
||||
vs.push(qkvz.narrow(0, base + 2 * head_k_dim, head_v_dim * r)?);
|
||||
zs.push(qkvz.narrow(0, base + 2 * head_k_dim + head_v_dim * r, head_v_dim * r)?);
|
||||
}
|
||||
let parts: Vec<Tensor> = qs.into_iter().chain(ks).chain(vs).collect();
|
||||
let qkv = Tensor::cat(&parts, 0)?.contiguous()?;
|
||||
let z = Tensor::cat(&zs, 0)?.contiguous()?;
|
||||
Ok((qkv, z))
|
||||
}
|
||||
|
||||
/// De-interleave a fused `in_proj_ba.weight` (qwen3_next layout, #92)
|
||||
/// into per-v-head `b` (beta) and `a` (decay) weights. Same per-key-head
|
||||
/// grouping as [`split_fused_qkvz`]: each group holds `[b (r) | a (r)]`
|
||||
/// rows, `r = num_v_heads / num_k_heads`.
|
||||
pub(crate) fn split_fused_ba(
|
||||
ba: &Tensor,
|
||||
num_k_heads: usize,
|
||||
num_v_heads: usize,
|
||||
) -> Result<(Tensor, Tensor)> {
|
||||
let r = num_v_heads / num_k_heads;
|
||||
let (mut bs, mut r#as) = (Vec::new(), Vec::new());
|
||||
for g in 0..num_k_heads {
|
||||
let base = g * 2 * r;
|
||||
bs.push(ba.narrow(0, base, r)?);
|
||||
r#as.push(ba.narrow(0, base + r, r)?);
|
||||
}
|
||||
let b = Tensor::cat(&bs, 0)?.contiguous()?;
|
||||
let a = Tensor::cat(&r#as, 0)?.contiguous()?;
|
||||
Ok((b, a))
|
||||
}
|
||||
|
||||
/// Numerically-stable `softplus(x) = ln(1 + exp(x))`. Matches PyTorch's
|
||||
/// `F.softplus` default (beta=1, threshold=20: for large positive x,
|
||||
/// returns x as-is to avoid overflow in the exp).
|
||||
@@ -1138,6 +1225,13 @@ mod tests {
|
||||
linear_key_head_dim: 4,
|
||||
linear_value_head_dim: 4,
|
||||
linear_conv_kernel_dim: 4,
|
||||
num_experts: 0,
|
||||
num_experts_per_tok: 0,
|
||||
moe_intermediate_size: 0,
|
||||
shared_expert_intermediate_size: 0,
|
||||
decoder_sparse_step: 1,
|
||||
mlp_only_layers: Vec::new(),
|
||||
norm_topk_prob: false,
|
||||
};
|
||||
|
||||
// Build a synthetic VarBuilder with all-zeros weights.
|
||||
@@ -1179,4 +1273,115 @@ mod tests {
|
||||
let v: Vec<f32> = y.flatten_all().unwrap().to_vec1().unwrap();
|
||||
assert!(v.iter().all(|x| x.is_finite()));
|
||||
}
|
||||
|
||||
/// Interleave known per-head Q/K/V/Z (and B/A) rows into the fused
|
||||
/// qwen3_next layout, split, and expect the original contiguous
|
||||
/// blocks back. Layout under test: per key-head group g,
|
||||
/// `[q_g | k_g | v_g | z_g]` with r = num_v/num_k value heads per
|
||||
/// group (upstream `fix_query_key_value_ordering`).
|
||||
#[test]
|
||||
fn split_fused_qkvz_and_ba_roundtrip() {
|
||||
let dev = Device::Cpu;
|
||||
let (num_k, num_v, head_k, head_v, hidden) = (2usize, 4usize, 3usize, 2usize, 5usize);
|
||||
let r = num_v / num_k;
|
||||
|
||||
// Distinct constant per logical row so any mis-slicing shows.
|
||||
let row = |tag: f32| Tensor::full(tag, (1, hidden), &dev).unwrap();
|
||||
let mut fused_rows: Vec<Tensor> = Vec::new();
|
||||
let (mut q_rows, mut k_rows, mut v_rows, mut z_rows) =
|
||||
(Vec::new(), Vec::new(), Vec::new(), Vec::new());
|
||||
for g in 0..num_k {
|
||||
let base = 1000.0 * (g as f32 + 1.0);
|
||||
for i in 0..head_k {
|
||||
let t = row(base + i as f32);
|
||||
fused_rows.push(t.clone());
|
||||
q_rows.push(t);
|
||||
}
|
||||
for i in 0..head_k {
|
||||
let t = row(base + 100.0 + i as f32);
|
||||
fused_rows.push(t.clone());
|
||||
k_rows.push(t);
|
||||
}
|
||||
for i in 0..head_v * r {
|
||||
let t = row(base + 200.0 + i as f32);
|
||||
fused_rows.push(t.clone());
|
||||
v_rows.push(t);
|
||||
}
|
||||
for i in 0..head_v * r {
|
||||
let t = row(base + 300.0 + i as f32);
|
||||
fused_rows.push(t.clone());
|
||||
z_rows.push(t);
|
||||
}
|
||||
}
|
||||
let fused = Tensor::cat(&fused_rows, 0).unwrap();
|
||||
let expected_qkv = Tensor::cat(
|
||||
&q_rows
|
||||
.iter()
|
||||
.chain(k_rows.iter())
|
||||
.chain(v_rows.iter())
|
||||
.cloned()
|
||||
.collect::<Vec<_>>(),
|
||||
0,
|
||||
)
|
||||
.unwrap();
|
||||
let expected_z = Tensor::cat(&z_rows, 0).unwrap();
|
||||
|
||||
let (qkv, z) = split_fused_qkvz(&fused, num_k, num_v, head_k, head_v).unwrap();
|
||||
assert_eq!(qkv.dims(), &[2 * num_k * head_k + num_v * head_v, hidden]);
|
||||
let diff_qkv: f32 = (qkv - expected_qkv)
|
||||
.unwrap()
|
||||
.abs()
|
||||
.unwrap()
|
||||
.max_all()
|
||||
.unwrap()
|
||||
.to_scalar()
|
||||
.unwrap();
|
||||
let diff_z: f32 = (z - expected_z)
|
||||
.unwrap()
|
||||
.abs()
|
||||
.unwrap()
|
||||
.max_all()
|
||||
.unwrap()
|
||||
.to_scalar()
|
||||
.unwrap();
|
||||
assert_eq!(diff_qkv, 0.0);
|
||||
assert_eq!(diff_z, 0.0);
|
||||
|
||||
// ba: per group, [b (r rows) | a (r rows)].
|
||||
let mut ba_rows = Vec::new();
|
||||
let (mut b_rows, mut a_rows) = (Vec::new(), Vec::new());
|
||||
for g in 0..num_k {
|
||||
let base = 10.0 * (g as f32 + 1.0);
|
||||
for i in 0..r {
|
||||
let t = row(base + i as f32);
|
||||
ba_rows.push(t.clone());
|
||||
b_rows.push(t);
|
||||
}
|
||||
for i in 0..r {
|
||||
let t = row(base + 5.0 + i as f32);
|
||||
ba_rows.push(t.clone());
|
||||
a_rows.push(t);
|
||||
}
|
||||
}
|
||||
let ba = Tensor::cat(&ba_rows, 0).unwrap();
|
||||
let (b, a) = split_fused_ba(&ba, num_k, num_v).unwrap();
|
||||
let diff_b: f32 = (b - Tensor::cat(&b_rows, 0).unwrap())
|
||||
.unwrap()
|
||||
.abs()
|
||||
.unwrap()
|
||||
.max_all()
|
||||
.unwrap()
|
||||
.to_scalar()
|
||||
.unwrap();
|
||||
let diff_a: f32 = (a - Tensor::cat(&a_rows, 0).unwrap())
|
||||
.unwrap()
|
||||
.abs()
|
||||
.unwrap()
|
||||
.max_all()
|
||||
.unwrap()
|
||||
.to_scalar()
|
||||
.unwrap();
|
||||
assert_eq!(diff_b, 0.0);
|
||||
assert_eq!(diff_a, 0.0);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -17,12 +17,32 @@ pub struct Qwen3_5MLP {
|
||||
}
|
||||
|
||||
impl Qwen3_5MLP {
|
||||
/// Construct directly from pre-built projections (MoE-block tests).
|
||||
#[cfg(test)]
|
||||
pub(crate) fn from_weights(gate_proj: Linear, up_proj: Linear, down_proj: Linear) -> Self {
|
||||
Self {
|
||||
gate_proj,
|
||||
up_proj,
|
||||
down_proj,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn load(cfg: &TextConfig, vb: &ShardedVarBuilder) -> Result<Self> {
|
||||
let h = cfg.hidden_size;
|
||||
let i = cfg.intermediate_size;
|
||||
let gate_proj = load_linear_no_bias(vb, "gate_proj", h, i)?;
|
||||
let up_proj = load_linear_no_bias(vb, "up_proj", h, i)?;
|
||||
let down_proj = load_linear_no_bias(vb, "down_proj", i, h)?;
|
||||
Self::load_with_dims(vb, cfg.hidden_size, cfg.intermediate_size)
|
||||
}
|
||||
|
||||
/// Load with explicit dims — the MoE block (#92) reuses this SwiGLU
|
||||
/// shape for routed experts (`moe_intermediate_size`) and the shared
|
||||
/// expert (`shared_expert_intermediate_size`), both narrower than
|
||||
/// the dense `intermediate_size`.
|
||||
pub fn load_with_dims(
|
||||
vb: &ShardedVarBuilder,
|
||||
hidden: usize,
|
||||
intermediate: usize,
|
||||
) -> Result<Self> {
|
||||
let gate_proj = load_linear_no_bias(vb, "gate_proj", hidden, intermediate)?;
|
||||
let up_proj = load_linear_no_bias(vb, "up_proj", hidden, intermediate)?;
|
||||
let down_proj = load_linear_no_bias(vb, "down_proj", intermediate, hidden)?;
|
||||
Ok(Self {
|
||||
gate_proj,
|
||||
up_proj,
|
||||
|
||||
@@ -76,6 +76,7 @@ pub mod decoder;
|
||||
pub mod full_attn;
|
||||
pub mod linear_attn;
|
||||
pub mod mlp;
|
||||
pub mod moe;
|
||||
pub mod rmsnorm;
|
||||
pub mod rope;
|
||||
pub mod snapshot;
|
||||
@@ -90,6 +91,13 @@ use rope::RotaryEmbedding;
|
||||
/// magic strings.
|
||||
pub const MODEL_TYPE: &str = "qwen3_5";
|
||||
|
||||
/// `model_type` of the MoE sibling family (Qwen3-Next-80B-A3B /
|
||||
/// Qwen3-Coder-Next): the same Gated DeltaNet hybrid layer stack with
|
||||
/// a high-sparsity MoE FFN per layer. Served by this same arch module —
|
||||
/// [`Config::from_config_json`] normalises the flat qwen3_next
|
||||
/// `config.json` layout into the nested shape used here.
|
||||
pub const MODEL_TYPE_NEXT: &str = "qwen3_next";
|
||||
|
||||
/// Top-level shape of Qwen3-Next's `config.json`. The real
|
||||
/// hyperparameters live in `text_config`; the rest is multimodal /
|
||||
/// tokeniser glue we don't need for the language-model forward.
|
||||
@@ -117,6 +125,79 @@ pub struct Config {
|
||||
pub image_token_id: Option<u32>,
|
||||
}
|
||||
|
||||
impl Config {
|
||||
/// Parse a `config.json` for either family this arch serves,
|
||||
/// normalising layout differences (#92):
|
||||
///
|
||||
/// - `model_type == "qwen3_5"` (Qwen3.6): hyperparameters nested
|
||||
/// under `text_config`, RoPE nested under `rope_parameters` —
|
||||
/// deserialises directly.
|
||||
/// - `model_type == "qwen3_next"` (Qwen3-Next-80B-A3B family):
|
||||
/// **flat** layout — hyperparameters at the top level,
|
||||
/// `rope_theta`/`partial_rotary_factor` flat, no vision block.
|
||||
/// Wrapped into the nested shape here. The output gate on full
|
||||
/// attention is unconditional in the upstream qwen3_next
|
||||
/// implementation (the config carries no flag), so
|
||||
/// `attn_output_gate` is forced on.
|
||||
///
|
||||
/// Both variants may omit `layer_types` (qwen3_next always does);
|
||||
/// it is derived from `full_attention_interval` using the upstream
|
||||
/// convention: layer `i` is `full_attention` iff
|
||||
/// `(i + 1) % interval == 0`, else `linear_attention`.
|
||||
pub fn from_config_json(json: &str) -> Result<Self> {
|
||||
let v: serde_json::Value =
|
||||
serde_json::from_str(json).context("parse config.json as JSON")?;
|
||||
let model_type = v
|
||||
.get("model_type")
|
||||
.and_then(|m| m.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
|
||||
let mut cfg: Config = if model_type == MODEL_TYPE_NEXT {
|
||||
let mut text = v.clone();
|
||||
if text.get("rope_parameters").is_none() {
|
||||
let mut rope = serde_json::Map::new();
|
||||
for key in ["rope_theta", "partial_rotary_factor", "rope_type"] {
|
||||
if let Some(val) = v.get(key) {
|
||||
rope.insert(key.to_string(), val.clone());
|
||||
}
|
||||
}
|
||||
text["rope_parameters"] = serde_json::Value::Object(rope);
|
||||
}
|
||||
let mut text_config: TextConfig = serde_json::from_value(text)
|
||||
.context("parse flat qwen3_next config.json hyperparameters")?;
|
||||
text_config.attn_output_gate = true;
|
||||
Config {
|
||||
model_type,
|
||||
text_config,
|
||||
vision_config: None,
|
||||
image_token_id: None,
|
||||
}
|
||||
} else {
|
||||
serde_json::from_str(json).context("parse nested qwen3_5 config.json")?
|
||||
};
|
||||
|
||||
if cfg.text_config.layer_types.is_empty() {
|
||||
let interval = cfg.text_config.full_attention_interval.unwrap_or(4);
|
||||
anyhow::ensure!(
|
||||
interval > 0,
|
||||
"full_attention_interval must be >= 1 to derive layer_types"
|
||||
);
|
||||
cfg.text_config.layer_types = (0..cfg.text_config.num_hidden_layers)
|
||||
.map(|i| {
|
||||
if (i + 1).is_multiple_of(interval) {
|
||||
"full_attention".to_string()
|
||||
} else {
|
||||
"linear_attention".to_string()
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
}
|
||||
|
||||
Ok(cfg)
|
||||
}
|
||||
}
|
||||
|
||||
/// Inner config (the `text_config` block). Mirrors the Qwen3 layout
|
||||
/// but with the extras Qwen3-Next adds (`attn_output_gate`,
|
||||
/// `layer_types`, `full_attention_interval`, larger `head_dim`).
|
||||
@@ -185,6 +266,56 @@ pub struct TextConfig {
|
||||
/// (Qwen3.6-27B: 4).
|
||||
#[serde(default)]
|
||||
pub linear_conv_kernel_dim: usize,
|
||||
|
||||
// --- High-sparsity MoE FFN (Qwen3-Next 80B-A3B family, #92) --------
|
||||
// All default to the dense case (0 experts) so existing dense
|
||||
// configs (Qwen3.6-27B) deserialise unchanged. A layer gets the MoE
|
||||
// FFN iff `layer_uses_moe` says so; otherwise the dense SwiGLU.
|
||||
/// Total routed experts per MoE layer (80B-A3B: 512). `0` → dense
|
||||
/// model, no MoE anywhere.
|
||||
#[serde(default)]
|
||||
pub num_experts: usize,
|
||||
/// Experts activated per token (80B-A3B: 10).
|
||||
#[serde(default)]
|
||||
pub num_experts_per_tok: usize,
|
||||
/// Per-expert FFN width (80B-A3B: 512). Distinct from the dense
|
||||
/// `intermediate_size`.
|
||||
#[serde(default)]
|
||||
pub moe_intermediate_size: usize,
|
||||
/// Width of the always-on shared expert (80B-A3B: 512). `0` → no
|
||||
/// shared expert (Qwen3-30B-A3B style).
|
||||
#[serde(default)]
|
||||
pub shared_expert_intermediate_size: usize,
|
||||
/// Every `decoder_sparse_step`-th layer is MoE (1 → all layers,
|
||||
/// the 80B-A3B case). Follows the upstream `(i+1) % step == 0`
|
||||
/// convention.
|
||||
#[serde(default = "default_decoder_sparse_step")]
|
||||
pub decoder_sparse_step: usize,
|
||||
/// Layer indices forced to the dense MLP even when MoE is on.
|
||||
/// Empty for 80B-A3B.
|
||||
#[serde(default)]
|
||||
pub mlp_only_layers: Vec<usize>,
|
||||
/// Renormalise the top-k routing weights to sum to 1 (80B-A3B:
|
||||
/// true). Upstream selects top-k *after* softmax over all experts.
|
||||
#[serde(default)]
|
||||
pub norm_topk_prob: bool,
|
||||
}
|
||||
|
||||
impl TextConfig {
|
||||
/// Whether decoder layer `layer_idx` carries the MoE FFN (vs the
|
||||
/// dense SwiGLU). Mirrors upstream `Qwen3NextDecoderLayer`:
|
||||
/// experts configured, layer not in `mlp_only_layers`, and on the
|
||||
/// `decoder_sparse_step` grid.
|
||||
pub fn layer_uses_moe(&self, layer_idx: usize) -> bool {
|
||||
self.num_experts > 0
|
||||
&& self.decoder_sparse_step > 0
|
||||
&& !self.mlp_only_layers.contains(&layer_idx)
|
||||
&& (layer_idx + 1).is_multiple_of(self.decoder_sparse_step)
|
||||
}
|
||||
}
|
||||
|
||||
fn default_decoder_sparse_step() -> usize {
|
||||
1
|
||||
}
|
||||
|
||||
fn default_hidden_act() -> String {
|
||||
@@ -330,17 +461,20 @@ pub struct Qwen3_5Model {
|
||||
}
|
||||
|
||||
impl Qwen3_5Model {
|
||||
pub fn load(cfg: &TextConfig, vb: &ShardedVarBuilder) -> Result<Self> {
|
||||
/// `text_prefix` is where the text core lives in the checkpoint:
|
||||
/// - Qwen3.6 (multimodal, `model_type = "qwen3_5"`):
|
||||
/// `model.language_model` — sibling to `model.visual.*` (the
|
||||
/// vision tower) and top-level `lm_head` / `mtp.*`.
|
||||
/// - Qwen3-Next-80B-A3B (text-only, `model_type = "qwen3_next"`):
|
||||
/// plain `model`.
|
||||
///
|
||||
/// [`Qwen3_5ForCausalLM::new`] picks by `Config::model_type` via
|
||||
/// [`text_weight_prefix`].
|
||||
pub fn load(cfg: &TextConfig, vb: &ShardedVarBuilder, text_prefix: &str) -> Result<Self> {
|
||||
let dtype = vb.dtype();
|
||||
let device = vb.device().clone();
|
||||
|
||||
// Qwen3-Next is a multimodal architecture whose text core lives
|
||||
// under `model.language_model.*` — sibling to `model.visual.*`
|
||||
// (the vision tower) and to top-level `lm_head` / `mtp.*`.
|
||||
// Every text-side tensor in the safetensors files is under
|
||||
// this prefix; we ignore the vision and MTP weights for
|
||||
// language-model inference.
|
||||
let text_vb = vb.pp("model.language_model");
|
||||
let text_vb = vb.pp(text_prefix);
|
||||
|
||||
let embed_vb = text_vb.pp("embed_tokens");
|
||||
let embed_weight = embed_vb
|
||||
@@ -444,6 +578,67 @@ impl Qwen3_5Model {
|
||||
self.forward_inner(input, offset, None, None, &[], None)
|
||||
}
|
||||
|
||||
/// Lockstep batched decode step (#98): `input` is `(B, 1)` — one
|
||||
/// new token per batch row — with each row at its own sequence
|
||||
/// position `positions[i]` (typically `prefix_lens[i] + step`).
|
||||
/// The cache must hold batched state (see
|
||||
/// `snapshot::assemble_batch`); `attn_mask` is the padding mask
|
||||
/// from [`Self::batch_decode_mask`] (or `None` when no row is
|
||||
/// padded). Text-only: `rope_delta` is ignored — positions are
|
||||
/// explicit and vision requests never enter the batch path.
|
||||
pub fn forward_batch_decode(
|
||||
&mut self,
|
||||
input: &Tensor,
|
||||
positions: &[usize],
|
||||
attn_mask: Option<&Tensor>,
|
||||
) -> candle_core::Result<Tensor> {
|
||||
let (b, l) = input.dims2()?;
|
||||
if l != 1 {
|
||||
candle_core::bail!("forward_batch_decode: expected (B, 1) input, got (B, {l})");
|
||||
}
|
||||
if positions.len() != b {
|
||||
candle_core::bail!(
|
||||
"forward_batch_decode: {} positions for batch of {b}",
|
||||
positions.len()
|
||||
);
|
||||
}
|
||||
let mut h = self.embed_tokens.forward(input)?;
|
||||
let (cos, sin) = self.rotary.batch_cos_sin(positions)?;
|
||||
for layer in &mut self.layers {
|
||||
h = layer.forward(&h, attn_mask, &cos, &sin)?;
|
||||
}
|
||||
self.norm.forward(&h)
|
||||
}
|
||||
|
||||
/// Additive padding mask for a batched decode step: shape
|
||||
/// `(B, 1, 1, total_len)`, `-inf` on each row's padding gap
|
||||
/// `[prefix_lens[i], padded_len)`, zero elsewhere. `total_len` is
|
||||
/// the KV length *after* this step's append (`padded_len + step +
|
||||
/// 1`). Returns `None` when no row is padded (uniform prefix
|
||||
/// lengths) — the decode step then needs no mask at all, matching
|
||||
/// the single-sequence fast path.
|
||||
pub fn batch_decode_mask(
|
||||
&self,
|
||||
prefix_lens: &[usize],
|
||||
padded_len: usize,
|
||||
total_len: usize,
|
||||
) -> candle_core::Result<Option<Tensor>> {
|
||||
if prefix_lens.iter().all(|&len| len == padded_len) {
|
||||
return Ok(None);
|
||||
}
|
||||
let minf = f32::NEG_INFINITY;
|
||||
let b = prefix_lens.len();
|
||||
let mask: Vec<f32> = prefix_lens
|
||||
.iter()
|
||||
.flat_map(|&len| {
|
||||
(0..total_len).map(move |j| if j >= len && j < padded_len { minf } else { 0. })
|
||||
})
|
||||
.collect();
|
||||
Ok(Some(
|
||||
Tensor::from_vec(mask, (b, 1, 1, total_len), &self.device)?.to_dtype(self.dtype)?,
|
||||
))
|
||||
}
|
||||
|
||||
/// Forward for a vision-prefill chunk: optional image-embedding
|
||||
/// splice plus explicit interleaved-M-RoPE `position_ids` (the
|
||||
/// chunk's slice of the full prompt's 3D positions). Mirrors the TP
|
||||
@@ -605,10 +800,20 @@ pub struct Qwen3_5ForCausalLM {
|
||||
image_token_id: Option<u32>,
|
||||
}
|
||||
|
||||
/// Checkpoint prefix of the text core for a given `model_type` — see
|
||||
/// [`Qwen3_5Model::load`].
|
||||
pub fn text_weight_prefix(model_type: &str) -> &'static str {
|
||||
if model_type == MODEL_TYPE_NEXT {
|
||||
"model"
|
||||
} else {
|
||||
"model.language_model"
|
||||
}
|
||||
}
|
||||
|
||||
impl Qwen3_5ForCausalLM {
|
||||
pub fn new(config: Config, vb: ShardedVarBuilder) -> Result<Self> {
|
||||
let cfg = &config.text_config;
|
||||
let base = Qwen3_5Model::load(cfg, &vb)?;
|
||||
let base = Qwen3_5Model::load(cfg, &vb, text_weight_prefix(&config.model_type))?;
|
||||
let lm_head = if cfg.tie_word_embeddings {
|
||||
Linear::new(base.embed_weight().clone(), None)
|
||||
} else {
|
||||
@@ -676,6 +881,34 @@ impl Qwen3_5ForCausalLM {
|
||||
hidden.i((.., l - 1.., ..))?.apply(&self.lm_head)
|
||||
}
|
||||
|
||||
/// Lockstep batched decode step (#98): `(B, 1)` input, per-row
|
||||
/// positions, padding mask from
|
||||
/// [`Qwen3_5Model::batch_decode_mask`]. Returns `(B, 1,
|
||||
/// vocab_size)` — one logits row per batch row.
|
||||
pub fn forward_batch_decode(
|
||||
&mut self,
|
||||
input: &Tensor,
|
||||
positions: &[usize],
|
||||
attn_mask: Option<&Tensor>,
|
||||
) -> candle_core::Result<Tensor> {
|
||||
let hidden = self
|
||||
.base
|
||||
.forward_batch_decode(input, positions, attn_mask)?;
|
||||
hidden.apply(&self.lm_head)
|
||||
}
|
||||
|
||||
/// Padding mask for a batched decode step — see
|
||||
/// [`Qwen3_5Model::batch_decode_mask`].
|
||||
pub fn batch_decode_mask(
|
||||
&self,
|
||||
prefix_lens: &[usize],
|
||||
padded_len: usize,
|
||||
total_len: usize,
|
||||
) -> candle_core::Result<Option<Tensor>> {
|
||||
self.base
|
||||
.batch_decode_mask(prefix_lens, padded_len, total_len)
|
||||
}
|
||||
|
||||
/// Stage B: forward with image-embedding splice. Mirrors `forward`
|
||||
/// but routes through `Qwen3_5Model::forward_with_vision` so the
|
||||
/// LM's input embeddings get the image patches spliced in at
|
||||
@@ -915,6 +1148,233 @@ mod tests {
|
||||
assert_eq!(cfg.text_config.layer_types.len(), 4);
|
||||
assert_eq!(cfg.text_config.rope_parameters.rope_theta, 10_000_000.0);
|
||||
assert!((cfg.text_config.rope_parameters.partial_rotary_factor - 0.25).abs() < 1e-6);
|
||||
// Dense config: no MoE anywhere.
|
||||
assert_eq!(cfg.text_config.num_experts, 0);
|
||||
assert!(!cfg.text_config.layer_uses_moe(0));
|
||||
|
||||
// The normalising entry point must agree with plain serde for
|
||||
// the nested shape (and leave the explicit layer_types alone).
|
||||
let via_norm = Config::from_config_json(raw).expect("normalised parse");
|
||||
assert_eq!(via_norm.text_config.layer_types.len(), 4);
|
||||
assert_eq!(via_norm.text_config.layer_types[3], "full_attention");
|
||||
}
|
||||
|
||||
/// The flat qwen3_next layout (Qwen3-Next-80B-A3B family): all
|
||||
/// hyperparameters top-level, flat rope fields, no `layer_types`,
|
||||
/// MoE fields present. Sample mirrors
|
||||
/// `Qwen/Qwen3-Next-80B-A3B-Instruct/config.json`.
|
||||
#[test]
|
||||
fn config_normalises_the_flat_qwen3_next_shape() {
|
||||
let raw = r#"{
|
||||
"architectures": ["Qwen3NextForCausalLM"],
|
||||
"model_type": "qwen3_next",
|
||||
"vocab_size": 151936,
|
||||
"hidden_size": 2048,
|
||||
"intermediate_size": 5120,
|
||||
"num_hidden_layers": 48,
|
||||
"num_attention_heads": 16,
|
||||
"num_key_value_heads": 2,
|
||||
"head_dim": 256,
|
||||
"max_position_embeddings": 262144,
|
||||
"partial_rotary_factor": 0.25,
|
||||
"rope_theta": 10000000,
|
||||
"rms_norm_eps": 1e-6,
|
||||
"tie_word_embeddings": false,
|
||||
"full_attention_interval": 4,
|
||||
"linear_conv_kernel_dim": 4,
|
||||
"linear_key_head_dim": 128,
|
||||
"linear_num_key_heads": 16,
|
||||
"linear_num_value_heads": 32,
|
||||
"linear_value_head_dim": 128,
|
||||
"decoder_sparse_step": 1,
|
||||
"mlp_only_layers": [],
|
||||
"moe_intermediate_size": 512,
|
||||
"norm_topk_prob": true,
|
||||
"num_experts": 512,
|
||||
"num_experts_per_tok": 10,
|
||||
"shared_expert_intermediate_size": 512
|
||||
}"#;
|
||||
let cfg = Config::from_config_json(raw).expect("parse qwen3_next config");
|
||||
assert_eq!(cfg.model_type, MODEL_TYPE_NEXT);
|
||||
assert!(cfg.vision_config.is_none());
|
||||
|
||||
let t = &cfg.text_config;
|
||||
assert_eq!(t.hidden_size, 2048);
|
||||
// Flat rope fields normalised into the nested block.
|
||||
assert_eq!(t.rope_parameters.rope_theta, 10_000_000.0);
|
||||
assert!((t.rope_parameters.partial_rotary_factor - 0.25).abs() < 1e-6);
|
||||
// Output-gated attention is unconditional for qwen3_next.
|
||||
assert!(t.attn_output_gate);
|
||||
// layer_types derived from the interval: (i+1) % 4 == 0 → full.
|
||||
assert_eq!(t.layer_types.len(), 48);
|
||||
assert_eq!(t.layer_types[3], "full_attention");
|
||||
assert_eq!(t.layer_types[47], "full_attention");
|
||||
assert_eq!(t.layer_types[0], "linear_attention");
|
||||
assert_eq!(t.layer_types[46], "linear_attention");
|
||||
assert_eq!(
|
||||
t.layer_types
|
||||
.iter()
|
||||
.filter(|s| *s == "full_attention")
|
||||
.count(),
|
||||
12
|
||||
);
|
||||
// MoE hyperparameters land.
|
||||
assert_eq!(t.num_experts, 512);
|
||||
assert_eq!(t.num_experts_per_tok, 10);
|
||||
assert_eq!(t.moe_intermediate_size, 512);
|
||||
assert_eq!(t.shared_expert_intermediate_size, 512);
|
||||
assert!(t.norm_topk_prob);
|
||||
// decoder_sparse_step 1 + empty mlp_only_layers → every layer MoE.
|
||||
assert!(t.layer_uses_moe(0));
|
||||
assert!(t.layer_uses_moe(47));
|
||||
}
|
||||
|
||||
/// End-to-end structural check for the qwen3_next path (#92): a
|
||||
/// tiny random-weight checkpoint in the **flat** layout (`model.*`
|
||||
/// prefix, fused `in_proj_qkvz`/`in_proj_ba`, per-expert MoE
|
||||
/// tensors, shared expert) loads through `Config::from_config_json`
|
||||
/// and `Qwen3_5ForCausalLM::new`, producing finite logits of the
|
||||
/// right shape. Numerical parity vs HF is pinned separately by the
|
||||
/// `qwen3_next_parity` fixture test.
|
||||
#[test]
|
||||
fn tiny_qwen3_next_checkpoint_loads_and_forwards() {
|
||||
use candle_core::Device;
|
||||
use std::collections::HashMap;
|
||||
|
||||
let raw = r#"{
|
||||
"model_type": "qwen3_next",
|
||||
"vocab_size": 32, "hidden_size": 8, "intermediate_size": 16,
|
||||
"num_hidden_layers": 2, "num_attention_heads": 2,
|
||||
"num_key_value_heads": 1, "head_dim": 4,
|
||||
"max_position_embeddings": 64, "rms_norm_eps": 1e-6,
|
||||
"full_attention_interval": 2,
|
||||
"linear_num_value_heads": 4, "linear_num_key_heads": 2,
|
||||
"linear_key_head_dim": 4, "linear_value_head_dim": 4,
|
||||
"linear_conv_kernel_dim": 4,
|
||||
"num_experts": 4, "num_experts_per_tok": 2,
|
||||
"moe_intermediate_size": 4,
|
||||
"shared_expert_intermediate_size": 4,
|
||||
"norm_topk_prob": true
|
||||
}"#;
|
||||
let cfg = Config::from_config_json(raw).expect("parse tiny qwen3_next config");
|
||||
assert_eq!(cfg.text_config.layer_types[0], "linear_attention");
|
||||
assert_eq!(cfg.text_config.layer_types[1], "full_attention");
|
||||
|
||||
let dev = Device::Cpu;
|
||||
let randn = |shape: &[usize]| Tensor::randn(0f32, 0.1f32, shape, &dev).unwrap();
|
||||
let ones = |shape: &[usize]| Tensor::ones(shape, DType::F32, &dev).unwrap();
|
||||
let mut t: HashMap<String, Tensor> = HashMap::new();
|
||||
|
||||
let (h, vocab) = (8usize, 32usize);
|
||||
t.insert("model.embed_tokens.weight".into(), randn(&[vocab, h]));
|
||||
t.insert("lm_head.weight".into(), randn(&[vocab, h]));
|
||||
t.insert("model.norm.weight".into(), ones(&[h]));
|
||||
|
||||
let moe = |t: &mut HashMap<String, Tensor>, p: &str| {
|
||||
t.insert(format!("{p}.gate.weight"), randn(&[4, h]));
|
||||
for e in 0..4 {
|
||||
t.insert(format!("{p}.experts.{e}.gate_proj.weight"), randn(&[4, h]));
|
||||
t.insert(format!("{p}.experts.{e}.up_proj.weight"), randn(&[4, h]));
|
||||
t.insert(format!("{p}.experts.{e}.down_proj.weight"), randn(&[h, 4]));
|
||||
}
|
||||
t.insert(
|
||||
format!("{p}.shared_expert.gate_proj.weight"),
|
||||
randn(&[4, h]),
|
||||
);
|
||||
t.insert(format!("{p}.shared_expert.up_proj.weight"), randn(&[4, h]));
|
||||
t.insert(
|
||||
format!("{p}.shared_expert.down_proj.weight"),
|
||||
randn(&[h, 4]),
|
||||
);
|
||||
t.insert(format!("{p}.shared_expert_gate.weight"), randn(&[1, h]));
|
||||
};
|
||||
|
||||
// Layer 0: linear_attention with the FUSED qwen3_next input
|
||||
// projections. key_dim = 2*4 = 8, value_dim = 4*4 = 16 →
|
||||
// qkvz rows = 2*8 + 2*16 = 48, ba rows = 2*4 = 8, conv_dim = 32.
|
||||
let l0 = "model.layers.0";
|
||||
t.insert(
|
||||
format!("{l0}.linear_attn.in_proj_qkvz.weight"),
|
||||
randn(&[48, h]),
|
||||
);
|
||||
t.insert(
|
||||
format!("{l0}.linear_attn.in_proj_ba.weight"),
|
||||
randn(&[8, h]),
|
||||
);
|
||||
t.insert(
|
||||
format!("{l0}.linear_attn.conv1d.weight"),
|
||||
randn(&[32, 1, 4]),
|
||||
);
|
||||
t.insert(format!("{l0}.linear_attn.dt_bias"), randn(&[4]));
|
||||
t.insert(format!("{l0}.linear_attn.A_log"), randn(&[4]));
|
||||
t.insert(format!("{l0}.linear_attn.norm.weight"), ones(&[4]));
|
||||
t.insert(format!("{l0}.linear_attn.out_proj.weight"), randn(&[h, 16]));
|
||||
t.insert(format!("{l0}.input_layernorm.weight"), ones(&[h]));
|
||||
t.insert(format!("{l0}.post_attention_layernorm.weight"), ones(&[h]));
|
||||
moe(&mut t, &format!("{l0}.mlp"));
|
||||
|
||||
// Layer 1: full_attention (output-gated: q_proj is 2×).
|
||||
let l1 = "model.layers.1";
|
||||
t.insert(
|
||||
format!("{l1}.self_attn.q_proj.weight"),
|
||||
randn(&[2 * 2 * 4, h]),
|
||||
);
|
||||
t.insert(format!("{l1}.self_attn.k_proj.weight"), randn(&[4, h]));
|
||||
t.insert(format!("{l1}.self_attn.v_proj.weight"), randn(&[4, h]));
|
||||
t.insert(format!("{l1}.self_attn.o_proj.weight"), randn(&[h, 8]));
|
||||
t.insert(format!("{l1}.self_attn.q_norm.weight"), ones(&[4]));
|
||||
t.insert(format!("{l1}.self_attn.k_norm.weight"), ones(&[4]));
|
||||
t.insert(format!("{l1}.input_layernorm.weight"), ones(&[h]));
|
||||
t.insert(format!("{l1}.post_attention_layernorm.weight"), ones(&[h]));
|
||||
moe(&mut t, &format!("{l1}.mlp"));
|
||||
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let path = dir.path().join("model.safetensors");
|
||||
candle_core::safetensors::save(&t, &path).expect("save safetensors");
|
||||
// SAFETY: mmap of a file this test just wrote; nothing mutates it.
|
||||
let vb = unsafe {
|
||||
candle_nn::var_builder::ShardedSafeTensors::var_builder(
|
||||
std::slice::from_ref(&path),
|
||||
DType::F32,
|
||||
&dev,
|
||||
)
|
||||
.expect("build ShardedVarBuilder")
|
||||
};
|
||||
|
||||
let mut model = Qwen3_5ForCausalLM::new(cfg, vb).expect("load tiny qwen3_next checkpoint");
|
||||
let input = Tensor::new(&[1u32, 5, 9], &dev)
|
||||
.unwrap()
|
||||
.unsqueeze(0)
|
||||
.unwrap();
|
||||
let logits = model.forward(&input, 0).expect("forward");
|
||||
assert_eq!(logits.dims(), &[1, 1, vocab]);
|
||||
let v: Vec<f32> = logits.flatten_all().unwrap().to_vec1().unwrap();
|
||||
assert!(v.iter().all(|x| x.is_finite()), "logits must be finite");
|
||||
}
|
||||
|
||||
/// `mlp_only_layers` and `decoder_sparse_step` gate `layer_uses_moe`
|
||||
/// per the upstream convention.
|
||||
#[test]
|
||||
fn layer_uses_moe_respects_step_and_exclusions() {
|
||||
let raw = r#"{
|
||||
"model_type": "qwen3_next",
|
||||
"vocab_size": 8, "hidden_size": 8, "intermediate_size": 8,
|
||||
"num_hidden_layers": 8, "num_attention_heads": 2,
|
||||
"num_key_value_heads": 1, "head_dim": 4,
|
||||
"max_position_embeddings": 128, "rms_norm_eps": 1e-6,
|
||||
"num_experts": 4, "num_experts_per_tok": 2,
|
||||
"moe_intermediate_size": 8,
|
||||
"decoder_sparse_step": 2,
|
||||
"mlp_only_layers": [3]
|
||||
}"#;
|
||||
let cfg = Config::from_config_json(raw).expect("parse");
|
||||
let t = &cfg.text_config;
|
||||
// step 2 → layers 1, 3, 5, 7 are on the sparse grid…
|
||||
assert!(!t.layer_uses_moe(0));
|
||||
assert!(t.layer_uses_moe(1));
|
||||
// …but 3 is excluded by mlp_only_layers.
|
||||
assert!(!t.layer_uses_moe(3));
|
||||
assert!(t.layer_uses_moe(5));
|
||||
}
|
||||
|
||||
/// `splice_runs` replaces (1, L, H) embedding rows at the given
|
||||
|
||||
360
crates/neuron/src/harness/arch/qwen3_5/moe.rs
Normal file
360
crates/neuron/src/harness/arch/qwen3_5/moe.rs
Normal file
@@ -0,0 +1,360 @@
|
||||
//! High-sparsity MoE FFN block for the qwen3_next family (#92).
|
||||
//!
|
||||
//! Qwen3-Next-80B-A3B replaces the dense SwiGLU in (almost) every
|
||||
//! decoder layer with `Qwen3NextSparseMoeBlock`: a top-k router over
|
||||
//! `num_experts` small SwiGLU experts, plus an always-on **shared
|
||||
//! expert** mixed in through a per-token sigmoid gate:
|
||||
//!
|
||||
//! ```text
|
||||
//! probs = softmax(gate(x)) # over ALL experts, f32
|
||||
//! w, idx = topk(probs, num_experts_per_tok)
|
||||
//! w = w / sum(w) # iff norm_topk_prob
|
||||
//! routed = Σ_j w_j · expert_{idx_j}(x)
|
||||
//! shared = sigmoid(shared_expert_gate(x)) · shared_expert(x)
|
||||
//! y = routed + shared
|
||||
//! ```
|
||||
//!
|
||||
//! Routing follows the upstream softmax-then-topk order (NOT
|
||||
//! topk-then-softmax — the renormalisation only equals softmax over
|
||||
//! the selected logits when `norm_topk_prob` is on, and the reference
|
||||
//! renormalises the *global* softmax values).
|
||||
//!
|
||||
//! ## Dispatch strategy
|
||||
//!
|
||||
//! This is the correctness-first implementation: a host-side scatter
|
||||
//! loop over the experts that actually received tokens (the pattern
|
||||
//! candle-transformers' `Qwen3SparseMoeBlock` uses). Batch-1 decode
|
||||
//! touches `num_experts_per_tok` experts per layer; prefill batches
|
||||
//! per-expert token groups. The fused grouped-GEMM path (slice 4)
|
||||
//! replaces the loop behind the same `forward` signature.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use candle_core::{DType, Module, Tensor};
|
||||
use candle_nn::Linear;
|
||||
use candle_nn::var_builder::ShardedVarBuilder;
|
||||
|
||||
use super::TextConfig;
|
||||
use super::mlp::Qwen3_5MLP;
|
||||
|
||||
pub struct Qwen3_5MoeBlock {
|
||||
/// Router: `(num_experts, hidden)`, checkpoint name `mlp.gate`.
|
||||
gate: Linear,
|
||||
/// Routed experts, checkpoint names `mlp.experts.{i}.{gate,up,down}_proj`.
|
||||
experts: Vec<Qwen3_5MLP>,
|
||||
/// Always-on expert, `mlp.shared_expert.*`. `None` when the config
|
||||
/// declares no shared expert (Qwen3-30B-A3B style).
|
||||
shared_expert: Option<Qwen3_5MLP>,
|
||||
/// Per-token sigmoid mix for the shared expert: `(1, hidden)`,
|
||||
/// checkpoint name `mlp.shared_expert_gate`.
|
||||
shared_expert_gate: Option<Linear>,
|
||||
num_experts_per_tok: usize,
|
||||
norm_topk_prob: bool,
|
||||
}
|
||||
|
||||
impl Qwen3_5MoeBlock {
|
||||
pub fn load(cfg: &TextConfig, vb: &ShardedVarBuilder) -> Result<Self> {
|
||||
anyhow::ensure!(
|
||||
cfg.num_experts > 0 && cfg.num_experts_per_tok > 0 && cfg.moe_intermediate_size > 0,
|
||||
"MoE block needs num_experts ({}), num_experts_per_tok ({}) and \
|
||||
moe_intermediate_size ({}) all > 0",
|
||||
cfg.num_experts,
|
||||
cfg.num_experts_per_tok,
|
||||
cfg.moe_intermediate_size,
|
||||
);
|
||||
anyhow::ensure!(
|
||||
cfg.num_experts_per_tok <= cfg.num_experts,
|
||||
"num_experts_per_tok ({}) exceeds num_experts ({})",
|
||||
cfg.num_experts_per_tok,
|
||||
cfg.num_experts,
|
||||
);
|
||||
|
||||
let h = cfg.hidden_size;
|
||||
|
||||
let gate_weight = vb
|
||||
.pp("gate")
|
||||
.get((cfg.num_experts, h), "weight")
|
||||
.with_context(|| format!("load '{}/gate/weight'", vb.prefix()))?;
|
||||
let gate = Linear::new(gate_weight, None);
|
||||
|
||||
let experts_vb = vb.pp("experts");
|
||||
let mut experts = Vec::with_capacity(cfg.num_experts);
|
||||
for i in 0..cfg.num_experts {
|
||||
experts.push(
|
||||
Qwen3_5MLP::load_with_dims(&experts_vb.pp(i), h, cfg.moe_intermediate_size)
|
||||
.with_context(|| format!("load expert {i}"))?,
|
||||
);
|
||||
}
|
||||
|
||||
let (shared_expert, shared_expert_gate) = if cfg.shared_expert_intermediate_size > 0 {
|
||||
let shared = Qwen3_5MLP::load_with_dims(
|
||||
&vb.pp("shared_expert"),
|
||||
h,
|
||||
cfg.shared_expert_intermediate_size,
|
||||
)
|
||||
.context("load shared_expert")?;
|
||||
let gate_w = vb
|
||||
.pp("shared_expert_gate")
|
||||
.get((1, h), "weight")
|
||||
.with_context(|| format!("load '{}/shared_expert_gate/weight'", vb.prefix()))?;
|
||||
(Some(shared), Some(Linear::new(gate_w, None)))
|
||||
} else {
|
||||
(None, None)
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
gate,
|
||||
experts,
|
||||
shared_expert,
|
||||
shared_expert_gate,
|
||||
num_experts_per_tok: cfg.num_experts_per_tok,
|
||||
norm_topk_prob: cfg.norm_topk_prob,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-expert routing assignment: `(token_rows, weights)` per expert,
|
||||
/// produced by [`route_scatter`].
|
||||
pub(crate) type ExpertAssignments = (Vec<Vec<u32>>, Vec<Vec<f32>>);
|
||||
|
||||
/// Router + host-side scatter shared by the single-GPU and TP MoE
|
||||
/// blocks (#92): softmax over ALL experts in f32 → descending-argsort
|
||||
/// top-k → renormalise iff `norm_topk_prob` → per-expert token-row and
|
||||
/// weight lists. Under TP the router weight is replicated, so every
|
||||
/// rank computes identical assignments with zero communication.
|
||||
pub(crate) fn route_scatter(
|
||||
gate: &Linear,
|
||||
xs_flat: &Tensor,
|
||||
num_experts: usize,
|
||||
num_experts_per_tok: usize,
|
||||
norm_topk_prob: bool,
|
||||
) -> candle_core::Result<ExpertAssignments> {
|
||||
let n_tokens = xs_flat.dim(0)?;
|
||||
// Router probabilities in f32 (reference uses float softmax
|
||||
// regardless of activations dtype).
|
||||
let router_logits = gate.forward(xs_flat)?;
|
||||
let probs = candle_nn::ops::softmax_last_dim(&router_logits.to_dtype(DType::F32)?)?;
|
||||
|
||||
// Top-k selection: descending argsort, take the first k. The
|
||||
// renormalisation (iff norm_topk_prob) divides by the sum of
|
||||
// the selected global-softmax values.
|
||||
let sorted = probs.arg_sort_last_dim(false)?;
|
||||
let topk_idx = sorted.narrow(1, 0, num_experts_per_tok)?.contiguous()?;
|
||||
let mut topk_w = probs.gather(&topk_idx, 1)?;
|
||||
if norm_topk_prob {
|
||||
let denom = topk_w.sum_keepdim(1)?;
|
||||
topk_w = topk_w.broadcast_div(&denom)?;
|
||||
}
|
||||
|
||||
// Host-side scatter: token row lists per expert. Cheap relative
|
||||
// to the expert GEMMs; replaced by grouped-GEMM in slice 4.
|
||||
let idx_host: Vec<Vec<u32>> = topk_idx.to_vec2()?;
|
||||
let w_host: Vec<Vec<f32>> = topk_w.to_vec2()?;
|
||||
let mut tokens_for: Vec<Vec<u32>> = vec![Vec::new(); num_experts];
|
||||
let mut weights_for: Vec<Vec<f32>> = vec![Vec::new(); num_experts];
|
||||
for t in 0..n_tokens {
|
||||
for j in 0..num_experts_per_tok {
|
||||
let e = idx_host[t][j] as usize;
|
||||
tokens_for[e].push(t as u32);
|
||||
weights_for[e].push(w_host[t][j]);
|
||||
}
|
||||
}
|
||||
Ok((tokens_for, weights_for))
|
||||
}
|
||||
|
||||
impl Module for Qwen3_5MoeBlock {
|
||||
fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||||
let (b, l, hidden) = xs.dims3()?;
|
||||
let xs_flat = xs.reshape(((), hidden))?;
|
||||
|
||||
let (tokens_for, weights_for) = route_scatter(
|
||||
&self.gate,
|
||||
&xs_flat,
|
||||
self.experts.len(),
|
||||
self.num_experts_per_tok,
|
||||
self.norm_topk_prob,
|
||||
)?;
|
||||
|
||||
let mut ys = xs_flat.zeros_like()?;
|
||||
for (e, expert) in self.experts.iter().enumerate() {
|
||||
if tokens_for[e].is_empty() {
|
||||
continue;
|
||||
}
|
||||
let rows = Tensor::new(tokens_for[e].as_slice(), xs.device())?;
|
||||
let picked = xs_flat.index_select(&rows, 0)?;
|
||||
let out = expert.forward(&picked)?;
|
||||
let w = Tensor::new(weights_for[e].as_slice(), xs.device())?
|
||||
.to_dtype(out.dtype())?
|
||||
.reshape(((), 1))?;
|
||||
ys = ys.index_add(&rows, &out.broadcast_mul(&w)?, 0)?;
|
||||
}
|
||||
|
||||
if let (Some(shared), Some(gate)) = (&self.shared_expert, &self.shared_expert_gate) {
|
||||
let mix = candle_nn::ops::sigmoid(&gate.forward(&xs_flat)?)?;
|
||||
let shared_out = shared.forward(&xs_flat)?.broadcast_mul(&mix)?;
|
||||
ys = (ys + shared_out)?;
|
||||
}
|
||||
|
||||
ys.reshape((b, l, hidden))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use candle_core::Device;
|
||||
|
||||
fn randn(shape: &[usize]) -> Tensor {
|
||||
Tensor::randn(0f32, 0.5f32, shape, &Device::Cpu).unwrap()
|
||||
}
|
||||
|
||||
fn rand_mlp(hidden: usize, inter: usize) -> Qwen3_5MLP {
|
||||
Qwen3_5MLP::from_weights(
|
||||
Linear::new(randn(&[inter, hidden]), None),
|
||||
Linear::new(randn(&[inter, hidden]), None),
|
||||
Linear::new(randn(&[hidden, inter]), None),
|
||||
)
|
||||
}
|
||||
|
||||
/// The batched scatter forward must equal a per-token dense
|
||||
/// reference: route each token independently (host softmax → top-k
|
||||
/// → renorm), run its selected experts one by one, and mix in the
|
||||
/// shared expert through the sigmoid gate. Catches indexing,
|
||||
/// weighting, and renormalisation bugs in the scatter path.
|
||||
#[test]
|
||||
fn scatter_forward_matches_per_token_reference() {
|
||||
let (hidden, inter, n_exp, top_k) = (8, 4, 6, 2);
|
||||
|
||||
let block = Qwen3_5MoeBlock {
|
||||
gate: Linear::new(randn(&[n_exp, hidden]), None),
|
||||
experts: (0..n_exp).map(|_| rand_mlp(hidden, inter)).collect(),
|
||||
shared_expert: Some(rand_mlp(hidden, inter)),
|
||||
shared_expert_gate: Some(Linear::new(randn(&[1, hidden]), None)),
|
||||
num_experts_per_tok: top_k,
|
||||
norm_topk_prob: true,
|
||||
};
|
||||
|
||||
let (b, l) = (2, 3);
|
||||
let xs = randn(&[b, l, hidden]);
|
||||
let got = block.forward(&xs).unwrap();
|
||||
assert_eq!(got.dims(), &[b, l, hidden]);
|
||||
|
||||
let xs_flat = xs.reshape(((), hidden)).unwrap();
|
||||
let logits: Vec<Vec<f32>> = block.gate.forward(&xs_flat).unwrap().to_vec2().unwrap();
|
||||
let got_flat: Vec<Vec<f32>> = got.reshape(((), hidden)).unwrap().to_vec2().unwrap();
|
||||
|
||||
for t in 0..b * l {
|
||||
// Host-side softmax over all experts, then top-k + renorm.
|
||||
let max = logits[t].iter().cloned().fold(f32::MIN, f32::max);
|
||||
let exps: Vec<f32> = logits[t].iter().map(|v| (v - max).exp()).collect();
|
||||
let sum: f32 = exps.iter().sum();
|
||||
let probs: Vec<f32> = exps.iter().map(|e| e / sum).collect();
|
||||
let mut order: Vec<usize> = (0..n_exp).collect();
|
||||
order.sort_by(|&a, &b| probs[b].partial_cmp(&probs[a]).unwrap());
|
||||
let selected = &order[..top_k];
|
||||
let denom: f32 = selected.iter().map(|&e| probs[e]).sum();
|
||||
|
||||
let row = xs_flat.narrow(0, t, 1).unwrap();
|
||||
let mut expect = vec![0f32; hidden];
|
||||
for &e in selected {
|
||||
let w = probs[e] / denom;
|
||||
let out: Vec<f32> = block.experts[e]
|
||||
.forward(&row)
|
||||
.unwrap()
|
||||
.flatten_all()
|
||||
.unwrap()
|
||||
.to_vec1()
|
||||
.unwrap();
|
||||
for (acc, o) in expect.iter_mut().zip(out) {
|
||||
*acc += w * o;
|
||||
}
|
||||
}
|
||||
let gate_v: f32 = block
|
||||
.shared_expert_gate
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.forward(&row)
|
||||
.unwrap()
|
||||
.flatten_all()
|
||||
.unwrap()
|
||||
.to_vec1::<f32>()
|
||||
.unwrap()[0];
|
||||
let mix = 1.0 / (1.0 + (-gate_v).exp());
|
||||
let shared: Vec<f32> = block
|
||||
.shared_expert
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.forward(&row)
|
||||
.unwrap()
|
||||
.flatten_all()
|
||||
.unwrap()
|
||||
.to_vec1()
|
||||
.unwrap();
|
||||
for (acc, s) in expect.iter_mut().zip(shared) {
|
||||
*acc += mix * s;
|
||||
}
|
||||
|
||||
for (i, (&g, &e)) in got_flat[t].iter().zip(expect.iter()).enumerate() {
|
||||
assert!(
|
||||
(g - e).abs() < 1e-4,
|
||||
"token {t} dim {i}: got {g}, expected {e}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Without a shared expert (Qwen3-30B-A3B shape) the block is pure
|
||||
/// routed output; without norm_topk_prob the raw global-softmax
|
||||
/// weights apply (they do NOT sum to 1 across the selected k).
|
||||
#[test]
|
||||
fn no_shared_expert_and_no_renorm() {
|
||||
let (hidden, inter, n_exp) = (4, 2, 3);
|
||||
let block = Qwen3_5MoeBlock {
|
||||
gate: Linear::new(randn(&[n_exp, hidden]), None),
|
||||
experts: (0..n_exp).map(|_| rand_mlp(hidden, inter)).collect(),
|
||||
shared_expert: None,
|
||||
shared_expert_gate: None,
|
||||
num_experts_per_tok: 1,
|
||||
norm_topk_prob: false,
|
||||
};
|
||||
let xs = randn(&[1, 1, hidden]);
|
||||
let got: Vec<f32> = block
|
||||
.forward(&xs)
|
||||
.unwrap()
|
||||
.flatten_all()
|
||||
.unwrap()
|
||||
.to_vec1()
|
||||
.unwrap();
|
||||
|
||||
// Reference: the argmax expert's output scaled by its raw
|
||||
// softmax probability.
|
||||
let logits: Vec<f32> = block
|
||||
.gate
|
||||
.forward(&xs.reshape(((), hidden)).unwrap())
|
||||
.unwrap()
|
||||
.flatten_all()
|
||||
.unwrap()
|
||||
.to_vec1()
|
||||
.unwrap();
|
||||
let max = logits.iter().cloned().fold(f32::MIN, f32::max);
|
||||
let exps: Vec<f32> = logits.iter().map(|v| (v - max).exp()).collect();
|
||||
let sum: f32 = exps.iter().sum();
|
||||
let best = (0..n_exp)
|
||||
.max_by(|&a, &b| exps[a].partial_cmp(&exps[b]).unwrap())
|
||||
.unwrap();
|
||||
let w = exps[best] / sum;
|
||||
let out: Vec<f32> = block.experts[best]
|
||||
.forward(&xs.reshape(((), hidden)).unwrap())
|
||||
.unwrap()
|
||||
.flatten_all()
|
||||
.unwrap()
|
||||
.to_vec1()
|
||||
.unwrap();
|
||||
for (i, (&g, &o)) in got.iter().zip(out.iter()).enumerate() {
|
||||
assert!(
|
||||
(g - w * o).abs() < 1e-5,
|
||||
"dim {i}: got {g}, expected {}",
|
||||
w * o
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -159,6 +159,19 @@ impl RotaryEmbedding {
|
||||
Ok((cos, sin))
|
||||
}
|
||||
|
||||
/// cos/sin gathered at arbitrary **per-row** positions — the
|
||||
/// batched-decode path (#98), where each batch row sits at its own
|
||||
/// sequence offset. Shape `(B, 1, rotary_dim/2)`: one position per
|
||||
/// row, one decode token per step. [`Self::apply_cos_sin`] detects
|
||||
/// the rank-3 shape and broadcasts per row instead of per position.
|
||||
pub fn batch_cos_sin(&self, positions: &[usize]) -> candle_core::Result<(Tensor, Tensor)> {
|
||||
let idx: Vec<u32> = positions.iter().map(|&p| p as u32).collect();
|
||||
let idx = Tensor::from_vec(idx, positions.len(), self.cos.device())?;
|
||||
let cos = self.cos.index_select(&idx, 0)?.unsqueeze(1)?;
|
||||
let sin = self.sin.index_select(&idx, 0)?.unsqueeze(1)?;
|
||||
Ok((cos, sin))
|
||||
}
|
||||
|
||||
/// cos/sin from explicit per-token 3D position ids, shape
|
||||
/// `(3, seq_len)` (axes: text, height, width). Builds each axis's
|
||||
/// frequencies and blends them at the interleave index sets, so
|
||||
@@ -185,7 +198,9 @@ impl RotaryEmbedding {
|
||||
}
|
||||
|
||||
/// Apply rotary to `q`, `k` (shape `(B, H, L, head_dim)`) using
|
||||
/// precomputed `cos`/`sin` of shape `(L, rotary_dim/2)`. Partial
|
||||
/// precomputed `cos`/`sin` of shape `(L, rotary_dim/2)` — or, for
|
||||
/// the batched-decode path (#98), `(B, L, rotary_dim/2)` with a
|
||||
/// distinct position per batch row (dispatch is on rank). Partial
|
||||
/// rotary: only the first `rotary_dim` dims rotate; the tail passes
|
||||
/// through unchanged.
|
||||
pub fn apply_cos_sin(
|
||||
@@ -197,9 +212,17 @@ impl RotaryEmbedding {
|
||||
) -> candle_core::Result<(Tensor, Tensor)> {
|
||||
let (_, _, _seq_len, head_dim_in) = q.dims4()?;
|
||||
debug_assert_eq!(head_dim_in, self.head_dim, "q head_dim mismatch");
|
||||
let per_row = cos.rank() == 3;
|
||||
let rope = |x: &Tensor| -> candle_core::Result<Tensor> {
|
||||
if per_row {
|
||||
rope_per_row(x, cos, sin)
|
||||
} else {
|
||||
candle_nn::rotary_emb::rope_slow(x, cos, sin)
|
||||
}
|
||||
};
|
||||
if self.rotary_dim == self.head_dim {
|
||||
let q_embed = candle_nn::rotary_emb::rope_slow(&q.contiguous()?, cos, sin)?;
|
||||
let k_embed = candle_nn::rotary_emb::rope_slow(&k.contiguous()?, cos, sin)?;
|
||||
let q_embed = rope(&q.contiguous()?)?;
|
||||
let k_embed = rope(&k.contiguous()?)?;
|
||||
Ok((q_embed, k_embed))
|
||||
} else {
|
||||
// Partial rotation: narrow → rotate → cat the untouched tail.
|
||||
@@ -212,8 +235,8 @@ impl RotaryEmbedding {
|
||||
.narrow(candle_core::D::Minus1, 0, self.rotary_dim)?
|
||||
.contiguous()?;
|
||||
let k_pass = k.narrow(candle_core::D::Minus1, self.rotary_dim, tail)?;
|
||||
let q_rotated = candle_nn::rotary_emb::rope_slow(&q_rot, cos, sin)?;
|
||||
let k_rotated = candle_nn::rotary_emb::rope_slow(&k_rot, cos, sin)?;
|
||||
let q_rotated = rope(&q_rot)?;
|
||||
let k_rotated = rope(&k_rot)?;
|
||||
let q_embed =
|
||||
Tensor::cat(&[&q_rotated, &q_pass.contiguous()?], candle_core::D::Minus1)?;
|
||||
let k_embed =
|
||||
@@ -223,6 +246,27 @@ impl RotaryEmbedding {
|
||||
}
|
||||
}
|
||||
|
||||
/// GLM rotate-half (same convention as candle's private
|
||||
/// `rotary_emb::rotate_half`: `cat(-x2, x1)`).
|
||||
fn rotate_half(x: &Tensor) -> candle_core::Result<Tensor> {
|
||||
let last = x.dim(candle_core::D::Minus1)?;
|
||||
let x1 = x.narrow(candle_core::D::Minus1, 0, last / 2)?;
|
||||
let x2 = x.narrow(candle_core::D::Minus1, last / 2, last - last / 2)?;
|
||||
Tensor::cat(&[&x2.neg()?, &x1], candle_core::D::Minus1)
|
||||
}
|
||||
|
||||
/// Per-row rope apply for batched decode: `x` is `(B, H, L, rot)`,
|
||||
/// `cos`/`sin` are `(B, L, rot/2)` — each batch row gets its own
|
||||
/// position's rotation (candle's `rope_slow` only broadcasts one
|
||||
/// `(L, rot/2)` table across the whole batch).
|
||||
fn rope_per_row(x: &Tensor, cos: &Tensor, sin: &Tensor) -> candle_core::Result<Tensor> {
|
||||
// (B, L, half) → duplicate pairs → (B, 1, L, rot) for broadcast
|
||||
// over the head dim.
|
||||
let cos = Tensor::cat(&[cos, cos], candle_core::D::Minus1)?.unsqueeze(1)?;
|
||||
let sin = Tensor::cat(&[sin, sin], candle_core::D::Minus1)?.unsqueeze(1)?;
|
||||
x.broadcast_mul(&cos)? + rotate_half(x)?.broadcast_mul(&sin)?
|
||||
}
|
||||
|
||||
/// Compute interleaved-M-RoPE 3D position ids for a full prompt that may
|
||||
/// contain image-placeholder runs, plus the decode `rope_delta`.
|
||||
///
|
||||
@@ -564,6 +608,49 @@ mod tests {
|
||||
assert!((last[1] - (2.0 * inv[1]).cos()).abs() < 1e-5);
|
||||
}
|
||||
|
||||
/// `batch_cos_sin` at positions [5, 9, 0] must gather exactly the
|
||||
/// rows `plain_cos_sin` would produce for each position alone.
|
||||
#[test]
|
||||
fn batch_cos_sin_gathers_per_row_positions() {
|
||||
let dev = Device::Cpu;
|
||||
let rope = RotaryEmbedding::new(DType::F32, &qwen36_cfg(), &dev).unwrap();
|
||||
let half = rope.inv_freq.dim(1).unwrap();
|
||||
let positions = [5usize, 9, 0];
|
||||
let (bc, bs) = rope.batch_cos_sin(&positions).unwrap();
|
||||
assert_eq!(bc.dims(), &[3, 1, half]);
|
||||
assert_eq!(bs.dims(), &[3, 1, half]);
|
||||
for (row, &p) in positions.iter().enumerate() {
|
||||
let (pc, ps) = rope.plain_cos_sin(p, 1).unwrap();
|
||||
let dc = (bc.i(row).unwrap() - pc).unwrap().abs().unwrap();
|
||||
let ds = (bs.i(row).unwrap() - ps).unwrap().abs().unwrap();
|
||||
assert!(dc.max_all().unwrap().to_scalar::<f32>().unwrap() < 1e-6);
|
||||
assert!(ds.max_all().unwrap().to_scalar::<f32>().unwrap() < 1e-6);
|
||||
}
|
||||
}
|
||||
|
||||
/// When every row sits at the same position, the per-row rank-3
|
||||
/// apply path must reproduce the shared rank-2 (`rope_slow`) path
|
||||
/// exactly — the invariant that makes the rank dispatch in
|
||||
/// `apply_cos_sin` safe.
|
||||
#[test]
|
||||
fn per_row_apply_matches_shared_when_uniform() {
|
||||
let dev = Device::Cpu;
|
||||
let rope = RotaryEmbedding::new(DType::F32, &qwen36_cfg(), &dev).unwrap();
|
||||
let q = Tensor::randn(0f32, 1f32, (2, 2, 1, 256), &dev).unwrap();
|
||||
let k = Tensor::randn(0f32, 1f32, (2, 2, 1, 256), &dev).unwrap();
|
||||
|
||||
let (c2, s2) = rope.plain_cos_sin(7, 1).unwrap();
|
||||
let (qa, ka) = rope.apply_cos_sin(&q, &k, &c2, &s2).unwrap();
|
||||
|
||||
let (c3, s3) = rope.batch_cos_sin(&[7, 7]).unwrap();
|
||||
let (qb, kb) = rope.apply_cos_sin(&q, &k, &c3, &s3).unwrap();
|
||||
|
||||
let dq = (qa - qb).unwrap().abs().unwrap().max_all().unwrap();
|
||||
let dk = (ka - kb).unwrap().abs().unwrap().max_all().unwrap();
|
||||
assert!(dq.to_scalar::<f32>().unwrap() < 1e-6, "q mismatch {dq:?}");
|
||||
assert!(dk.to_scalar::<f32>().unwrap() < 1e-6, "k mismatch {dk:?}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_rope_index_196_is_14x14() {
|
||||
let mut ids = vec![1u32]; // one text token
|
||||
|
||||
@@ -84,6 +84,232 @@ impl KvCacheSnapshot {
|
||||
}
|
||||
}
|
||||
|
||||
/// Batched cache state assembled from per-sequence snapshots (#98).
|
||||
/// Install with `Qwen3_5Model::restore_kv_cache(&self.snapshot)`; the
|
||||
/// forward then runs lockstep batched decode with
|
||||
/// `forward_batch_decode` using `prefix_lens[i] + step` positions and
|
||||
/// the padding mask from `Qwen3_5Model::batch_decode_mask`.
|
||||
pub struct BatchedKvState {
|
||||
/// Per-layer `(B, …)` state: attention K/V right-padded along the
|
||||
/// sequence axis to `padded_len` and `cat`ed on dim 0; GDN
|
||||
/// conv/recurrent states `cat`ed on dim 0 (position-free).
|
||||
pub snapshot: KvCacheSnapshot,
|
||||
/// The uniform KV sequence length every row was padded to — the
|
||||
/// max prefix length in the batch. Decode appends start here.
|
||||
pub padded_len: usize,
|
||||
/// Each row's true prefix length. Columns `[prefix_lens[i],
|
||||
/// padded_len)` of row `i` are zero-padding and must stay masked.
|
||||
pub prefix_lens: Vec<usize>,
|
||||
}
|
||||
|
||||
/// Assemble per-sequence snapshots into one batched cache state.
|
||||
/// `seqs` pairs each snapshot with its true token length (the caller
|
||||
/// tracks prompt token counts; attention K/V lengths are validated
|
||||
/// against it). All snapshots must come from single-sequence (`B=1`)
|
||||
/// prefills of the same model, with `rope_delta == 0` (text-only —
|
||||
/// vision requests don't batch, #98 v1).
|
||||
///
|
||||
/// Keys are stored post-RoPE, so right-padding does not disturb
|
||||
/// position correctness: a row's cached keys keep the rotation of
|
||||
/// their true positions, the garbage columns are masked, and new
|
||||
/// tokens rotate at `prefix_len + step` while landing at storage
|
||||
/// column `padded_len + step`.
|
||||
pub fn assemble_batch(seqs: &[(&KvCacheSnapshot, usize)]) -> candle_core::Result<BatchedKvState> {
|
||||
let Some((first, _)) = seqs.first() else {
|
||||
candle_core::bail!("assemble_batch: empty batch");
|
||||
};
|
||||
let n_layers = first.layers.len();
|
||||
let prefix_lens: Vec<usize> = seqs.iter().map(|&(_, len)| len).collect();
|
||||
let padded_len = *prefix_lens.iter().max().expect("non-empty");
|
||||
for (snap, len) in seqs {
|
||||
if snap.layers.len() != n_layers {
|
||||
candle_core::bail!(
|
||||
"assemble_batch: snapshot layer count mismatch ({} vs {n_layers})",
|
||||
snap.layers.len()
|
||||
);
|
||||
}
|
||||
if snap.rope_delta != 0 {
|
||||
candle_core::bail!(
|
||||
"assemble_batch: rope_delta {} != 0 — vision-positioned sequences cannot batch",
|
||||
snap.rope_delta
|
||||
);
|
||||
}
|
||||
if *len == 0 {
|
||||
candle_core::bail!("assemble_batch: zero-length sequence");
|
||||
}
|
||||
}
|
||||
|
||||
let mut layers = Vec::with_capacity(n_layers);
|
||||
for li in 0..n_layers {
|
||||
layers.push(assemble_layer(seqs, li, padded_len)?);
|
||||
}
|
||||
Ok(BatchedKvState {
|
||||
snapshot: KvCacheSnapshot {
|
||||
layers,
|
||||
rope_delta: 0,
|
||||
},
|
||||
padded_len,
|
||||
prefix_lens,
|
||||
})
|
||||
}
|
||||
|
||||
fn assemble_layer(
|
||||
seqs: &[(&KvCacheSnapshot, usize)],
|
||||
li: usize,
|
||||
padded_len: usize,
|
||||
) -> candle_core::Result<LayerKvSnapshot> {
|
||||
match &seqs[0].0.layers[li] {
|
||||
LayerKvSnapshot::Full(_) => {
|
||||
let mut ks = Vec::with_capacity(seqs.len());
|
||||
let mut vs = Vec::with_capacity(seqs.len());
|
||||
for (row, (snap, len)) in seqs.iter().enumerate() {
|
||||
let LayerKvSnapshot::Full(Some((k, v))) = &snap.layers[li] else {
|
||||
candle_core::bail!(
|
||||
"assemble_batch: row {row} layer {li} is not a populated \
|
||||
full-attention snapshot"
|
||||
);
|
||||
};
|
||||
let (b, _h, s, _d) = k.dims4()?;
|
||||
if b != 1 {
|
||||
candle_core::bail!(
|
||||
"assemble_batch: row {row} layer {li} has batch dim {b}, want 1"
|
||||
);
|
||||
}
|
||||
if s != *len {
|
||||
candle_core::bail!(
|
||||
"assemble_batch: row {row} layer {li} KV length {s} != declared \
|
||||
sequence length {len}"
|
||||
);
|
||||
}
|
||||
ks.push(pad_seq(k, padded_len)?);
|
||||
vs.push(pad_seq(v, padded_len)?);
|
||||
}
|
||||
let k = Tensor::cat(&ks, 0)?;
|
||||
let v = Tensor::cat(&vs, 0)?;
|
||||
Ok(LayerKvSnapshot::Full(Some((k, v))))
|
||||
}
|
||||
LayerKvSnapshot::Linear { .. } => {
|
||||
let mut convs = Vec::with_capacity(seqs.len());
|
||||
let mut recs = Vec::with_capacity(seqs.len());
|
||||
for (row, (snap, _)) in seqs.iter().enumerate() {
|
||||
let LayerKvSnapshot::Linear {
|
||||
conv_state: Some(conv),
|
||||
recurrent_state: Some(rec),
|
||||
} = &snap.layers[li]
|
||||
else {
|
||||
candle_core::bail!(
|
||||
"assemble_batch: row {row} layer {li} is not a populated \
|
||||
linear-attention snapshot"
|
||||
);
|
||||
};
|
||||
convs.push(conv.clone());
|
||||
recs.push(rec.clone());
|
||||
}
|
||||
Ok(LayerKvSnapshot::Linear {
|
||||
conv_state: Some(Tensor::cat(&convs, 0)?),
|
||||
recurrent_state: Some(Tensor::cat(&recs, 0)?),
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract one row of a batched cache snapshot back into a contiguous
|
||||
/// single-sequence snapshot — the defragment half of batch membership
|
||||
/// changes (#98). A join/leave rebatches by extracting every surviving
|
||||
/// row (each becomes a gap-free `B=1` snapshot of length `prefix_len +
|
||||
/// steps`) and running [`assemble_batch`] over them again, so the
|
||||
/// "each row has exactly one padding gap" invariant holds for the new
|
||||
/// batch too.
|
||||
///
|
||||
/// `row` indexes the batch dim. The row's valid attention KV is its
|
||||
/// prefix `[0, prefix_len)` plus the lockstep decode columns
|
||||
/// `[padded_len, padded_len + steps)`; the gap between them is padding
|
||||
/// and is dropped. GDN states are position-free — the row is sliced
|
||||
/// out whole and **deep-copied** (the live buffers are mutated in
|
||||
/// place by the CUDA kernels; see the module notes on copy semantics).
|
||||
pub fn extract_row(
|
||||
snap: &KvCacheSnapshot,
|
||||
row: usize,
|
||||
prefix_len: usize,
|
||||
padded_len: usize,
|
||||
steps: usize,
|
||||
) -> candle_core::Result<KvCacheSnapshot> {
|
||||
if prefix_len == 0 || prefix_len > padded_len {
|
||||
candle_core::bail!(
|
||||
"extract_row: prefix_len {prefix_len} out of range (padded_len {padded_len})"
|
||||
);
|
||||
}
|
||||
let mut layers = Vec::with_capacity(snap.layers.len());
|
||||
for (li, layer) in snap.layers.iter().enumerate() {
|
||||
layers.push(match layer {
|
||||
LayerKvSnapshot::Full(Some((k, v))) => {
|
||||
let (b, _h, s, _d) = k.dims4()?;
|
||||
if row >= b {
|
||||
candle_core::bail!("extract_row: row {row} out of range (batch {b})");
|
||||
}
|
||||
if s != padded_len + steps {
|
||||
candle_core::bail!(
|
||||
"extract_row: layer {li} KV length {s} != padded_len {padded_len} + \
|
||||
steps {steps}"
|
||||
);
|
||||
}
|
||||
LayerKvSnapshot::Full(Some((
|
||||
unpad_row(k, row, prefix_len, padded_len, steps)?,
|
||||
unpad_row(v, row, prefix_len, padded_len, steps)?,
|
||||
)))
|
||||
}
|
||||
LayerKvSnapshot::Full(None) => {
|
||||
candle_core::bail!("extract_row: layer {li} has an empty attention cache")
|
||||
}
|
||||
LayerKvSnapshot::Linear {
|
||||
conv_state: Some(conv),
|
||||
recurrent_state: Some(rec),
|
||||
} => LayerKvSnapshot::Linear {
|
||||
conv_state: Some(conv.narrow(0, row, 1)?.copy()?),
|
||||
recurrent_state: Some(rec.narrow(0, row, 1)?.copy()?),
|
||||
},
|
||||
LayerKvSnapshot::Linear { .. } => {
|
||||
candle_core::bail!("extract_row: layer {li} has unpopulated GDN state")
|
||||
}
|
||||
});
|
||||
}
|
||||
Ok(KvCacheSnapshot {
|
||||
layers,
|
||||
rope_delta: 0,
|
||||
})
|
||||
}
|
||||
|
||||
/// Slice `row` out of a batched `(B, H, padded_len + steps, D)` K or V
|
||||
/// tensor and drop its padding gap, yielding an owned contiguous
|
||||
/// `(1, H, prefix_len + steps, D)` tensor.
|
||||
fn unpad_row(
|
||||
t: &Tensor,
|
||||
row: usize,
|
||||
prefix_len: usize,
|
||||
padded_len: usize,
|
||||
steps: usize,
|
||||
) -> candle_core::Result<Tensor> {
|
||||
let r = t.narrow(0, row, 1)?;
|
||||
let prefix = r.narrow(2, 0, prefix_len)?;
|
||||
if steps == 0 {
|
||||
return prefix.contiguous()?.copy();
|
||||
}
|
||||
let decoded = r.narrow(2, padded_len, steps)?;
|
||||
Tensor::cat(&[&prefix.contiguous()?, &decoded.contiguous()?], 2)
|
||||
}
|
||||
|
||||
/// Right-pad a `(1, H, S, D)` K or V tensor with zeros along the
|
||||
/// sequence axis to `padded_len`. Zero columns are inert: the padding
|
||||
/// mask keeps every query from attending to them.
|
||||
fn pad_seq(t: &Tensor, padded_len: usize) -> candle_core::Result<Tensor> {
|
||||
let (b, h, s, d) = t.dims4()?;
|
||||
if s == padded_len {
|
||||
return Ok(t.clone());
|
||||
}
|
||||
let pad = Tensor::zeros((b, h, padded_len - s, d), t.dtype(), t.device())?;
|
||||
Tensor::cat(&[t, &pad], 2)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::super::{Qwen3_5Model, RopeParameters, TextConfig};
|
||||
@@ -119,6 +345,13 @@ mod tests {
|
||||
linear_key_head_dim: 4,
|
||||
linear_value_head_dim: 4,
|
||||
linear_conv_kernel_dim: 4,
|
||||
num_experts: 0,
|
||||
num_experts_per_tok: 0,
|
||||
moe_intermediate_size: 0,
|
||||
shared_expert_intermediate_size: 0,
|
||||
decoder_sparse_step: 1,
|
||||
mlp_only_layers: Vec::new(),
|
||||
norm_topk_prob: false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -201,7 +434,7 @@ mod tests {
|
||||
)
|
||||
.expect("build ShardedVarBuilder")
|
||||
};
|
||||
Qwen3_5Model::load(cfg, &vb).expect("load tiny qwen3_5 model")
|
||||
Qwen3_5Model::load(cfg, &vb, "model.language_model").expect("load tiny qwen3_5 model")
|
||||
}
|
||||
|
||||
fn forward_tokens(model: &mut Qwen3_5Model, tokens: &[u32], offset: usize) -> Vec<f32> {
|
||||
@@ -271,6 +504,228 @@ mod tests {
|
||||
assert!(diff < 1e-6, "second restore diverged: {diff}");
|
||||
}
|
||||
|
||||
/// The gold test for #98 slice 1: ragged sequences decoded in one
|
||||
/// lockstep batch (assembled from per-sequence snapshots, per-row
|
||||
/// positions, padding mask) must match the same sequences decoded
|
||||
/// sequentially, hidden-state for hidden-state at every step.
|
||||
#[test]
|
||||
fn batched_decode_matches_sequential() {
|
||||
use candle_core::IndexOp;
|
||||
let cfg = tiny_config();
|
||||
let mut model = tiny_model(&cfg);
|
||||
|
||||
let prompts: [&[u32]; 3] = [&[1, 2, 3], &[4, 5], &[7, 7, 2, 5, 6]];
|
||||
let steps: [&[u32]; 3] = [&[11, 12, 13, 14], &[9, 8, 7, 6], &[21, 22, 23, 24]];
|
||||
let n_steps = 4;
|
||||
|
||||
// Sequential reference: each sequence decoded alone.
|
||||
let mut expected: Vec<Vec<Vec<f32>>> = Vec::new(); // [row][step]
|
||||
for (prompt, toks) in prompts.iter().zip(steps.iter()) {
|
||||
model.clear_kv_cache();
|
||||
forward_tokens(&mut model, prompt, 0);
|
||||
let mut per_step = Vec::new();
|
||||
for (t, tok) in toks.iter().enumerate() {
|
||||
per_step.push(forward_tokens(&mut model, &[*tok], prompt.len() + t));
|
||||
}
|
||||
expected.push(per_step);
|
||||
}
|
||||
|
||||
// Batched: prefill each sequence alone, snapshot, assemble.
|
||||
let mut snaps = Vec::new();
|
||||
for prompt in prompts.iter() {
|
||||
model.clear_kv_cache();
|
||||
forward_tokens(&mut model, prompt, 0);
|
||||
snaps.push(model.snapshot_kv_cache().expect("snapshot"));
|
||||
}
|
||||
let seqs: Vec<(&super::KvCacheSnapshot, usize)> = snaps
|
||||
.iter()
|
||||
.zip(prompts.iter())
|
||||
.map(|(s, p)| (s, p.len()))
|
||||
.collect();
|
||||
let batch = super::assemble_batch(&seqs).expect("assemble");
|
||||
assert_eq!(batch.padded_len, 5);
|
||||
assert_eq!(batch.prefix_lens, vec![3, 2, 5]);
|
||||
model
|
||||
.restore_kv_cache(&batch.snapshot)
|
||||
.expect("install batched state");
|
||||
|
||||
for t in 0..n_steps {
|
||||
let toks: Vec<u32> = steps.iter().map(|s| s[t]).collect();
|
||||
let input = Tensor::from_vec(toks, (3, 1), &Device::Cpu).unwrap();
|
||||
let positions: Vec<usize> = prompts.iter().map(|p| p.len() + t).collect();
|
||||
let total_len = batch.padded_len + t + 1;
|
||||
let mask = model
|
||||
.batch_decode_mask(&batch.prefix_lens, batch.padded_len, total_len)
|
||||
.expect("mask");
|
||||
assert!(mask.is_some(), "ragged batch must be masked");
|
||||
let h = model
|
||||
.forward_batch_decode(&input, &positions, mask.as_ref())
|
||||
.expect("batched step");
|
||||
assert_eq!(h.dims()[0], 3);
|
||||
for row in 0..3 {
|
||||
let got: Vec<f32> = h.i((row, 0, ..)).unwrap().to_vec1().unwrap();
|
||||
let diff = max_abs_diff(&expected[row][t], &got);
|
||||
assert!(diff < 1e-4, "row {row} step {t} diverged: {diff}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Uniform-length batch: no padding → `batch_decode_mask` returns
|
||||
/// `None`, and unmasked lockstep decode still matches sequential.
|
||||
#[test]
|
||||
fn batched_decode_uniform_lengths_needs_no_mask() {
|
||||
use candle_core::IndexOp;
|
||||
let cfg = tiny_config();
|
||||
let mut model = tiny_model(&cfg);
|
||||
|
||||
let prompts: [&[u32]; 2] = [&[1, 2, 3], &[6, 5, 4]];
|
||||
let toks = [13u32, 17];
|
||||
|
||||
let mut expected = Vec::new();
|
||||
for (prompt, tok) in prompts.iter().zip(toks.iter()) {
|
||||
model.clear_kv_cache();
|
||||
forward_tokens(&mut model, prompt, 0);
|
||||
expected.push(forward_tokens(&mut model, &[*tok], prompt.len()));
|
||||
}
|
||||
|
||||
let mut snaps = Vec::new();
|
||||
for prompt in prompts.iter() {
|
||||
model.clear_kv_cache();
|
||||
forward_tokens(&mut model, prompt, 0);
|
||||
snaps.push(model.snapshot_kv_cache().expect("snapshot"));
|
||||
}
|
||||
let seqs: Vec<(&super::KvCacheSnapshot, usize)> = snaps
|
||||
.iter()
|
||||
.zip(prompts.iter())
|
||||
.map(|(s, p)| (s, p.len()))
|
||||
.collect();
|
||||
let batch = super::assemble_batch(&seqs).expect("assemble");
|
||||
let mask = model
|
||||
.batch_decode_mask(&batch.prefix_lens, batch.padded_len, batch.padded_len + 1)
|
||||
.expect("mask");
|
||||
assert!(mask.is_none(), "uniform lengths must not build a mask");
|
||||
model.restore_kv_cache(&batch.snapshot).expect("install");
|
||||
|
||||
let input = Tensor::from_vec(toks.to_vec(), (2, 1), &Device::Cpu).unwrap();
|
||||
let h = model
|
||||
.forward_batch_decode(&input, &[3, 3], None)
|
||||
.expect("step");
|
||||
for row in 0..2 {
|
||||
let got: Vec<f32> = h.i((row, 0, ..)).unwrap().to_vec1().unwrap();
|
||||
let diff = max_abs_diff(&expected[row], &got);
|
||||
assert!(diff < 1e-4, "row {row} diverged: {diff}");
|
||||
}
|
||||
}
|
||||
|
||||
/// Round-trip for batch membership changes (#98): decode two steps
|
||||
/// in a lockstep batch, extract each row back to a contiguous
|
||||
/// single-sequence snapshot, restore it alone, and continue
|
||||
/// decoding at B=1 — the continuation must match a pure-sequential
|
||||
/// run of the same tokens. This is the primitive a join/leave
|
||||
/// rebatch is built from.
|
||||
#[test]
|
||||
fn extract_row_continues_like_sequential() {
|
||||
use candle_core::IndexOp;
|
||||
let cfg = tiny_config();
|
||||
let mut model = tiny_model(&cfg);
|
||||
|
||||
let prompts: [&[u32]; 2] = [&[1, 2, 3], &[4, 5]];
|
||||
let toks: [&[u32]; 2] = [&[11, 12, 13, 14], &[9, 8, 7, 6]];
|
||||
let batched_steps = 2; // decoded in the batch
|
||||
let solo_steps = 2; // decoded after extraction, B=1
|
||||
|
||||
// Pure-sequential reference over all four steps.
|
||||
let mut expected: Vec<Vec<Vec<f32>>> = Vec::new();
|
||||
for (prompt, t) in prompts.iter().zip(toks.iter()) {
|
||||
model.clear_kv_cache();
|
||||
forward_tokens(&mut model, prompt, 0);
|
||||
let mut per_step = Vec::new();
|
||||
for (i, tok) in t.iter().enumerate() {
|
||||
per_step.push(forward_tokens(&mut model, &[*tok], prompt.len() + i));
|
||||
}
|
||||
expected.push(per_step);
|
||||
}
|
||||
|
||||
// Prefill + assemble + two lockstep batched steps.
|
||||
let mut snaps = Vec::new();
|
||||
for prompt in prompts.iter() {
|
||||
model.clear_kv_cache();
|
||||
forward_tokens(&mut model, prompt, 0);
|
||||
snaps.push(model.snapshot_kv_cache().expect("snapshot"));
|
||||
}
|
||||
let seqs: Vec<(&super::KvCacheSnapshot, usize)> = snaps
|
||||
.iter()
|
||||
.zip(prompts.iter())
|
||||
.map(|(s, p)| (s, p.len()))
|
||||
.collect();
|
||||
let batch = super::assemble_batch(&seqs).expect("assemble");
|
||||
model.restore_kv_cache(&batch.snapshot).expect("install");
|
||||
for t in 0..batched_steps {
|
||||
let step_toks: Vec<u32> = toks.iter().map(|s| s[t]).collect();
|
||||
let input = Tensor::from_vec(step_toks, (2, 1), &Device::Cpu).unwrap();
|
||||
let positions: Vec<usize> = prompts.iter().map(|p| p.len() + t).collect();
|
||||
let mask = model
|
||||
.batch_decode_mask(
|
||||
&batch.prefix_lens,
|
||||
batch.padded_len,
|
||||
batch.padded_len + t + 1,
|
||||
)
|
||||
.expect("mask");
|
||||
let h = model
|
||||
.forward_batch_decode(&input, &positions, mask.as_ref())
|
||||
.expect("batched step");
|
||||
for row in 0..2 {
|
||||
let got: Vec<f32> = h.i((row, 0, ..)).unwrap().to_vec1().unwrap();
|
||||
let diff = max_abs_diff(&expected[row][t], &got);
|
||||
assert!(diff < 1e-4, "batched row {row} step {t}: {diff}");
|
||||
}
|
||||
}
|
||||
|
||||
// Extract each row from the live batched state, restore it
|
||||
// alone, and continue at B=1.
|
||||
let live = model.snapshot_kv_cache().expect("snapshot live batch");
|
||||
for row in 0..2 {
|
||||
let solo = super::extract_row(
|
||||
&live,
|
||||
row,
|
||||
prompts[row].len(),
|
||||
batch.padded_len,
|
||||
batched_steps,
|
||||
)
|
||||
.expect("extract row");
|
||||
model.restore_kv_cache(&solo).expect("restore solo");
|
||||
for i in 0..solo_steps {
|
||||
let t = batched_steps + i;
|
||||
let got = forward_tokens(&mut model, &[toks[row][t]], prompts[row].len() + t);
|
||||
let diff = max_abs_diff(&expected[row][t], &got);
|
||||
assert!(diff < 1e-4, "solo row {row} step {t}: {diff}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Mask geometry: `-inf` exactly on `[prefix_len, padded_len)` per
|
||||
/// row, zero elsewhere (including the decode columns past
|
||||
/// `padded_len`).
|
||||
#[test]
|
||||
fn batch_decode_mask_covers_only_padding_gap() {
|
||||
let model = tiny_model(&tiny_config());
|
||||
let m = model
|
||||
.batch_decode_mask(&[3, 5], 5, 7)
|
||||
.unwrap()
|
||||
.expect("ragged → mask");
|
||||
assert_eq!(m.dims(), &[2, 1, 1, 7]);
|
||||
let flat: Vec<f32> = m.flatten_all().unwrap().to_vec1().unwrap();
|
||||
let (row0, row1) = flat.split_at(7);
|
||||
for (j, &v) in row0.iter().enumerate() {
|
||||
if (3..5).contains(&j) {
|
||||
assert_eq!(v, f32::NEG_INFINITY, "row0 col {j} must be masked");
|
||||
} else {
|
||||
assert_eq!(v, 0.0, "row0 col {j} must be open");
|
||||
}
|
||||
}
|
||||
assert!(row1.iter().all(|&v| v == 0.0), "unpadded row must be open");
|
||||
}
|
||||
|
||||
/// Restoring must fully replace the live state, not blend with it
|
||||
/// — a divergent continuation after restore equals the same
|
||||
/// continuation after a fresh prefill of the prefix.
|
||||
|
||||
@@ -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.
|
||||
///
|
||||
@@ -800,19 +850,28 @@ impl ModelArch {
|
||||
/// 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. Otherwise (non-reasoning
|
||||
/// model, thinking disabled, or a generation truncated mid-reasoning) 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.
|
||||
/// 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
|
||||
&& let Some(idx) = generated_ids.iter().rposition(|&t| t == pair.close_id)
|
||||
{
|
||||
return (&generated_ids[idx + 1..], (idx + 1) as 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)
|
||||
}
|
||||
@@ -894,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
|
||||
@@ -949,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
|
||||
@@ -1020,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,
|
||||
@@ -1061,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;
|
||||
@@ -1121,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
|
||||
@@ -1495,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,
|
||||
@@ -1554,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],
|
||||
@@ -1605,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,
|
||||
@@ -2007,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.
|
||||
@@ -2054,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,
|
||||
@@ -2302,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;
|
||||
@@ -2373,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,
|
||||
@@ -2425,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,
|
||||
@@ -2491,7 +2563,11 @@ impl CandleHarness {
|
||||
// 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());
|
||||
split_off_reasoning(
|
||||
&generated_ids,
|
||||
loaded.reasoning_tokens.as_ref(),
|
||||
prompt_opened_reasoning,
|
||||
);
|
||||
let completion_text = loaded
|
||||
.tokenizer
|
||||
.decode(content_ids, true)
|
||||
@@ -2831,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();
|
||||
@@ -2848,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,
|
||||
@@ -2913,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,
|
||||
@@ -3266,10 +3364,15 @@ impl Harness for CandleHarness {
|
||||
self.resolve_dense_files(spec, &source_id).await?;
|
||||
let meta = VisionMeta::from_config_path(&config_path);
|
||||
// Prefix snapshots (#11) exist only for the in-tree
|
||||
// qwen3_5 arch; the worker holds the ModelArch so
|
||||
// the async side decides from config.json instead.
|
||||
let snapshot_capable = config_model_type(&config_path).as_deref()
|
||||
== Some(super::arch::qwen3_5::MODEL_TYPE);
|
||||
// qwen3_5 arch — which serves BOTH its model_types
|
||||
// (qwen3_5 and the flat-config qwen3_next MoE family);
|
||||
// the worker holds the ModelArch so the async side
|
||||
// decides from config.json instead.
|
||||
let snapshot_capable = matches!(
|
||||
config_model_type(&config_path).as_deref(),
|
||||
Some(super::arch::qwen3_5::MODEL_TYPE)
|
||||
| Some(super::arch::qwen3_5::MODEL_TYPE_NEXT)
|
||||
);
|
||||
// Context-limit physics (#67): single-GPU → world_size 1.
|
||||
// `None` for non-qwen3_5 dense archs.
|
||||
let context_profile =
|
||||
@@ -3357,6 +3460,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,
|
||||
@@ -3364,10 +3504,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,
|
||||
@@ -3376,11 +3516,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!(
|
||||
@@ -3461,12 +3602,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");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -3621,7 +3785,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(),
|
||||
@@ -3637,10 +3801,11 @@ impl 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(
|
||||
config_model_type(&config_path).as_deref()
|
||||
== Some(super::arch::qwen3_5::MODEL_TYPE),
|
||||
),
|
||||
prefix_cache: self.new_prefix_cache(matches!(
|
||||
config_model_type(&config_path).as_deref(),
|
||||
Some(super::arch::qwen3_5::MODEL_TYPE)
|
||||
| Some(super::arch::qwen3_5::MODEL_TYPE_NEXT)
|
||||
)),
|
||||
// Context-limit physics (#67): per-rank KV cost sharded across
|
||||
// the TP world. `None` for non-qwen3_5 dense archs.
|
||||
context_profile: super::context_limit::profile_from_qwen3_5_config(
|
||||
@@ -3651,7 +3816,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,
|
||||
@@ -3971,6 +4162,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 {
|
||||
@@ -4788,8 +5011,11 @@ async fn chat_completion_tp_inner(
|
||||
|
||||
// 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());
|
||||
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(content_ids, true)
|
||||
@@ -4852,8 +5078,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;
|
||||
}
|
||||
@@ -4889,7 +5118,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,
|
||||
@@ -4914,7 +5143,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 {
|
||||
@@ -4929,7 +5161,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,
|
||||
@@ -4946,7 +5178,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,
|
||||
@@ -4984,7 +5216,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}}}}}`).
|
||||
@@ -5028,7 +5261,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,
|
||||
@@ -5635,7 +5868,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)
|
||||
@@ -5659,8 +5895,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>,
|
||||
@@ -5709,8 +5944,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>,
|
||||
@@ -5773,7 +6007,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],
|
||||
@@ -5830,7 +6064,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>,
|
||||
@@ -6613,16 +6847,16 @@ mod tests {
|
||||
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()));
|
||||
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 / model never closed the span: return as-is.
|
||||
// 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()));
|
||||
let (content, reasoning) = split_off_reasoning(&ids, Some(&think_pair()), false);
|
||||
assert_eq!(content, &ids);
|
||||
assert_eq!(reasoning, 0);
|
||||
}
|
||||
@@ -6630,7 +6864,7 @@ mod tests {
|
||||
#[test]
|
||||
fn split_off_reasoning_no_marker_pair_is_noop() {
|
||||
let ids = [1, 2, 3];
|
||||
let (content, reasoning) = split_off_reasoning(&ids, None);
|
||||
let (content, reasoning) = split_off_reasoning(&ids, None, true);
|
||||
assert_eq!(content, &ids);
|
||||
assert_eq!(reasoning, 0);
|
||||
}
|
||||
@@ -6639,7 +6873,18 @@ mod tests {
|
||||
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()));
|
||||
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);
|
||||
}
|
||||
@@ -6649,7 +6894,7 @@ mod tests {
|
||||
// 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()));
|
||||
let (content, reasoning) = split_off_reasoning(&ids, Some(&think_pair()), true);
|
||||
assert_eq!(content, &[42]);
|
||||
assert_eq!(reasoning, 3);
|
||||
}
|
||||
@@ -6729,6 +6974,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 }"#;
|
||||
|
||||
@@ -145,10 +145,15 @@ pub fn profile_from_qwen3_5_config(config_path: &Path, world_size: u32) -> Optio
|
||||
.get("model_type")?
|
||||
.as_str()?
|
||||
.to_owned();
|
||||
if model_type != super::arch::qwen3_5::MODEL_TYPE {
|
||||
if model_type != super::arch::qwen3_5::MODEL_TYPE
|
||||
&& model_type != super::arch::qwen3_5::MODEL_TYPE_NEXT
|
||||
{
|
||||
return None;
|
||||
}
|
||||
let cfg: super::arch::qwen3_5::Config = serde_json::from_str(&text).ok()?;
|
||||
// `from_config_json` normalises both layouts (nested qwen3_5, flat
|
||||
// qwen3_next) — a plain serde parse would reject the flat family,
|
||||
// which is how Coder-Next went without an advertised limit (#126).
|
||||
let cfg = super::arch::qwen3_5::Config::from_config_json(&text).ok()?;
|
||||
let tc = &cfg.text_config;
|
||||
let n_full_attn_layers = {
|
||||
let counted = tc
|
||||
@@ -257,6 +262,22 @@ pub fn derive_limit(
|
||||
}
|
||||
context = round_down(context, CONTEXT_GRANULARITY);
|
||||
|
||||
// Observability (#126): every input and intermediate term, so a
|
||||
// surprising advertised limit is diagnosable from the journal
|
||||
// instead of re-deriving by hand. DEBUG — this runs on every
|
||||
// `GET /models` poll (a few lines per poll cycle).
|
||||
tracing::debug!(
|
||||
max_pos = profile.max_position_embeddings,
|
||||
kv_bytes_per_token_per_card = profile.kv_bytes_per_token_per_card,
|
||||
free_tightest_mb,
|
||||
prefill_tok_per_sec = tok_per_sec,
|
||||
vram_ceiling,
|
||||
throughput_ceiling,
|
||||
?hard_ceiling,
|
||||
context,
|
||||
"derive_limit"
|
||||
);
|
||||
|
||||
let input = context.saturating_sub(output);
|
||||
ModelLimit {
|
||||
context,
|
||||
|
||||
@@ -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).
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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");
|
||||
|
||||
1131
crates/neuron/src/harness/engine.rs
Normal file
1131
crates/neuron/src/harness/engine.rs
Normal file
File diff suppressed because it is too large
Load Diff
@@ -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;
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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 },
|
||||
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
56
crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/config.json
vendored
Normal file
56
crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/config.json
vendored
Normal file
@@ -0,0 +1,56 @@
|
||||
{
|
||||
"architectures": [
|
||||
"Qwen3NextForCausalLM"
|
||||
],
|
||||
"attention_bias": false,
|
||||
"attention_dropout": 0.0,
|
||||
"bos_token_id": null,
|
||||
"decoder_sparse_step": 1,
|
||||
"dtype": "float32",
|
||||
"eos_token_id": null,
|
||||
"head_dim": 32,
|
||||
"hidden_act": "silu",
|
||||
"hidden_size": 64,
|
||||
"initializer_range": 0.02,
|
||||
"intermediate_size": 128,
|
||||
"layer_types": [
|
||||
"linear_attention",
|
||||
"linear_attention",
|
||||
"linear_attention",
|
||||
"full_attention",
|
||||
"linear_attention",
|
||||
"linear_attention",
|
||||
"linear_attention",
|
||||
"full_attention"
|
||||
],
|
||||
"linear_conv_kernel_dim": 4,
|
||||
"linear_key_head_dim": 16,
|
||||
"linear_num_key_heads": 2,
|
||||
"linear_num_value_heads": 4,
|
||||
"linear_value_head_dim": 16,
|
||||
"max_position_embeddings": 512,
|
||||
"mlp_only_layers": [],
|
||||
"model_type": "qwen3_next",
|
||||
"moe_intermediate_size": 32,
|
||||
"norm_topk_prob": true,
|
||||
"num_attention_heads": 4,
|
||||
"num_experts": 16,
|
||||
"num_experts_per_tok": 4,
|
||||
"num_hidden_layers": 8,
|
||||
"num_key_value_heads": 2,
|
||||
"output_router_logits": false,
|
||||
"pad_token_id": null,
|
||||
"partial_rotary_factor": 0.25,
|
||||
"rms_norm_eps": 1e-06,
|
||||
"rope_parameters": {
|
||||
"partial_rotary_factor": 0.25,
|
||||
"rope_theta": 10000000,
|
||||
"rope_type": "default"
|
||||
},
|
||||
"router_aux_loss_coef": 0.001,
|
||||
"shared_expert_intermediate_size": 32,
|
||||
"tie_word_embeddings": false,
|
||||
"transformers_version": "5.9.0",
|
||||
"use_cache": true,
|
||||
"vocab_size": 512
|
||||
}
|
||||
7
crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/generation_config.json
vendored
Normal file
7
crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/generation_config.json
vendored
Normal file
@@ -0,0 +1,7 @@
|
||||
{
|
||||
"_from_model_config": true,
|
||||
"output_attentions": false,
|
||||
"output_hidden_states": false,
|
||||
"transformers_version": "5.9.0",
|
||||
"use_cache": true
|
||||
}
|
||||
BIN
crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/logits.f32
vendored
Normal file
BIN
crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/logits.f32
vendored
Normal file
Binary file not shown.
114
crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/manifest.json
vendored
Normal file
114
crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/manifest.json
vendored
Normal file
@@ -0,0 +1,114 @@
|
||||
{
|
||||
"case": "qwen3_next-tiny",
|
||||
"seed": 92,
|
||||
"token_ids": [
|
||||
3,
|
||||
10,
|
||||
17,
|
||||
24,
|
||||
31,
|
||||
38,
|
||||
45,
|
||||
52,
|
||||
59,
|
||||
66,
|
||||
73,
|
||||
80,
|
||||
87,
|
||||
94,
|
||||
101,
|
||||
108,
|
||||
115,
|
||||
122,
|
||||
129,
|
||||
136,
|
||||
143,
|
||||
150,
|
||||
157,
|
||||
164,
|
||||
171,
|
||||
178,
|
||||
185,
|
||||
192,
|
||||
199,
|
||||
206,
|
||||
213,
|
||||
220,
|
||||
227,
|
||||
234,
|
||||
241,
|
||||
248,
|
||||
255,
|
||||
262,
|
||||
269,
|
||||
276,
|
||||
283,
|
||||
290,
|
||||
297,
|
||||
304,
|
||||
311,
|
||||
318,
|
||||
325,
|
||||
332,
|
||||
339,
|
||||
346,
|
||||
353,
|
||||
360,
|
||||
367,
|
||||
374,
|
||||
381,
|
||||
388,
|
||||
395,
|
||||
402,
|
||||
409,
|
||||
416,
|
||||
423,
|
||||
430,
|
||||
437,
|
||||
444,
|
||||
451,
|
||||
458,
|
||||
465,
|
||||
472,
|
||||
479,
|
||||
486,
|
||||
493,
|
||||
500,
|
||||
507,
|
||||
2,
|
||||
9,
|
||||
16,
|
||||
23,
|
||||
30,
|
||||
37,
|
||||
44,
|
||||
51,
|
||||
58,
|
||||
65,
|
||||
72,
|
||||
79,
|
||||
86,
|
||||
93,
|
||||
100,
|
||||
107,
|
||||
114,
|
||||
121,
|
||||
128,
|
||||
135,
|
||||
142,
|
||||
149,
|
||||
156
|
||||
],
|
||||
"files": {
|
||||
"logits": {
|
||||
"file": "logits.f32",
|
||||
"shape": [
|
||||
512
|
||||
]
|
||||
}
|
||||
},
|
||||
"versions": {
|
||||
"transformers": "5.9.0",
|
||||
"torch": "2.9.1+cu128"
|
||||
}
|
||||
}
|
||||
BIN
crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/model.safetensors
vendored
Normal file
BIN
crates/neuron/tests/fixtures/numerical/qwen3_next-tiny/model.safetensors
vendored
Normal file
Binary file not shown.
122
crates/neuron/tests/qwen3_next_parity.rs
Normal file
122
crates/neuron/tests/qwen3_next_parity.rs
Normal file
@@ -0,0 +1,122 @@
|
||||
//! Numerical parity for the qwen3_next path (#92) against the HF
|
||||
//! transformers reference, via the tiny self-contained fixture
|
||||
//! generated by `script/dump_qwen3_next_tiny.py`.
|
||||
//!
|
||||
//! The fixture directory carries the WHOLE checkpoint (tiny
|
||||
//! random-weight `Qwen3NextForCausalLM`: config.json +
|
||||
//! model.safetensors, a few hundred KB) plus the reference
|
||||
//! final-position logits, so this test needs no snapshot, no env var,
|
||||
//! and runs in CI. It pins: flat-config normalisation, the `model.*`
|
||||
//! weight prefix, the fused `in_proj_qkvz`/`in_proj_ba` de-interleave,
|
||||
//! hybrid full/linear layer interleaving, and the MoE block (routing,
|
||||
//! per-expert SwiGLU, shared expert + sigmoid gate).
|
||||
//!
|
||||
//! Self-skips (with a loud eprintln) while the fixture has not been
|
||||
//! generated yet — regeneration instructions in the script docstring.
|
||||
|
||||
use candle_core::{DType, Device, Tensor};
|
||||
use serde::Deserialize;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Manifest {
|
||||
token_ids: Vec<u32>,
|
||||
files: std::collections::HashMap<String, FileEntry>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct FileEntry {
|
||||
file: String,
|
||||
shape: Vec<usize>,
|
||||
}
|
||||
|
||||
fn fixture_dir() -> PathBuf {
|
||||
Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/numerical/qwen3_next-tiny")
|
||||
}
|
||||
|
||||
fn read_f32(path: &Path) -> Vec<f32> {
|
||||
let bytes = std::fs::read(path).unwrap_or_else(|e| panic!("read {path:?}: {e}"));
|
||||
bytes
|
||||
.chunks_exact(4)
|
||||
.map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tiny_qwen3_next_logits_match_hf_reference() {
|
||||
let dir = fixture_dir();
|
||||
let manifest_path = dir.join("manifest.json");
|
||||
if !manifest_path.exists() {
|
||||
eprintln!(
|
||||
"SKIP qwen3_next parity: fixture not generated yet — run \
|
||||
script/dump_qwen3_next_tiny.py --out {} on a host with \
|
||||
torch + transformers>=4.57",
|
||||
dir.display()
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
let manifest: Manifest =
|
||||
serde_json::from_str(&std::fs::read_to_string(&manifest_path).expect("read manifest"))
|
||||
.expect("parse manifest");
|
||||
let logits_entry = &manifest.files["logits"];
|
||||
let reference = read_f32(&dir.join(&logits_entry.file));
|
||||
assert_eq!(
|
||||
reference.len(),
|
||||
logits_entry.shape.iter().product::<usize>()
|
||||
);
|
||||
|
||||
let config_json = std::fs::read_to_string(dir.join("config.json")).expect("read config.json");
|
||||
let cfg = neuron::harness::arch::qwen3_5::Config::from_config_json(&config_json)
|
||||
.expect("normalise qwen3_next config");
|
||||
|
||||
let dev = Device::Cpu;
|
||||
let st_path = dir.join("model.safetensors");
|
||||
// SAFETY: mmap of committed fixture files; nothing mutates them.
|
||||
let vb = unsafe {
|
||||
candle_nn::var_builder::ShardedSafeTensors::var_builder(
|
||||
std::slice::from_ref(&st_path),
|
||||
DType::F32,
|
||||
&dev,
|
||||
)
|
||||
.expect("build ShardedVarBuilder over fixture checkpoint")
|
||||
};
|
||||
let mut model = neuron::harness::arch::qwen3_5::Qwen3_5ForCausalLM::new(cfg, vb)
|
||||
.expect("load tiny qwen3_next checkpoint through neuron");
|
||||
|
||||
let input = Tensor::new(manifest.token_ids.as_slice(), &dev)
|
||||
.unwrap()
|
||||
.unsqueeze(0)
|
||||
.unwrap();
|
||||
let logits = model.forward(&input, 0).expect("forward");
|
||||
let got: Vec<f32> = logits.flatten_all().unwrap().to_vec1().unwrap();
|
||||
assert_eq!(got.len(), reference.len());
|
||||
|
||||
// f32-vs-f32 through an 8-layer doll-house model: agreement should
|
||||
// be tight (the qwen3_5 text fixtures observe max_abs ≈ 0.000,
|
||||
// cosine ≈ 1.0). Thresholds sit far above rounding noise and far
|
||||
// below any real wiring bug (a swapped de-interleave region, a
|
||||
// topk-before-softmax, a missing shared-expert gate all blow past
|
||||
// them instantly).
|
||||
let max_abs = got
|
||||
.iter()
|
||||
.zip(&reference)
|
||||
.map(|(a, b)| (a - b).abs())
|
||||
.fold(0f32, f32::max);
|
||||
let dot: f64 = got
|
||||
.iter()
|
||||
.zip(&reference)
|
||||
.map(|(a, b)| (*a as f64) * (*b as f64))
|
||||
.sum();
|
||||
let na: f64 = got.iter().map(|a| (*a as f64).powi(2)).sum::<f64>().sqrt();
|
||||
let nb: f64 = reference
|
||||
.iter()
|
||||
.map(|b| (*b as f64).powi(2))
|
||||
.sum::<f64>()
|
||||
.sqrt();
|
||||
let cosine = dot / (na * nb);
|
||||
|
||||
eprintln!("qwen3_next parity: max_abs={max_abs:.6} cosine={cosine:.8}");
|
||||
assert!(max_abs < 1e-3, "max abs diff {max_abs} exceeds 1e-3");
|
||||
assert!(cosine > 0.9999, "cosine {cosine} below 0.9999");
|
||||
}
|
||||
@@ -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`).
|
||||
|
||||
156
script/dump_qwen3_next_tiny.py
Normal file
156
script/dump_qwen3_next_tiny.py
Normal 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()
|
||||
Reference in New Issue
Block a user