Compare commits
32 Commits
577781de8d
...
chore/cuda
| Author | SHA1 | Date | |
|---|---|---|---|
|
60f5598542
|
|||
|
7945240646
|
|||
|
0c74d89d15
|
|||
|
c94a2ae755
|
|||
|
99920dd322
|
|||
|
c4f239ceb9
|
|||
|
ac445c1569
|
|||
|
abc6e605b8
|
|||
|
4f2957af9e
|
|||
|
75cd088b61
|
|||
|
d311c8ca7a
|
|||
|
c97a8654f5
|
|||
|
dc048ffcc9
|
|||
|
7ebcfba5ca
|
|||
|
825bf4e905
|
|||
|
4c12c7e2f0
|
|||
|
ba1b5ba408
|
|||
|
5731f4c318
|
|||
|
fa013505d1
|
|||
|
c8bcaabc38
|
|||
|
7ad56c6a86
|
|||
|
1b0e36c119
|
|||
|
ed2d09864e
|
|||
|
4994b94c84
|
|||
|
9a24b05866
|
|||
|
7bb033b4ed
|
|||
|
f8c0da0ebf
|
|||
|
dd592d918d
|
|||
|
766c20ba47
|
|||
|
4972c7d1e7
|
|||
|
a26bb9f04b
|
|||
|
ea1fdf8aa6
|
@@ -70,6 +70,16 @@ jobs:
|
||||
sudo /usr/bin/systemctl daemon-reload
|
||||
sudo /usr/bin/systemctl start cortex.service'
|
||||
|
||||
# Wait for the service to either come up or wedge, then capture
|
||||
# the latest-invocation journal. Runs even on prior failure so a
|
||||
# failed start step still leaves a usable record in the deploy log.
|
||||
- name: Capture cortex.service startup journal
|
||||
if: always()
|
||||
run: |
|
||||
sleep 10
|
||||
ssh gitea_ci@hanzalova.internal \
|
||||
'journalctl --unit cortex.service -I --no-pager'
|
||||
|
||||
deploy-neurons:
|
||||
needs: [deploy-cortex]
|
||||
runs-on: fedora-43
|
||||
@@ -124,3 +134,13 @@ jobs:
|
||||
ssh gitea_ci@${{ matrix.host }} '
|
||||
sudo /usr/bin/systemctl daemon-reload
|
||||
sudo /usr/bin/systemctl start neuron.service'
|
||||
|
||||
# Wait for the service to either come up or wedge, then capture
|
||||
# the latest-invocation journal. Runs even on prior failure so a
|
||||
# failed start step still leaves a usable record in the deploy log.
|
||||
- name: Capture neuron.service startup journal
|
||||
if: always()
|
||||
run: |
|
||||
sleep 10
|
||||
ssh gitea_ci@${{ matrix.host }} \
|
||||
'journalctl --unit neuron.service -I --no-pager'
|
||||
|
||||
1
.gitignore
vendored
1
.gitignore
vendored
@@ -7,3 +7,4 @@ cortex.toml
|
||||
models.toml
|
||||
doc/plan/*
|
||||
/target-cuda/
|
||||
.claude/
|
||||
|
||||
14
Cargo.lock
generated
14
Cargo.lock
generated
@@ -905,8 +905,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "cudarc"
|
||||
version = "0.19.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1cea5f10a99e025c1b44ae2354c2d8326b25ddbd0baf76bde8e55cfd4018a2cc"
|
||||
source = "git+https://github.com/grenade/cudarc?rev=63327a256059f8252641ae46c6bb9eefe707f382#63327a256059f8252641ae46c6bb9eefe707f382"
|
||||
dependencies = [
|
||||
"float8",
|
||||
"half",
|
||||
@@ -2508,6 +2507,16 @@ dependencies = [
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "minijinja-contrib"
|
||||
version = "2.20.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "99df5123c54391e2a228014c1dbbd85a3dab08a25e776c810526f2f47542b3de"
|
||||
dependencies = [
|
||||
"minijinja",
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "minimal-lexical"
|
||||
version = "0.2.1"
|
||||
@@ -2605,6 +2614,7 @@ dependencies = [
|
||||
"hf-hub",
|
||||
"image",
|
||||
"minijinja",
|
||||
"minijinja-contrib",
|
||||
"reqwest",
|
||||
"safetensors 0.7.0",
|
||||
"serde",
|
||||
|
||||
@@ -61,3 +61,12 @@ eventsource-stream = "0.2"
|
||||
# workspace crates
|
||||
cortex-core = { path = "crates/cortex-core" }
|
||||
cortex-gateway = { path = "crates/cortex-gateway" }
|
||||
|
||||
# Patched cudarc (affects neuron's 0.19.x only; candle's 0.17.x is
|
||||
# untouched since the fork is 0.19.7 and doesn't satisfy a 0.17 req). Adds
|
||||
# Comm::abort / get_async_error / raw comm() — needed for #17 Stage 2 TP
|
||||
# hang-recovery (abort a wedged collective from another thread, then
|
||||
# rebuild the comm). Pinned to a fork revision pending upstream review
|
||||
# (grenade/cudarc @ nccl-comm-abort).
|
||||
[patch.crates-io]
|
||||
cudarc = { git = "https://github.com/grenade/cudarc", rev = "63327a256059f8252641ae46c6bb9eefe707f382" }
|
||||
|
||||
@@ -1,30 +0,0 @@
|
||||
# Helexa fleet manifest.
|
||||
#
|
||||
# Drives rolling deploys via script/deploy.sh and serves as the source
|
||||
# of truth for which hosts run cortex vs neuron, and which CUDA
|
||||
# compute-capability flavour each neuron host needs.
|
||||
#
|
||||
# Flavour ↔ NVIDIA generation ↔ compute cap:
|
||||
# ampere sm_86 (RTX 30 series — e.g. 3060)
|
||||
# ada sm_89 (RTX 40 series — e.g. 4090)
|
||||
# blackwell sm_120 (RTX 50 series — e.g. 5090)
|
||||
#
|
||||
# The flavour determines which RPM is installed on a given neuron host:
|
||||
# helexa-neuron-<flavour>. Only one flavour may be installed at a time
|
||||
# (the packages Conflict: with each other).
|
||||
|
||||
cortex:
|
||||
host: hanzalova.internal
|
||||
|
||||
neurons:
|
||||
- host: beast.hanzalova.internal
|
||||
flavour: blackwell
|
||||
gpu: "2x RTX 5090"
|
||||
|
||||
- host: benjy.hanzalova.internal
|
||||
flavour: ada
|
||||
gpu: "RTX 4090"
|
||||
|
||||
- host: quadbrat.hanzalova.internal
|
||||
flavour: ampere
|
||||
gpu: "RTX 3060"
|
||||
@@ -5,9 +5,9 @@
|
||||
# invocation: `validate-neuron.sh beast.hanzalova.internal
|
||||
# Qwen/Qwen3.6-27B q5k 2`.
|
||||
#
|
||||
# Synced by script/deploy.sh from asset/neuron/<short-host>.toml. Edits
|
||||
# take effect on the next deploy.sh run (which stops + restarts the
|
||||
# service so default_models is re-read at activation).
|
||||
# Synced to /etc/neuron/neuron.toml by script/infra-setup.sh. Edits
|
||||
# take effect after the next deploy workflow run restarts the service
|
||||
# (default_models is read at activation).
|
||||
|
||||
port = 13131
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
# Qwen3-8B (bf16, ~18 GB), leaving ~6 GB for KV cache + activations on
|
||||
# moderate-length contexts.
|
||||
#
|
||||
# Synced by script/deploy.sh from asset/neuron/<short-host>.toml.
|
||||
# Synced to /etc/neuron/neuron.toml by script/infra-setup.sh.
|
||||
|
||||
port = 13131
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
# (bf16, ~4 GB), leaving ~7 GB for KV cache so long contexts on a small
|
||||
# model still have plenty of room.
|
||||
#
|
||||
# Synced by script/deploy.sh from asset/neuron/<short-host>.toml.
|
||||
# Synced to /etc/neuron/neuron.toml by script/infra-setup.sh.
|
||||
|
||||
port = 13131
|
||||
|
||||
|
||||
@@ -37,6 +37,12 @@ pub struct ModelEntry {
|
||||
pub last_accessed: Option<DateTime<Utc>>,
|
||||
/// Estimated VRAM usage in MB when loaded.
|
||||
pub vram_estimate_mb: Option<u64>,
|
||||
/// Modalities the loaded model advertises (e.g. `["text", "vision"]`),
|
||||
/// copied verbatim from the neuron's `ModelInfo.capabilities` at poll
|
||||
/// time. Empty when the neuron reports none. `#[serde(default)]` keeps
|
||||
/// older persisted/serialised entries deserialisable.
|
||||
#[serde(default)]
|
||||
pub capabilities: Vec<String>,
|
||||
}
|
||||
|
||||
/// Model lifecycle status.
|
||||
@@ -85,6 +91,12 @@ pub struct CortexModelEntry {
|
||||
/// disjoint from) `feasible_on` depending on whether the catalogue
|
||||
/// covers this model.
|
||||
pub locations: Vec<ModelLocation>,
|
||||
/// Union of the modalities advertised by every neuron that has this
|
||||
/// model loaded (e.g. `["text", "vision"]`). Empty for catalogue-only
|
||||
/// entries with no loaded location — the catalogue profile doesn't
|
||||
/// declare capabilities yet (tracked separately from C3).
|
||||
#[serde(default)]
|
||||
pub capabilities: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
|
||||
@@ -414,6 +414,9 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
||||
loaded: false,
|
||||
feasible_on,
|
||||
locations: Vec::new(),
|
||||
// Catalogue profiles don't declare capabilities yet;
|
||||
// the union is filled in Pass 2 from loaded locations.
|
||||
capabilities: Vec::new(),
|
||||
},
|
||||
);
|
||||
}
|
||||
@@ -438,6 +441,14 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
||||
if was_loaded {
|
||||
e.loaded = true;
|
||||
}
|
||||
// Union the per-node capabilities so a model loaded
|
||||
// on several neurons reports every modality any of
|
||||
// them advertises.
|
||||
for cap in &entry.capabilities {
|
||||
if !e.capabilities.contains(cap) {
|
||||
e.capabilities.push(cap.clone());
|
||||
}
|
||||
}
|
||||
})
|
||||
.or_insert_with(|| CortexModelEntry {
|
||||
id: model_id.clone(),
|
||||
@@ -449,6 +460,7 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
||||
// feasibility; leave empty.
|
||||
feasible_on: Vec::new(),
|
||||
locations: vec![location],
|
||||
capabilities: entry.capabilities.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -498,6 +510,9 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
||||
loaded: false,
|
||||
feasible_on: Vec::new(),
|
||||
locations: vec![location],
|
||||
// A model that's only mid-prewarm has no loaded
|
||||
// location to read capabilities from yet.
|
||||
capabilities: Vec::new(),
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -527,6 +542,7 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
||||
loaded: target_entry.loaded,
|
||||
feasible_on: target_entry.feasible_on,
|
||||
locations: target_entry.locations,
|
||||
capabilities: target_entry.capabilities,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
@@ -107,12 +107,14 @@ async fn poll_neuron(fleet: &CortexState, name: &str, endpoint: &str) {
|
||||
.and_modify(|e| {
|
||||
e.status = status;
|
||||
e.vram_estimate_mb = upstream.vram_used_mb;
|
||||
e.capabilities = upstream.capabilities.clone();
|
||||
})
|
||||
.or_insert_with(|| ModelEntry {
|
||||
id: upstream.id.clone(),
|
||||
status,
|
||||
last_accessed: None,
|
||||
vram_estimate_mb: upstream.vram_used_mb,
|
||||
capabilities: upstream.capabilities.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
@@ -244,6 +244,7 @@ async fn cold_load(
|
||||
status: ModelStatus::Loaded,
|
||||
last_accessed: Some(chrono::Utc::now()),
|
||||
vram_estimate_mb: profile.vram_mb,
|
||||
capabilities: Vec::new(),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
@@ -74,6 +74,7 @@ async fn test_alias_resolves_in_chat_completions() {
|
||||
status: ModelStatus::Loaded,
|
||||
last_accessed: None,
|
||||
vram_estimate_mb: None,
|
||||
capabilities: Vec::new(),
|
||||
},
|
||||
);
|
||||
}
|
||||
@@ -154,6 +155,7 @@ async fn test_aliases_surface_in_v1_models() {
|
||||
status: ModelStatus::Loaded,
|
||||
last_accessed: None,
|
||||
vram_estimate_mb: Some(2000),
|
||||
capabilities: Vec::new(),
|
||||
},
|
||||
);
|
||||
}
|
||||
@@ -235,6 +237,7 @@ async fn test_alias_falls_through_for_unmapped_model() {
|
||||
status: ModelStatus::Loaded,
|
||||
last_accessed: None,
|
||||
vram_estimate_mb: None,
|
||||
capabilities: Vec::new(),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
@@ -305,6 +305,7 @@ pub async fn spawn_gateway_with_state(mock_url: &str) -> (Arc<CortexState>, Stri
|
||||
status: ModelStatus::Loaded,
|
||||
last_accessed: None,
|
||||
vram_estimate_mb: Some(8000),
|
||||
capabilities: Vec::new(),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
@@ -91,6 +91,7 @@ async fn test_evict_lru_model() {
|
||||
status: ModelStatus::Loaded,
|
||||
last_accessed: Some(Utc::now() - chrono::Duration::hours(2)),
|
||||
vram_estimate_mb: Some(8000),
|
||||
capabilities: Vec::new(),
|
||||
},
|
||||
);
|
||||
node.models.insert(
|
||||
@@ -100,6 +101,7 @@ async fn test_evict_lru_model() {
|
||||
status: ModelStatus::Loaded,
|
||||
last_accessed: Some(Utc::now()),
|
||||
vram_estimate_mb: Some(8000),
|
||||
capabilities: Vec::new(),
|
||||
},
|
||||
);
|
||||
}
|
||||
@@ -163,6 +165,7 @@ async fn test_eviction_increments_lifecycle_cycles() {
|
||||
status: ModelStatus::Loaded,
|
||||
last_accessed: None,
|
||||
vram_estimate_mb: None,
|
||||
capabilities: Vec::new(),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
@@ -118,6 +118,87 @@ async fn test_poller_updates_gateway_models_endpoint() {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_models_endpoint_unions_capabilities_across_nodes() {
|
||||
// C3: two neurons each have the same model loaded but advertise
|
||||
// different capability sets. The gateway's /v1/models must report
|
||||
// the union — a model loaded text-only on one node and
|
||||
// text+vision on another is vision-capable to the fleet.
|
||||
let node_a = common::spawn_mock_neuron_with_models(json!([
|
||||
{"id": "shared-model", "harness": "candle", "status": "loaded", "devices": [0], "vram_used_mb": null, "capabilities": ["text"]}
|
||||
]))
|
||||
.await;
|
||||
let node_b = common::spawn_mock_neuron_with_models(json!([
|
||||
{"id": "shared-model", "harness": "candle", "status": "loaded", "devices": [1], "vram_used_mb": null, "capabilities": ["text", "vision"]}
|
||||
]))
|
||||
.await;
|
||||
|
||||
let config = GatewayConfig {
|
||||
gateway: GatewaySettings {
|
||||
listen: "127.0.0.1:0".into(),
|
||||
metrics_listen: "127.0.0.1:0".into(),
|
||||
},
|
||||
eviction: EvictionSettings {
|
||||
strategy: EvictionStrategy::Lru,
|
||||
defrag_after_cycles: 0,
|
||||
},
|
||||
neurons: vec![
|
||||
NeuronEndpoint {
|
||||
name: "node-a".into(),
|
||||
endpoint: node_a,
|
||||
},
|
||||
NeuronEndpoint {
|
||||
name: "node-b".into(),
|
||||
endpoint: node_b,
|
||||
},
|
||||
],
|
||||
models_config: "/dev/null".into(),
|
||||
};
|
||||
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
cortex_gateway::poller::poll_once(&fleet).await;
|
||||
|
||||
let app = cortex_gateway::build_app(Arc::clone(&fleet));
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let body: serde_json::Value = client
|
||||
.get(format!("http://{addr}/v1/models"))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed")
|
||||
.json()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let model = body["data"]
|
||||
.as_array()
|
||||
.expect("data array")
|
||||
.iter()
|
||||
.find(|m| m["id"] == "shared-model")
|
||||
.expect("shared-model should be present");
|
||||
|
||||
let caps: Vec<&str> = model["capabilities"]
|
||||
.as_array()
|
||||
.expect("capabilities array")
|
||||
.iter()
|
||||
.filter_map(|c| c.as_str())
|
||||
.collect();
|
||||
assert!(caps.contains(&"text"), "union must include text: {caps:?}");
|
||||
assert!(
|
||||
caps.contains(&"vision"),
|
||||
"union must include vision: {caps:?}"
|
||||
);
|
||||
assert_eq!(caps.len(), 2, "union must not duplicate text: {caps:?}");
|
||||
|
||||
// Both nodes hold the model, so two locations regardless of caps.
|
||||
assert_eq!(model["locations"].as_array().unwrap().len(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_poller_marks_unreachable_node_unhealthy() {
|
||||
let config = GatewayConfig {
|
||||
@@ -216,6 +297,7 @@ async fn test_poller_removes_stale_models() {
|
||||
status: ModelStatus::Loaded,
|
||||
last_accessed: None,
|
||||
vram_estimate_mb: None,
|
||||
capabilities: Vec::new(),
|
||||
},
|
||||
);
|
||||
node.models.insert(
|
||||
@@ -225,6 +307,7 @@ async fn test_poller_removes_stale_models() {
|
||||
status: ModelStatus::Loaded,
|
||||
last_accessed: None,
|
||||
vram_estimate_mb: None,
|
||||
capabilities: Vec::new(),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
@@ -76,15 +76,19 @@ cudarc = { version = "0.19", optional = true, default-features = false, features
|
||||
half = { version = "2.5", optional = true }
|
||||
tokenizers = { version = "0.22", default-features = false, features = ["onig"] }
|
||||
hf-hub = { version = "0.4", features = ["tokio"] }
|
||||
# Jinja-compatible template renderer for the model's
|
||||
# `tokenizer_config.json::chat_template`. Hugging Face's chat
|
||||
# templates use a strict subset of Jinja2 that minijinja supports
|
||||
# out of the box. ~80KB compiled; pure Rust, no async surface.
|
||||
# Features: `builtins` for the `is defined` / `default` filters HF
|
||||
# templates use; `json` for `tojson` (some Qwen3 templates emit
|
||||
# tool definitions via tojson); `serde` so we can hand it a
|
||||
# serde_json::Value as the context.
|
||||
# Jinja-compatible template renderer for the model's chat template
|
||||
# (standalone `chat_template.jinja` or `tokenizer_config.json::chat_template`).
|
||||
# Hugging Face's chat templates lean on Python string semantics; we
|
||||
# bridge them with `minijinja-contrib`'s `pycompat` callback (str
|
||||
# methods like `startswith`/`split`/`strip`) plus a `raise_exception`
|
||||
# global. Features: `builtins` for `is defined` / `default`; `json`
|
||||
# for `tojson`; `serde` so we can hand it a serde_json::Value context.
|
||||
minijinja = { version = "2", features = ["builtins", "json", "serde"] }
|
||||
# Python-compatibility shim: the Qwen3-VL / Qwen3.6 template uses
|
||||
# `content.startswith(...)`, `.endswith(...)`, `.split(...)`,
|
||||
# `.rstrip(...)`, `.lstrip(...)` — Python str methods minijinja doesn't
|
||||
# implement natively. `pycompat::unknown_method_callback` supplies them.
|
||||
minijinja-contrib = { version = "2", features = ["pycompat"] }
|
||||
# Direct dep on `safetensors` (re-exported by candle but its `TensorView`
|
||||
# / `slice::IndexOp` types are public-but-not-re-exported). Used by the
|
||||
# tp `fused_load` module to read per-rank slices of fused QKV tensors
|
||||
|
||||
@@ -93,12 +93,13 @@ impl Qwen3_5DecoderLayer {
|
||||
&mut self,
|
||||
x: &Tensor,
|
||||
attn_mask: Option<&Tensor>,
|
||||
offset: usize,
|
||||
cos: &Tensor,
|
||||
sin: &Tensor,
|
||||
) -> candle_core::Result<Tensor> {
|
||||
let h = self.input_layernorm.forward(x)?;
|
||||
let attn_out = match &mut self.attention {
|
||||
AttentionKind::Full(attn) => attn.forward(&h, attn_mask, offset)?,
|
||||
// Linear attention ignores attn_mask + offset; its causal
|
||||
AttentionKind::Full(attn) => attn.forward(&h, attn_mask, cos, sin)?,
|
||||
// Linear attention ignores attn_mask + rope; its causal
|
||||
// structure is baked into the recurrent state lifecycle.
|
||||
AttentionKind::Linear(net) => net.forward(&h)?,
|
||||
};
|
||||
|
||||
@@ -96,7 +96,8 @@ impl Qwen3_5Attention {
|
||||
&mut self,
|
||||
x: &Tensor,
|
||||
attn_mask: Option<&Tensor>,
|
||||
offset: usize,
|
||||
cos: &Tensor,
|
||||
sin: &Tensor,
|
||||
) -> candle_core::Result<Tensor> {
|
||||
let (b, l, _) = x.dims3()?;
|
||||
|
||||
@@ -131,8 +132,9 @@ impl Qwen3_5Attention {
|
||||
.transpose(1, 2)?
|
||||
.contiguous()?;
|
||||
|
||||
// 3. RoPE on q, k.
|
||||
let (q, k) = self.rotary.apply(&q, &k, offset)?;
|
||||
// 3. RoPE on q, k (cos/sin built once per forward by the model —
|
||||
// interleaved M-RoPE for image tokens, plain for text).
|
||||
let (q, k) = self.rotary.apply_cos_sin(&q, &k, cos, sin)?;
|
||||
|
||||
// 4. KV cache.
|
||||
let (k, v) = self.kv_cache.append(&k, &v)?;
|
||||
|
||||
@@ -737,6 +737,8 @@ mod tests {
|
||||
rope_theta: 10000.0,
|
||||
partial_rotary_factor: 1.0,
|
||||
rope_type: None,
|
||||
mrope_section: Vec::new(),
|
||||
mrope_interleaved: false,
|
||||
},
|
||||
rms_norm_eps: 1e-6,
|
||||
tie_word_embeddings: false,
|
||||
|
||||
@@ -191,11 +191,12 @@ fn default_hidden_act() -> String {
|
||||
}
|
||||
|
||||
/// Nested `rope_parameters` block from a Qwen3-Next `config.json`.
|
||||
/// `mrope_section` and `mrope_interleaved` are accepted via the
|
||||
/// `#[serde(default)]` flatten-tolerance below but ignored — we treat
|
||||
/// MRoPE as plain RoPE for text-only inference (the three position
|
||||
/// grids carry identical ids when there's no vision input, so the
|
||||
/// interleaving is a no-op).
|
||||
///
|
||||
/// For text-only inference the three MRoPE position grids carry
|
||||
/// identical ids, so the interleave is a no-op and plain RoPE applies.
|
||||
/// For vision inputs `mrope_section` + `mrope_interleaved` drive the
|
||||
/// per-axis (text/height/width) rotary used by image tokens — see
|
||||
/// `rope.rs`.
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct RopeParameters {
|
||||
/// Base for the inverse-frequency computation. Qwen3.6: 10_000_000.
|
||||
@@ -211,6 +212,16 @@ pub struct RopeParameters {
|
||||
/// implemented here.
|
||||
#[serde(default)]
|
||||
pub rope_type: Option<String>,
|
||||
/// MRoPE per-axis section sizes `[text, height, width]` — e.g.
|
||||
/// `[11, 11, 10]` for Qwen3.6, summing to the rotary half-dim.
|
||||
/// Empty for models that don't declare MRoPE (→ plain RoPE).
|
||||
#[serde(default)]
|
||||
pub mrope_section: Vec<usize>,
|
||||
/// Whether the three MRoPE axes are interleaved per-frequency
|
||||
/// (Qwen3-VL / Qwen3.6 style, `true`) rather than block-concatenated
|
||||
/// (Qwen2-VL style, `false`).
|
||||
#[serde(default)]
|
||||
pub mrope_interleaved: bool,
|
||||
}
|
||||
|
||||
fn default_rope_theta() -> f64 {
|
||||
@@ -236,7 +247,11 @@ fn default_partial_rotary_factor() -> f32 {
|
||||
/// `slice_assign` per run. For typical Qwen3.6 requests this is one
|
||||
/// or two runs per image; `slice_assign` does one tensor copy per
|
||||
/// run, which is cheap relative to the decoder forward pass.
|
||||
fn splice_runs(h: &Tensor, img: &Tensor, positions: &[u32]) -> candle_core::Result<Tensor> {
|
||||
pub(crate) fn splice_runs(
|
||||
h: &Tensor,
|
||||
img: &Tensor,
|
||||
positions: &[u32],
|
||||
) -> candle_core::Result<Tensor> {
|
||||
debug_assert!(
|
||||
!positions.is_empty(),
|
||||
"splice_runs precondition: non-empty positions"
|
||||
@@ -299,6 +314,16 @@ pub struct Qwen3_5Model {
|
||||
embed_tokens: Embedding,
|
||||
layers: Vec<Qwen3_5DecoderLayer>,
|
||||
norm: Qwen3_5RmsNorm,
|
||||
/// Shared with every full-attention layer; the model uses it to
|
||||
/// build the per-forward cos/sin (interleaved M-RoPE for image
|
||||
/// tokens, plain for text) once, which the layers then apply.
|
||||
rotary: Arc<RotaryEmbedding>,
|
||||
/// `offset + rope_delta` is the text-axis position during decode.
|
||||
/// 0 for text-only; set from `get_rope_index` during a vision
|
||||
/// prefill (image tokens compress the position space, so text after
|
||||
/// the image resumes from a smaller counter than the sequence
|
||||
/// index). Reset in `clear_kv_cache`.
|
||||
rope_delta: i64,
|
||||
device: Device,
|
||||
dtype: DType,
|
||||
}
|
||||
@@ -350,6 +375,8 @@ impl Qwen3_5Model {
|
||||
embed_tokens,
|
||||
layers,
|
||||
norm,
|
||||
rotary,
|
||||
rope_delta: 0,
|
||||
device,
|
||||
dtype,
|
||||
})
|
||||
@@ -363,6 +390,9 @@ impl Qwen3_5Model {
|
||||
for l in &mut self.layers {
|
||||
l.clear_kv_cache();
|
||||
}
|
||||
// New request → no image-compressed position offset until the
|
||||
// next vision prefill sets one.
|
||||
self.rope_delta = 0;
|
||||
}
|
||||
|
||||
fn causal_mask(&self, b: usize, tgt: usize, offset: usize) -> candle_core::Result<Tensor> {
|
||||
@@ -374,7 +404,7 @@ impl Qwen3_5Model {
|
||||
}
|
||||
|
||||
pub fn forward(&mut self, input: &Tensor, offset: usize) -> candle_core::Result<Tensor> {
|
||||
self.forward_inner(input, offset, None, None)
|
||||
self.forward_inner(input, offset, None, None, &[])
|
||||
}
|
||||
|
||||
/// Forward with image-embedding splice. Stage B of the vision plan.
|
||||
@@ -392,23 +422,25 @@ impl Qwen3_5Model {
|
||||
///
|
||||
/// The splice replaces the LM's text-side embedding at each
|
||||
/// `image_token_id` position with the corresponding row from
|
||||
/// `image_embeds`. After the splice the decoder runs unchanged.
|
||||
///
|
||||
/// **MRoPE gap.** Qwen3.6's `rope_parameters` declares MRoPE
|
||||
/// (interleaved text/height/width axes); Stage B applies plain
|
||||
/// text-position RoPE to image tokens. The model still attends
|
||||
/// to image content but loses spatial structure that MRoPE-aware
|
||||
/// position encoding would preserve. Tracked under issue #15
|
||||
/// (numerical validation) — quality benchmark from Stage D should
|
||||
/// surface the impact, and the fix lives in `rope::RotaryEmbedding`.
|
||||
/// `image_embeds`. After the splice the decoder runs the interleaved
|
||||
/// M-RoPE path: `grids` carries each image's post-merge LM grid
|
||||
/// `(lm_gh, lm_gw)` so `get_rope_index` assigns image tokens their 2D
|
||||
/// coordinates (dynamic resolution, #14).
|
||||
pub fn forward_with_vision(
|
||||
&mut self,
|
||||
input_ids: &Tensor,
|
||||
offset: usize,
|
||||
image_embeds: &Tensor,
|
||||
image_token_id: u32,
|
||||
grids: &[(usize, usize)],
|
||||
) -> candle_core::Result<Tensor> {
|
||||
self.forward_inner(input_ids, offset, Some(image_embeds), Some(image_token_id))
|
||||
self.forward_inner(
|
||||
input_ids,
|
||||
offset,
|
||||
Some(image_embeds),
|
||||
Some(image_token_id),
|
||||
grids,
|
||||
)
|
||||
}
|
||||
|
||||
fn forward_inner(
|
||||
@@ -417,19 +449,20 @@ impl Qwen3_5Model {
|
||||
offset: usize,
|
||||
image_embeds: Option<&Tensor>,
|
||||
image_token_id: Option<u32>,
|
||||
grids: &[(usize, usize)],
|
||||
) -> candle_core::Result<Tensor> {
|
||||
let (b, l) = input.dims2()?;
|
||||
let mut h = self.embed_tokens.forward(input)?;
|
||||
// Splice image embeddings at `image_token_id` positions. The
|
||||
// caller pre-expanded the prompt so every patch token in the
|
||||
// image_embeds tensor has a matching position in `input`. We
|
||||
// index_put the rows in place.
|
||||
if let (Some(img), Some(tok_id)) = (image_embeds, image_token_id) {
|
||||
// Locate image-token positions in input_ids. Operate on
|
||||
// CPU since the input ids are tiny (max ~10k entries
|
||||
// including the patch expansion) and the comparison is
|
||||
// not in the per-step hot path.
|
||||
|
||||
// Vision path: splice image embeddings at `image_token_id`
|
||||
// positions and build interleaved M-RoPE cos/sin so image tokens
|
||||
// carry their 2D (lm_gh × lm_gw) grid coordinates. Text / decode skip the
|
||||
// device→host id copy entirely and take the plain-RoPE fast path
|
||||
// — bit-for-bit the pre-M-RoPE behaviour when `rope_delta == 0`.
|
||||
let (cos, sin) = if let (Some(img), Some(tok_id)) = (image_embeds, image_token_id) {
|
||||
// Token ids on CPU — reused for the splice + position ids.
|
||||
let ids: Vec<u32> = input.flatten_all()?.to_vec1()?;
|
||||
|
||||
let mut positions: Vec<u32> = Vec::with_capacity(img.dim(0)?);
|
||||
for (idx, id) in ids.iter().enumerate() {
|
||||
if *id == tok_id {
|
||||
@@ -447,22 +480,22 @@ impl Qwen3_5Model {
|
||||
);
|
||||
}
|
||||
if !positions.is_empty() {
|
||||
// Cast image_embeds to the LM's dtype so the splice
|
||||
// produces a uniform tensor for the decoder stack.
|
||||
// Cast image_embeds to the LM's dtype, then splice the
|
||||
// contiguous `<|image_pad|>` runs in place.
|
||||
let img = img.to_dtype(self.dtype)?;
|
||||
// index_select would return the rows; we want to put.
|
||||
// candle's slice_assign with explicit positions ranges
|
||||
// doesn't exist; use scatter via index_select + an
|
||||
// accumulator: build a `(B, L, hidden)` zero tensor,
|
||||
// scatter the image rows in, then add to a masked
|
||||
// version of `h`. Simpler approach: walk positions
|
||||
// and use `slice_assign` for contiguous runs. Since
|
||||
// image_pad runs are contiguous (template emits
|
||||
// `<|vision_start|><|image_pad|>×N<|vision_end|>`),
|
||||
// we group positions and assign per run.
|
||||
h = splice_runs(&h, &img, &positions)?;
|
||||
}
|
||||
}
|
||||
|
||||
let (text, height, width, delta) = rope::get_rope_index(&ids, tok_id, grids)
|
||||
.map_err(|e| candle_core::Error::Msg(format!("get_rope_index: {e}")))?;
|
||||
self.rope_delta = delta;
|
||||
let pos = rope::mrope_position_tensor(&text, &height, &width, &self.device)?;
|
||||
self.rotary.mrope_cos_sin(&pos)?
|
||||
} else {
|
||||
let base = (offset as i64 + self.rope_delta).max(0) as usize;
|
||||
self.rotary.plain_cos_sin(base, l)?
|
||||
};
|
||||
|
||||
// Causal mask only needed for L > 1 prefill; full-attention
|
||||
// layers consume it via broadcast_add. Linear-attention layers
|
||||
// ignore the mask.
|
||||
@@ -472,7 +505,7 @@ impl Qwen3_5Model {
|
||||
Some(self.causal_mask(b, l, offset)?)
|
||||
};
|
||||
for layer in &mut self.layers {
|
||||
h = layer.forward(&h, causal.as_ref(), offset)?;
|
||||
h = layer.forward(&h, causal.as_ref(), &cos, &sin)?;
|
||||
}
|
||||
self.norm.forward(&h)
|
||||
}
|
||||
@@ -573,11 +606,12 @@ impl Qwen3_5ForCausalLM {
|
||||
offset: usize,
|
||||
image_embeds: &Tensor,
|
||||
image_token_id: u32,
|
||||
grids: &[(usize, usize)],
|
||||
) -> candle_core::Result<Tensor> {
|
||||
let (_, l) = input.dims2()?;
|
||||
let hidden = self
|
||||
.base
|
||||
.forward_with_vision(input, offset, image_embeds, image_token_id)?;
|
||||
let hidden =
|
||||
self.base
|
||||
.forward_with_vision(input, offset, image_embeds, image_token_id, grids)?;
|
||||
hidden.i((.., l - 1.., ..))?.apply(&self.lm_head)
|
||||
}
|
||||
|
||||
|
||||
@@ -1,19 +1,27 @@
|
||||
//! Rotary position embedding for Qwen3-Next's full-attention layers.
|
||||
//!
|
||||
//! Qwen3.6 ships with MRoPE (multimodal RoPE) machinery in the
|
||||
//! reference Python — three position grids interleaved per
|
||||
//! `mrope_section`. For text-only inference all three grids carry the
|
||||
//! same position ids and the interleave is a no-op, so this module
|
||||
//! implements the plain (non-mrope) flavour: the standard inv_freq
|
||||
//! cosine/sine tables driven by `rope_theta` and `head_dim`.
|
||||
//! Qwen3.6 declares **interleaved M-RoPE** (multimodal RoPE): the
|
||||
//! rotary half-dimension is split across three position axes —
|
||||
//! `[text, height, width]` per `mrope_section` (`[11,11,10]` for
|
||||
//! Qwen3.6) — interleaved per-frequency. For **text** every token's
|
||||
//! three axes carry the same position id, so the interleave is a no-op
|
||||
//! and this reduces exactly to plain RoPE. For **image** tokens the
|
||||
//! height/width axes carry the patch's 2D grid coordinates, which is
|
||||
//! how the model reads the 14×14 patch layout (without it, all patches
|
||||
//! share a height position and the image reads as vertical repetition).
|
||||
//!
|
||||
//! Rotation flavour: **GLM-style** rotate-half (the second half of the
|
||||
//! head dim is negated and swapped into the first). The reference
|
||||
//! Python uses `apply_rotary_pos_emb` with `rotate_half`; candle's
|
||||
//! `rope_slow` is the matching helper.
|
||||
//! Two cos/sin builders feed a shared [`RotaryEmbedding::apply`]:
|
||||
//! - [`RotaryEmbedding::plain_cos_sin`] narrows the precomputed tables
|
||||
//! at a scalar position — the text / decode fast path.
|
||||
//! - [`RotaryEmbedding::mrope_cos_sin`] builds per-token cos/sin from a
|
||||
//! `(3, seq)` position-id tensor, blending the three axes' frequencies
|
||||
//! at the interleave index sets — the vision-prefill path.
|
||||
//!
|
||||
//! Rotation flavour: **GLM-style** rotate-half (candle's `rope_slow`),
|
||||
//! matching the reference Python's `apply_rotary_pos_emb` + `rotate_half`.
|
||||
|
||||
use anyhow::Result;
|
||||
use candle_core::{DType, Device, Tensor};
|
||||
use candle_core::{DType, Device, IndexOp, Tensor};
|
||||
|
||||
use super::TextConfig;
|
||||
|
||||
@@ -21,6 +29,18 @@ use super::TextConfig;
|
||||
pub struct RotaryEmbedding {
|
||||
sin: Tensor,
|
||||
cos: Tensor,
|
||||
/// Inverse frequencies, shape `(1, rotary_dim/2)`. Retained (beyond
|
||||
/// the precomputed `sin`/`cos` tables) so [`Self::mrope_cos_sin`] can
|
||||
/// build cos/sin from arbitrary per-axis position ids.
|
||||
inv_freq: Tensor,
|
||||
/// Per-axis column masks over the rotary half-dim, shape `(1, half)`,
|
||||
/// f32 0/1. `mask_t + mask_h + mask_w` partitions the columns; a
|
||||
/// column belongs to exactly one axis. For a non-MRoPE config
|
||||
/// `mask_t` is all-ones and the others all-zero (→ plain RoPE).
|
||||
mask_t: Tensor,
|
||||
mask_h: Tensor,
|
||||
mask_w: Tensor,
|
||||
dtype: DType,
|
||||
/// Number of dims at the head's leading edge that the rotation
|
||||
/// covers. The remaining `head_dim - rotary_dim` dims pass through
|
||||
/// unchanged. Qwen3-Next uses `partial_rotary_factor = 0.25`, so
|
||||
@@ -29,6 +49,52 @@ pub struct RotaryEmbedding {
|
||||
head_dim: usize,
|
||||
}
|
||||
|
||||
/// Build the per-axis 0/1 column masks over the rotary half-dim from
|
||||
/// `mrope_section`. Returns `(temporal, height, width)` each length
|
||||
/// `half`. Temporal is the complement of height ∪ width, so the three
|
||||
/// masks always partition `0..half` and reduce to all-temporal (plain
|
||||
/// RoPE) when no usable section is given.
|
||||
fn mrope_masks(
|
||||
half: usize,
|
||||
section: &[usize],
|
||||
interleaved: bool,
|
||||
) -> (Vec<f32>, Vec<f32>, Vec<f32>) {
|
||||
let mut mh = vec![0f32; half];
|
||||
let mut mw = vec![0f32; half];
|
||||
if section.len() == 3 {
|
||||
if interleaved {
|
||||
// Qwen3-VL: height at columns 1,4,7,… ; width at 2,5,8,… ;
|
||||
// temporal keeps 0,3,6,… — each `take`n from `mrope_section`.
|
||||
for i in (1..half).step_by(3).take(section[1]) {
|
||||
mh[i] = 1.0;
|
||||
}
|
||||
for i in (2..half).step_by(3).take(section[2]) {
|
||||
mw[i] = 1.0;
|
||||
}
|
||||
} else {
|
||||
// Qwen2-VL: contiguous blocks [text | height | width].
|
||||
let h_start = section[0].min(half);
|
||||
let h_end = (section[0] + section[1]).min(half);
|
||||
for m in mh.iter_mut().take(h_end).skip(h_start) {
|
||||
*m = 1.0;
|
||||
}
|
||||
for m in mw.iter_mut().take(half).skip(h_end) {
|
||||
*m = 1.0;
|
||||
}
|
||||
}
|
||||
}
|
||||
let mt: Vec<f32> = (0..half)
|
||||
.map(|i| {
|
||||
if mh[i] == 0.0 && mw[i] == 0.0 {
|
||||
1.0
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
(mt, mh, mw)
|
||||
}
|
||||
|
||||
impl RotaryEmbedding {
|
||||
pub fn new(dtype: DType, cfg: &TextConfig, dev: &Device) -> Result<Self> {
|
||||
let head_dim = cfg.head_dim;
|
||||
@@ -52,44 +118,88 @@ impl RotaryEmbedding {
|
||||
.step_by(2)
|
||||
.map(|i| 1f32 / rope.rope_theta.powf(i as f64 / rotary_dim as f64) as f32)
|
||||
.collect();
|
||||
let n = inv_freq.len();
|
||||
let inv_freq = Tensor::from_vec(inv_freq, (1, n), dev)?.to_dtype(DType::F32)?;
|
||||
let half = inv_freq.len();
|
||||
let inv_freq = Tensor::from_vec(inv_freq, (1, half), dev)?.to_dtype(DType::F32)?;
|
||||
let t = Tensor::arange(0u32, max_seq_len as u32, dev)?
|
||||
.to_dtype(DType::F32)?
|
||||
.reshape((max_seq_len, 1))?;
|
||||
let freqs = t.matmul(&inv_freq)?;
|
||||
|
||||
// MRoPE axis masks. `sum(mrope_section)` should equal `half`;
|
||||
// warn-tolerant: any shortfall just stays on the temporal axis.
|
||||
let (mt, mh, mw) = mrope_masks(half, &rope.mrope_section, rope.mrope_interleaved);
|
||||
let mask_t = Tensor::from_vec(mt, (1, half), dev)?;
|
||||
let mask_h = Tensor::from_vec(mh, (1, half), dev)?;
|
||||
let mask_w = Tensor::from_vec(mw, (1, half), dev)?;
|
||||
|
||||
Ok(Self {
|
||||
sin: freqs.sin()?.to_dtype(dtype)?,
|
||||
cos: freqs.cos()?.to_dtype(dtype)?,
|
||||
inv_freq,
|
||||
mask_t,
|
||||
mask_h,
|
||||
mask_w,
|
||||
dtype,
|
||||
rotary_dim,
|
||||
head_dim,
|
||||
})
|
||||
}
|
||||
|
||||
/// Apply RoPE to q, k.
|
||||
///
|
||||
/// `q`, `k` shape: `(B, H, L, head_dim)`. `offset` is the index
|
||||
/// into the cached cos/sin table — the position of the first token
|
||||
/// in the current step.
|
||||
///
|
||||
/// When `rotary_dim < head_dim` the rotation is applied only to the
|
||||
/// first `rotary_dim` dims of each head; the tail passes through
|
||||
/// unchanged (matches the reference Python's
|
||||
/// `apply_rotary_pos_emb` with non-trivial `partial_rotary_factor`).
|
||||
pub fn apply(
|
||||
/// cos/sin for a contiguous run of `seq_len` positions starting at
|
||||
/// `pos`, by narrowing the precomputed tables. The text / decode
|
||||
/// path (all three MRoPE axes equal → plain RoPE). Shape
|
||||
/// `(seq_len, rotary_dim/2)`.
|
||||
pub fn plain_cos_sin(
|
||||
&self,
|
||||
pos: usize,
|
||||
seq_len: usize,
|
||||
) -> candle_core::Result<(Tensor, Tensor)> {
|
||||
let cos = self.cos.narrow(0, pos, seq_len)?;
|
||||
let sin = self.sin.narrow(0, pos, seq_len)?;
|
||||
Ok((cos, sin))
|
||||
}
|
||||
|
||||
/// cos/sin from explicit per-token 3D position ids, shape
|
||||
/// `(3, seq_len)` (axes: text, height, width). Builds each axis's
|
||||
/// frequencies and blends them at the interleave index sets, so
|
||||
/// every rotary frequency slot is driven by exactly one axis.
|
||||
/// Reduces exactly to [`Self::plain_cos_sin`] when the three axes are
|
||||
/// equal. Returns cos/sin of shape `(seq_len, rotary_dim/2)`.
|
||||
pub fn mrope_cos_sin(&self, position_ids: &Tensor) -> candle_core::Result<(Tensor, Tensor)> {
|
||||
let pos = position_ids.to_dtype(DType::F32)?;
|
||||
let (axes, seq_len) = pos.dims2()?;
|
||||
debug_assert_eq!(axes, 3, "mrope position_ids must have 3 axes");
|
||||
// Per-axis freqs: pos[a] (seq,1) @ inv_freq (1,half) → (seq,half).
|
||||
let ft = pos.i(0)?.reshape((seq_len, 1))?.matmul(&self.inv_freq)?;
|
||||
let fh = pos.i(1)?.reshape((seq_len, 1))?.matmul(&self.inv_freq)?;
|
||||
let fw = pos.i(2)?.reshape((seq_len, 1))?.matmul(&self.inv_freq)?;
|
||||
// Blend: each column belongs to exactly one axis (masks partition
|
||||
// the half-dim), so this picks the right axis per frequency slot.
|
||||
let blended = ft
|
||||
.broadcast_mul(&self.mask_t)?
|
||||
.add(&fh.broadcast_mul(&self.mask_h)?)?
|
||||
.add(&fw.broadcast_mul(&self.mask_w)?)?;
|
||||
let cos = blended.cos()?.to_dtype(self.dtype)?;
|
||||
let sin = blended.sin()?.to_dtype(self.dtype)?;
|
||||
Ok((cos, sin))
|
||||
}
|
||||
|
||||
/// Apply rotary to `q`, `k` (shape `(B, H, L, head_dim)`) using
|
||||
/// precomputed `cos`/`sin` of shape `(L, rotary_dim/2)`. Partial
|
||||
/// rotary: only the first `rotary_dim` dims rotate; the tail passes
|
||||
/// through unchanged.
|
||||
pub fn apply_cos_sin(
|
||||
&self,
|
||||
q: &Tensor,
|
||||
k: &Tensor,
|
||||
offset: usize,
|
||||
cos: &Tensor,
|
||||
sin: &Tensor,
|
||||
) -> candle_core::Result<(Tensor, Tensor)> {
|
||||
let (_, _, seq_len, head_dim_in) = q.dims4()?;
|
||||
let (_, _, _seq_len, head_dim_in) = q.dims4()?;
|
||||
debug_assert_eq!(head_dim_in, self.head_dim, "q head_dim mismatch");
|
||||
let cos = self.cos.narrow(0, offset, seq_len)?;
|
||||
let sin = self.sin.narrow(0, offset, seq_len)?;
|
||||
if self.rotary_dim == self.head_dim {
|
||||
// Full rotation.
|
||||
let q_embed = candle_nn::rotary_emb::rope_slow(&q.contiguous()?, &cos, &sin)?;
|
||||
let k_embed = candle_nn::rotary_emb::rope_slow(&k.contiguous()?, &cos, &sin)?;
|
||||
let q_embed = candle_nn::rotary_emb::rope_slow(&q.contiguous()?, cos, sin)?;
|
||||
let k_embed = candle_nn::rotary_emb::rope_slow(&k.contiguous()?, cos, sin)?;
|
||||
Ok((q_embed, k_embed))
|
||||
} else {
|
||||
// Partial rotation: narrow → rotate → cat the untouched tail.
|
||||
@@ -102,8 +212,8 @@ impl RotaryEmbedding {
|
||||
.narrow(candle_core::D::Minus1, 0, self.rotary_dim)?
|
||||
.contiguous()?;
|
||||
let k_pass = k.narrow(candle_core::D::Minus1, self.rotary_dim, tail)?;
|
||||
let q_rotated = candle_nn::rotary_emb::rope_slow(&q_rot, &cos, &sin)?;
|
||||
let k_rotated = candle_nn::rotary_emb::rope_slow(&k_rot, &cos, &sin)?;
|
||||
let q_rotated = candle_nn::rotary_emb::rope_slow(&q_rot, cos, sin)?;
|
||||
let k_rotated = candle_nn::rotary_emb::rope_slow(&k_rot, cos, sin)?;
|
||||
let q_embed =
|
||||
Tensor::cat(&[&q_rotated, &q_pass.contiguous()?], candle_core::D::Minus1)?;
|
||||
let k_embed =
|
||||
@@ -112,3 +222,358 @@ impl RotaryEmbedding {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute interleaved-M-RoPE 3D position ids for a full prompt that may
|
||||
/// contain image-placeholder runs, plus the decode `rope_delta`.
|
||||
///
|
||||
/// Mirrors the reference `get_rope_index`:
|
||||
/// - text tokens advance a single running counter `c`, all three axes
|
||||
/// equal (`[c, c, c]`);
|
||||
/// - each contiguous run of `image_token_id` is one image; its tokens get
|
||||
/// `[base + t, base + h, base + w]` in row-major (t outer, h, w inner),
|
||||
/// where `base` is the counter at the run's start; after the run the
|
||||
/// counter resumes from `base + max(grid_t, grid_h, grid_w)`.
|
||||
///
|
||||
/// Returns `(text_pos, height_pos, width_pos, rope_delta)`, each pos `Vec`
|
||||
/// length `input_ids.len()`. `rope_delta = final_counter - seq_len`: add it
|
||||
/// to a plain decode offset so text resumes from the counter after the
|
||||
/// (position-compressed) image blocks.
|
||||
///
|
||||
/// Whether interleaved M-RoPE for image tokens is enabled. Default
|
||||
/// **on** — Qwen3.6 was trained with interleaved M-RoPE, and this
|
||||
/// implementation matches the HF `apply_interleaved_mrope` /
|
||||
/// `get_rope_index` reference exactly (verified column-for-column). The
|
||||
/// env var is a **kill switch**: `NEURON_MROPE=0` falls back to plain
|
||||
/// sequential positions for image tokens (the pre-M-RoPE behaviour).
|
||||
pub(crate) fn mrope_enabled() -> bool {
|
||||
std::env::var("NEURON_MROPE")
|
||||
.map(|v| {
|
||||
!matches!(
|
||||
v.trim().to_ascii_lowercase().as_str(),
|
||||
"0" | "false" | "no" | "off"
|
||||
)
|
||||
})
|
||||
.unwrap_or(true)
|
||||
}
|
||||
|
||||
/// Position ids for the forward path. Gated by [`mrope_enabled`]: when
|
||||
/// off, returns plain sequential identity positions on all three axes
|
||||
/// (`mrope_cos_sin` then reduces exactly to plain RoPE), restoring the
|
||||
/// pre-M-RoPE behaviour without touching the rest of the forward.
|
||||
pub(crate) fn get_rope_index(
|
||||
input_ids: &[u32],
|
||||
image_token_id: u32,
|
||||
grids: &[(usize, usize)],
|
||||
) -> Result<MRopeIndex> {
|
||||
if !mrope_enabled() {
|
||||
let seq: Vec<i64> = (0..input_ids.len() as i64).collect();
|
||||
return Ok((seq.clone(), seq.clone(), seq, 0));
|
||||
}
|
||||
compute_mrope_index(input_ids, image_token_id, grids)
|
||||
}
|
||||
|
||||
/// The real interleaved-M-RoPE position-id computation (always active in
|
||||
/// unit tests; gated behind [`get_rope_index`] at runtime).
|
||||
///
|
||||
/// `grids` carries the post-merge LM grid `(lm_gh, lm_gw)` for each image
|
||||
/// run, in prompt order — a run length alone cannot recover its
|
||||
/// factorisation, so the grids must be passed (#14 dynamic resolution).
|
||||
/// Each image is a still frame (`grid_t = 1`); its tokens get
|
||||
/// `[base, base + hh, base + ww]` row-major and the shared counter
|
||||
/// resumes at `base + max(lm_gh, lm_gw)`. Multi-image is correct because
|
||||
/// the counter threads across images and interleaved text.
|
||||
pub(crate) fn compute_mrope_index(
|
||||
input_ids: &[u32],
|
||||
image_token_id: u32,
|
||||
grids: &[(usize, usize)],
|
||||
) -> Result<MRopeIndex> {
|
||||
let n = input_ids.len();
|
||||
let mut text = Vec::with_capacity(n);
|
||||
let mut height = Vec::with_capacity(n);
|
||||
let mut width = Vec::with_capacity(n);
|
||||
let mut counter: i64 = 0;
|
||||
let mut i = 0;
|
||||
let mut k = 0; // index into `grids`, one per image run
|
||||
while i < n {
|
||||
if input_ids[i] == image_token_id {
|
||||
let start = i;
|
||||
while i < n && input_ids[i] == image_token_id {
|
||||
i += 1;
|
||||
}
|
||||
let run = i - start;
|
||||
let (grid_h, grid_w) = *grids.get(k).ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"get_rope_index: image run #{k} (len {run}) has no matching grid \
|
||||
({} grids supplied)",
|
||||
grids.len()
|
||||
)
|
||||
})?;
|
||||
k += 1;
|
||||
if grid_h * grid_w != run {
|
||||
anyhow::bail!(
|
||||
"get_rope_index: image run #{} length {run} != grid {grid_h}×{grid_w} = {}",
|
||||
k - 1,
|
||||
grid_h * grid_w
|
||||
);
|
||||
}
|
||||
let base = counter;
|
||||
for hh in 0..grid_h {
|
||||
for ww in 0..grid_w {
|
||||
text.push(base); // grid_t = 1 → temporal axis const
|
||||
height.push(base + hh as i64);
|
||||
width.push(base + ww as i64);
|
||||
}
|
||||
}
|
||||
counter = base + grid_h.max(grid_w) as i64;
|
||||
} else {
|
||||
text.push(counter);
|
||||
height.push(counter);
|
||||
width.push(counter);
|
||||
counter += 1;
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
if k != grids.len() {
|
||||
anyhow::bail!(
|
||||
"get_rope_index: prompt has {k} image run(s) but {} grid(s) were supplied",
|
||||
grids.len()
|
||||
);
|
||||
}
|
||||
let delta = counter - n as i64;
|
||||
Ok((text, height, width, delta))
|
||||
}
|
||||
|
||||
/// `(text_pos, height_pos, width_pos, rope_delta)` returned by
|
||||
/// [`get_rope_index`]; the three vectors combine into the `(3, seq)`
|
||||
/// MRoPE position-id tensor.
|
||||
pub(crate) type MRopeIndex = (Vec<i64>, Vec<i64>, Vec<i64>, i64);
|
||||
|
||||
/// Build the `(3, seq)` position-id tensor consumed by
|
||||
/// [`RotaryEmbedding::mrope_cos_sin`] from the three axis vectors.
|
||||
///
|
||||
/// Built directly as **f32** (positions are small integers, exact in
|
||||
/// f32 well past any context length): the freqs matmul needs float
|
||||
/// anyway, and this avoids an i64 tensor / i64→f32 cast on the GPU.
|
||||
pub(crate) fn mrope_position_tensor(
|
||||
text: &[i64],
|
||||
height: &[i64],
|
||||
width: &[i64],
|
||||
dev: &Device,
|
||||
) -> candle_core::Result<Tensor> {
|
||||
let seq = text.len();
|
||||
let mut flat = Vec::with_capacity(3 * seq);
|
||||
flat.extend(text.iter().map(|&x| x as f32));
|
||||
flat.extend(height.iter().map(|&x| x as f32));
|
||||
flat.extend(width.iter().map(|&x| x as f32));
|
||||
Tensor::from_vec(flat, (3, seq), dev)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use candle_core::IndexOp;
|
||||
|
||||
/// A TextConfig stub with Qwen3.6's rope params (head_dim 256,
|
||||
/// partial 0.25 → rotary_dim 64 → half 32; section [11,11,10]).
|
||||
fn qwen36_cfg() -> TextConfig {
|
||||
serde_json::from_value(serde_json::json!({
|
||||
"hidden_size": 5120,
|
||||
"num_hidden_layers": 1,
|
||||
"num_attention_heads": 64,
|
||||
"num_key_value_heads": 8,
|
||||
"head_dim": 256,
|
||||
"intermediate_size": 1,
|
||||
"vocab_size": 10,
|
||||
"rms_norm_eps": 1e-6,
|
||||
"max_position_embeddings": 64,
|
||||
"layer_types": ["full_attention"],
|
||||
"rope_parameters": {
|
||||
"rope_theta": 10000000.0,
|
||||
"partial_rotary_factor": 0.25,
|
||||
"mrope_section": [11, 11, 10],
|
||||
"mrope_interleaved": true
|
||||
}
|
||||
}))
|
||||
.expect("cfg")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mrope_masks_partition_the_half_dim() {
|
||||
let (mt, mh, mw) = mrope_masks(32, &[11, 11, 10], true);
|
||||
// Each column belongs to exactly one axis.
|
||||
for i in 0..32 {
|
||||
let s = mt[i] + mh[i] + mw[i];
|
||||
assert_eq!(s, 1.0, "column {i} covered {s} times");
|
||||
}
|
||||
assert_eq!(mt.iter().sum::<f32>(), 11.0);
|
||||
assert_eq!(mh.iter().sum::<f32>(), 11.0);
|
||||
assert_eq!(mw.iter().sum::<f32>(), 10.0);
|
||||
// Interleave: temporal 0,3,…; height 1,4,…; width 2,5,…
|
||||
assert_eq!(mt[0], 1.0);
|
||||
assert_eq!(mh[1], 1.0);
|
||||
assert_eq!(mw[2], 1.0);
|
||||
assert_eq!(mt[3], 1.0);
|
||||
}
|
||||
|
||||
/// The load-bearing invariant: when all three position axes are
|
||||
/// equal (text), `mrope_cos_sin` must reproduce `plain_cos_sin`
|
||||
/// bit-for-bit — i.e. M-RoPE is a no-op for text, so text inference
|
||||
/// is unchanged.
|
||||
#[test]
|
||||
fn mrope_reduces_to_plain_for_equal_axes() {
|
||||
let dev = Device::Cpu;
|
||||
let rope = RotaryEmbedding::new(DType::F32, &qwen36_cfg(), &dev).unwrap();
|
||||
|
||||
// positions 5,6,7 on all three axes.
|
||||
let base: Vec<i64> = vec![5, 6, 7];
|
||||
let pos =
|
||||
Tensor::from_vec([base.clone(), base.clone(), base].concat(), (3, 3), &dev).unwrap();
|
||||
|
||||
let (mc, ms) = rope.mrope_cos_sin(&pos).unwrap();
|
||||
let (pc, ps) = rope.plain_cos_sin(5, 3).unwrap();
|
||||
|
||||
let dcos = (mc - pc).unwrap().abs().unwrap().max_all().unwrap();
|
||||
let dsin = (ms - ps).unwrap().abs().unwrap().max_all().unwrap();
|
||||
assert!(
|
||||
dcos.to_scalar::<f32>().unwrap() < 1e-6,
|
||||
"cos mismatch {dcos:?}"
|
||||
);
|
||||
assert!(
|
||||
dsin.to_scalar::<f32>().unwrap() < 1e-6,
|
||||
"sin mismatch {dsin:?}"
|
||||
);
|
||||
}
|
||||
|
||||
/// Hand-checked interleave: a width-axis column (index 2) must track
|
||||
/// the WIDTH position, while a temporal column (index 0) tracks the
|
||||
/// TEXT position, even when the axes differ.
|
||||
#[test]
|
||||
fn mrope_blends_axes_at_interleave_columns() {
|
||||
let dev = Device::Cpu;
|
||||
let rope = RotaryEmbedding::new(DType::F32, &qwen36_cfg(), &dev).unwrap();
|
||||
let half = rope.inv_freq.dim(1).unwrap();
|
||||
let inv: Vec<f32> = rope.inv_freq.i(0).unwrap().to_vec1().unwrap();
|
||||
|
||||
// One token: text=10, height=3, width=7 — all distinct.
|
||||
let pos = Tensor::from_vec(vec![10i64, 3, 7], (3, 1), &dev).unwrap();
|
||||
let (cos, _sin) = rope.mrope_cos_sin(&pos).unwrap();
|
||||
let cos_row: Vec<f32> = cos.i(0).unwrap().to_vec1().unwrap();
|
||||
assert_eq!(cos_row.len(), half);
|
||||
|
||||
// Column 0 (temporal) → text pos 10. Column 1 (height) → 3.
|
||||
// Column 2 (width) → 7.
|
||||
assert!((cos_row[0] - (10.0 * inv[0]).cos()).abs() < 1e-5);
|
||||
assert!((cos_row[1] - (3.0 * inv[1]).cos()).abs() < 1e-5);
|
||||
assert!((cos_row[2] - (7.0 * inv[2]).cos()).abs() < 1e-5);
|
||||
assert!((cos_row[3] - (10.0 * inv[3]).cos()).abs() < 1e-5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_rope_index_text_only_is_sequential() {
|
||||
let (t, h, w, delta) = compute_mrope_index(&[1, 2, 3, 4], 99, &[]).unwrap();
|
||||
assert_eq!(t, vec![0, 1, 2, 3]);
|
||||
assert_eq!(h, vec![0, 1, 2, 3]);
|
||||
assert_eq!(w, vec![0, 1, 2, 3]);
|
||||
assert_eq!(delta, 0, "no image → delta 0 → plain decode positions");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_rope_index_text_image_text() {
|
||||
// [text, image(2x2 run of 4), text]. image_token = 99, grid (2,2).
|
||||
let ids = [1u32, 99, 99, 99, 99, 2];
|
||||
let (t, h, w, delta) = compute_mrope_index(&ids, 99, &[(2, 2)]).unwrap();
|
||||
// token 0: text → 0. image base=1, grid 2x2:
|
||||
// t all = 1; h = base+row = [1,1,2,2]; w = base+col = [1,2,1,2].
|
||||
// resume from base + max(2,2) = 3. trailing text → 3.
|
||||
assert_eq!(t, vec![0, 1, 1, 1, 1, 3]);
|
||||
assert_eq!(h, vec![0, 1, 1, 2, 2, 3]);
|
||||
assert_eq!(w, vec![0, 1, 2, 1, 2, 3]);
|
||||
// final counter = 4, seq_len = 6 → delta = -2 (the 4 image tokens
|
||||
// advanced the counter by only 2).
|
||||
assert_eq!(delta, -2);
|
||||
// Decode after the prompt (offset = 6) → text position 6 + (-2) = 4.
|
||||
assert_eq!(6 + delta, 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_rope_index_nonsquare_single_image() {
|
||||
// text + image(2 rows × 3 cols = 6 tokens). grid (2,3).
|
||||
let ids = [1u32, 99, 99, 99, 99, 99, 99];
|
||||
let (t, h, w, delta) = compute_mrope_index(&ids, 99, &[(2, 3)]).unwrap();
|
||||
// base = 1; row-major h = [0,0,0,1,1,1]+1, w = [0,1,2,0,1,2]+1.
|
||||
assert_eq!(t, vec![0, 1, 1, 1, 1, 1, 1]);
|
||||
assert_eq!(h, vec![0, 1, 1, 1, 2, 2, 2]);
|
||||
assert_eq!(w, vec![0, 1, 2, 3, 1, 2, 3]);
|
||||
// resume from base + max(2,3) = 4; seq_len 7, counter 4 → delta -3.
|
||||
assert_eq!(delta, 4 - 7);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_rope_index_two_images_different_grids() {
|
||||
// img(2x2)=4, text, img(1x3)=3. grids [(2,2),(1,3)].
|
||||
let ids = [99, 99, 99, 99, 7, 99, 99, 99];
|
||||
let (t, h, w, delta) = compute_mrope_index(&ids, 99, &[(2, 2), (1, 3)]).unwrap();
|
||||
// img1 base=0 → t=0, h=[0,0,1,1], w=[0,1,0,1]; resume max(2,2)=2.
|
||||
// text at counter 2. img2 base=3 → t=3, h=[3,3,3], w=[3,4,5];
|
||||
// resume 3+max(1,3)=6.
|
||||
assert_eq!(t, vec![0, 0, 0, 0, 2, 3, 3, 3]);
|
||||
assert_eq!(h, vec![0, 0, 1, 1, 2, 3, 3, 3]);
|
||||
assert_eq!(w, vec![0, 1, 0, 1, 2, 3, 4, 5]);
|
||||
assert_eq!(delta, 6 - 8);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_rope_index_on_by_default() {
|
||||
// With NEURON_MROPE unset (default ON), the runtime path returns
|
||||
// the real interleaved-M-RoPE positions. (NEURON_MROPE=0 would fall
|
||||
// back to identity; not asserted here since it depends on env.)
|
||||
let (t, h, w, _delta) = get_rope_index(&[1, 99, 99, 99, 99, 2], 99, &[(2, 2)]).unwrap();
|
||||
assert_eq!(t, vec![0, 1, 1, 1, 1, 3]);
|
||||
assert_eq!(h, vec![0, 1, 1, 2, 2, 3]);
|
||||
assert_eq!(w, vec![0, 1, 2, 1, 2, 3]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_rope_index_grid_mismatches_error() {
|
||||
// run length != grid product.
|
||||
assert!(compute_mrope_index(&[99u32; 6], 99, &[(2, 2)]).is_err());
|
||||
// too few grids for the number of image runs.
|
||||
assert!(compute_mrope_index(&[99, 99, 7, 99], 99, &[(1, 2)]).is_err());
|
||||
// too many grids.
|
||||
assert!(compute_mrope_index(&[99, 99], 99, &[(1, 2), (1, 1)]).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn position_tensor_round_trips_through_mrope_cos_sin() {
|
||||
// get_rope_index → (3,seq) tensor → mrope_cos_sin, and confirm an
|
||||
// image token's height column tracks its grid row (not the text
|
||||
// counter), i.e. the end-to-end position plumbing is wired right.
|
||||
let dev = Device::Cpu;
|
||||
let rope = RotaryEmbedding::new(DType::F32, &qwen36_cfg(), &dev).unwrap();
|
||||
let ids = [1u32, 99, 99, 99, 99]; // text + 2x2 image
|
||||
let (t, h, w, _d) = compute_mrope_index(&ids, 99, &[(2, 2)]).unwrap();
|
||||
let pos = mrope_position_tensor(&t, &h, &w, &dev).unwrap();
|
||||
assert_eq!(pos.dims(), &[3, 5]);
|
||||
let (cos, _sin) = rope.mrope_cos_sin(&pos).unwrap();
|
||||
assert_eq!(cos.dims(), &[5, rope.inv_freq.dim(1).unwrap()]);
|
||||
|
||||
let inv: Vec<f32> = rope.inv_freq.i(0).unwrap().to_vec1().unwrap();
|
||||
// Last image token (index 4): grid (h=1, w=1) → base 1 → h=2, w=2.
|
||||
// Height column (index 1) must track h-position 2, not text.
|
||||
let last: Vec<f32> = cos.i(4).unwrap().to_vec1().unwrap();
|
||||
assert!((last[1] - (2.0 * inv[1]).cos()).abs() < 1e-5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_rope_index_196_is_14x14() {
|
||||
let mut ids = vec![1u32]; // one text token
|
||||
ids.extend(std::iter::repeat_n(99u32, 196));
|
||||
let (t, h, w, _delta) = compute_mrope_index(&ids, 99, &[(14, 14)]).unwrap();
|
||||
// image base = 1. Last image token (index 196) is grid (h=13,w=13).
|
||||
assert_eq!(*t.last().unwrap(), 1, "grid_t=1 → temporal const at base");
|
||||
assert_eq!(h[1], 1, "first image row at base");
|
||||
assert_eq!(w[1], 1, "first image col at base");
|
||||
assert_eq!(h[196], 1 + 13, "last image row = base + 13");
|
||||
assert_eq!(w[196], 1 + 13, "last image col = base + 13");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -48,6 +48,31 @@ use candle_nn::var_builder::ShardedVarBuilder;
|
||||
use candle_nn::{Conv2d, Conv2dConfig, Embedding, LayerNorm, Linear};
|
||||
use serde::Deserialize;
|
||||
|
||||
fn env_truthy(name: &str) -> bool {
|
||||
std::env::var(name)
|
||||
.map(|v| {
|
||||
matches!(
|
||||
v.trim().to_ascii_lowercase().as_str(),
|
||||
"1" | "true" | "yes" | "on"
|
||||
)
|
||||
})
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
/// Legacy escape hatch: when set, use the original Stage-A sequential
|
||||
/// `pos_embed` lookup instead of the bilinear grid interpolation.
|
||||
/// Default off (interpolation on) — for A/B comparison only.
|
||||
fn vision_legacy_pos() -> bool {
|
||||
env_truthy("NEURON_VISION_LEGACY_POS")
|
||||
}
|
||||
|
||||
/// Legacy escape hatch: when set, skip the 2D vision rotary in the ViT
|
||||
/// attention (the original Stage-A behaviour). Default off (rotary on)
|
||||
/// — for A/B comparison only.
|
||||
fn vision_legacy_rope() -> bool {
|
||||
env_truthy("NEURON_VISION_LEGACY_ROPE")
|
||||
}
|
||||
|
||||
/// Qwen3.6 vision tower hyperparameters. Mirrors the `vision_config`
|
||||
/// block of `config.json`. Only the fields we actually need are
|
||||
/// captured; serde tolerates the rest.
|
||||
@@ -118,10 +143,12 @@ impl VisionBlock {
|
||||
})
|
||||
}
|
||||
|
||||
/// `x`: `(N, hidden_size)` un-batched. Returns same shape.
|
||||
fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
/// `x`: `(N, hidden_size)` un-batched. `rotary`: optional
|
||||
/// `(cos, sin)` each `(N, head_dim/2)` — the 2D vision rotary applied
|
||||
/// to q/k. Returns same shape.
|
||||
fn forward(&self, x: &Tensor, rotary: Option<&(Tensor, Tensor)>) -> Result<Tensor> {
|
||||
let attn_in = self.norm1.forward(x)?;
|
||||
let attn_out = self.attention(&attn_in)?;
|
||||
let attn_out = self.attention(&attn_in, rotary)?;
|
||||
let x = x.add(&attn_out)?;
|
||||
let mlp_in = self.norm2.forward(&x)?;
|
||||
let mlp_out = self.fc2.forward(&gelu_tanh(&self.fc1.forward(&mlp_in)?)?)?;
|
||||
@@ -129,8 +156,11 @@ impl VisionBlock {
|
||||
}
|
||||
|
||||
/// Multi-head self-attention over the patch sequence. No causal
|
||||
/// mask — every patch attends to every other patch.
|
||||
fn attention(&self, x: &Tensor) -> Result<Tensor> {
|
||||
/// mask — every patch attends to every other patch. When `rotary` is
|
||||
/// given, the 2D vision rotary (row/col position) is applied to q, k
|
||||
/// before the scores, matching HF `apply_rotary_pos_emb_vision`
|
||||
/// (`rope_slow` is the same rotate-half form).
|
||||
fn attention(&self, x: &Tensor, rotary: Option<&(Tensor, Tensor)>) -> Result<Tensor> {
|
||||
let (n, hidden) = x.dims2()?;
|
||||
// qkv: (N, 3*hidden). Split into Q, K, V each (N, hidden).
|
||||
let qkv = self.qkv.forward(x)?;
|
||||
@@ -140,6 +170,15 @@ impl VisionBlock {
|
||||
let q = qkv.i(0)?;
|
||||
let k = qkv.i(1)?;
|
||||
let v = qkv.i(2)?;
|
||||
// 2D vision rotary on q, k (full head_dim; rotate-half form).
|
||||
let (q, k) = match rotary {
|
||||
Some((cos, sin)) => {
|
||||
let q = candle_nn::rotary_emb::rope_slow(&q.unsqueeze(0)?, cos, sin)?.squeeze(0)?;
|
||||
let k = candle_nn::rotary_emb::rope_slow(&k.unsqueeze(0)?, cos, sin)?.squeeze(0)?;
|
||||
(q, k)
|
||||
}
|
||||
None => (q, k),
|
||||
};
|
||||
let scale = 1.0 / (self.head_dim as f64).sqrt();
|
||||
// (num_heads, N, head_dim) @ (num_heads, head_dim, N) -> (num_heads, N, N)
|
||||
let scores = q.matmul(&k.transpose(D::Minus2, D::Minus1)?)?;
|
||||
@@ -210,11 +249,65 @@ impl VisionMerger {
|
||||
}
|
||||
}
|
||||
|
||||
/// 2D rotary position embedding for the vision tower. Each patch's
|
||||
/// `head_dim` rotates by its `(row, col)` grid coordinates: the first
|
||||
/// half of the rotary freqs are driven by the row position, the second
|
||||
/// half by the column. Mirrors HF `Qwen3VLVisionRotaryEmbedding` +
|
||||
/// `rot_pos_emb` (θ = 10000, `dim = head_dim/2`).
|
||||
struct VisionRotaryEmbedding {
|
||||
/// `(half,)` f32, `half = head_dim/4` freqs per spatial axis.
|
||||
inv_freq: Vec<f32>,
|
||||
}
|
||||
|
||||
impl VisionRotaryEmbedding {
|
||||
fn new(head_dim: usize) -> Self {
|
||||
// HF: Qwen3VLVisionRotaryEmbedding(head_dim // 2), theta 10000.
|
||||
let dim = head_dim / 2;
|
||||
let theta = 10000f32;
|
||||
let inv_freq = (0..dim)
|
||||
.step_by(2)
|
||||
.map(|i| 1f32 / theta.powf(i as f32 / dim as f32))
|
||||
.collect();
|
||||
Self { inv_freq }
|
||||
}
|
||||
|
||||
/// cos/sin for a `gh×gw` patch grid in **row-major** order. Returns
|
||||
/// `(cos, sin)` each `(gh*gw, head_dim/2)`: per patch, the row-axis
|
||||
/// freqs `row·inv_freq` followed by the col-axis freqs `col·inv_freq`
|
||||
/// (then `rope_slow` duplicates them across the full head_dim).
|
||||
fn cos_sin(
|
||||
&self,
|
||||
gh: usize,
|
||||
gw: usize,
|
||||
dev: &Device,
|
||||
dtype: DType,
|
||||
) -> candle_core::Result<(Tensor, Tensor)> {
|
||||
let half = self.inv_freq.len();
|
||||
let n = gh * gw;
|
||||
let mut data = Vec::with_capacity(n * 2 * half);
|
||||
for hi in 0..gh {
|
||||
for wi in 0..gw {
|
||||
for &f in &self.inv_freq {
|
||||
data.push(hi as f32 * f);
|
||||
}
|
||||
for &f in &self.inv_freq {
|
||||
data.push(wi as f32 * f);
|
||||
}
|
||||
}
|
||||
}
|
||||
let freqs = Tensor::from_vec(data, (n, 2 * half), dev)?;
|
||||
let cos = freqs.cos()?.to_dtype(dtype)?;
|
||||
let sin = freqs.sin()?.to_dtype(dtype)?;
|
||||
Ok((cos, sin))
|
||||
}
|
||||
}
|
||||
|
||||
/// The vision tower itself.
|
||||
pub struct VisionTower {
|
||||
/// Sum-collapsed temporal kernel (Conv2d, see module doc).
|
||||
patch_embed: Conv2d,
|
||||
pos_embed: Embedding,
|
||||
rotary: VisionRotaryEmbedding,
|
||||
blocks: Vec<VisionBlock>,
|
||||
merger: VisionMerger,
|
||||
config: VisionConfig,
|
||||
@@ -265,6 +358,7 @@ impl VisionTower {
|
||||
.get((cfg.num_position_embeddings, cfg.hidden_size), "weight")
|
||||
.context("load model.visual.pos_embed.weight")?;
|
||||
let pos_embed = Embedding::new(pos_embed_weight, cfg.hidden_size);
|
||||
let rotary = VisionRotaryEmbedding::new(cfg.hidden_size / cfg.num_heads);
|
||||
|
||||
let blocks_vb = vb.pp("blocks");
|
||||
let mut blocks = Vec::with_capacity(cfg.depth);
|
||||
@@ -279,6 +373,7 @@ impl VisionTower {
|
||||
Ok(Self {
|
||||
patch_embed,
|
||||
pos_embed,
|
||||
rotary,
|
||||
blocks,
|
||||
merger,
|
||||
config: cfg,
|
||||
@@ -302,6 +397,81 @@ impl VisionTower {
|
||||
gh * gw * LM_TOKENS_PER_MERGE_GROUP
|
||||
}
|
||||
|
||||
/// Bilinearly interpolate the learned `pos_embed` grid (a
|
||||
/// `num_grid_per_side × num_grid_per_side` table, 48×48 for Qwen3.6)
|
||||
/// onto the actual `gh × gw` patch grid, in **row-major** patch
|
||||
/// order. Port of the HF `fast_pos_embed_interpolate`: for each patch
|
||||
/// at fractional grid coord `(linspace(0, ngrid-1, gh)[hi],
|
||||
/// linspace(0, ngrid-1, gw)[wi])`, blend the 4 surrounding grid
|
||||
/// entries by bilinear weights. Returns `(gh*gw, hidden)` in
|
||||
/// `self.dtype`.
|
||||
fn interpolated_pos_embed(&self, gh: usize, gw: usize) -> Result<Tensor> {
|
||||
let ngrid = (self.config.num_position_embeddings as f64).sqrt().round() as usize;
|
||||
anyhow::ensure!(
|
||||
ngrid * ngrid == self.config.num_position_embeddings,
|
||||
"num_position_embeddings {} is not a perfect square",
|
||||
self.config.num_position_embeddings
|
||||
);
|
||||
// Evenly-spaced fractional indices into the [0, ngrid-1] grid.
|
||||
let lin = |n: usize| -> Vec<f64> {
|
||||
if n <= 1 {
|
||||
vec![0.0]
|
||||
} else {
|
||||
let step = (ngrid - 1) as f64 / (n - 1) as f64;
|
||||
(0..n).map(|i| i as f64 * step).collect()
|
||||
}
|
||||
};
|
||||
let hs = lin(gh);
|
||||
let ws = lin(gw);
|
||||
let n = gh * gw;
|
||||
|
||||
// Four corner index sets + bilinear weight sets, row-major.
|
||||
let mut idx: [Vec<u32>; 4] = [
|
||||
Vec::with_capacity(n),
|
||||
Vec::with_capacity(n),
|
||||
Vec::with_capacity(n),
|
||||
Vec::with_capacity(n),
|
||||
];
|
||||
let mut wts: [Vec<f32>; 4] = [
|
||||
Vec::with_capacity(n),
|
||||
Vec::with_capacity(n),
|
||||
Vec::with_capacity(n),
|
||||
Vec::with_capacity(n),
|
||||
];
|
||||
for &hv in &hs {
|
||||
let hf = hv as usize; // floor (hv >= 0)
|
||||
let hc = (hf + 1).min(ngrid - 1);
|
||||
let dh = (hv - hf as f64) as f32;
|
||||
for &wv in &ws {
|
||||
let wf = wv as usize;
|
||||
let wc = (wf + 1).min(ngrid - 1);
|
||||
let dw = (wv - wf as f64) as f32;
|
||||
idx[0].push((hf * ngrid + wf) as u32);
|
||||
wts[0].push((1.0 - dh) * (1.0 - dw));
|
||||
idx[1].push((hf * ngrid + wc) as u32);
|
||||
wts[1].push((1.0 - dh) * dw);
|
||||
idx[2].push((hc * ngrid + wf) as u32);
|
||||
wts[2].push(dh * (1.0 - dw));
|
||||
idx[3].push((hc * ngrid + wc) as u32);
|
||||
wts[3].push(dh * dw);
|
||||
}
|
||||
}
|
||||
|
||||
let mut acc: Option<Tensor> = None;
|
||||
for corner in 0..4 {
|
||||
let idx_t = Tensor::from_vec(std::mem::take(&mut idx[corner]), (n,), &self.device)?;
|
||||
let emb = self.pos_embed.forward(&idx_t)?; // (n, hidden), pos_embed dtype
|
||||
let wt = Tensor::from_vec(std::mem::take(&mut wts[corner]), (n, 1), &self.device)?
|
||||
.to_dtype(self.dtype)?;
|
||||
let term = emb.broadcast_mul(&wt)?;
|
||||
acc = Some(match acc {
|
||||
Some(a) => a.add(&term)?,
|
||||
None => term,
|
||||
});
|
||||
}
|
||||
Ok(acc.expect("4 corners accumulated"))
|
||||
}
|
||||
|
||||
/// Encode one image.
|
||||
///
|
||||
/// `image`: row-major `(3, H, W)` f32 tensor on `self.device`,
|
||||
@@ -339,16 +509,34 @@ impl VisionTower {
|
||||
let x = x.permute((1, 2, 0))?.contiguous()?;
|
||||
let x = x.reshape((n_patches, self.config.hidden_size))?;
|
||||
|
||||
// Add learned positional embeddings (sequential indices for
|
||||
// Stage A's fixed-resolution path; full 2D positional logic
|
||||
// lands with variable resolution, issue #14).
|
||||
let positions = Tensor::arange(0u32, n_patches as u32, &self.device)?;
|
||||
let pos = self.pos_embed.forward(&positions)?;
|
||||
// Learned absolute position embeddings. The `pos_embed` table is
|
||||
// a `num_position_embeddings = num_grid_per_side²` learned grid
|
||||
// (48×48 for Qwen3.6); for a `gh×gw` patch grid the reference
|
||||
// (`fast_pos_embed_interpolate`) bilinearly interpolates that
|
||||
// grid to `gh×gw`. The legacy path (a naive sequential lookup of
|
||||
// the first `n_patches` rows) mis-maps the grid stride and
|
||||
// scrambles spatial structure — kept only behind
|
||||
// `NEURON_VISION_LEGACY_POS=1` for A/B comparison.
|
||||
let pos = if vision_legacy_pos() {
|
||||
let positions = Tensor::arange(0u32, n_patches as u32, &self.device)?;
|
||||
self.pos_embed.forward(&positions)?
|
||||
} else {
|
||||
self.interpolated_pos_embed(gh, gw)?
|
||||
};
|
||||
let mut x = x.add(&pos)?;
|
||||
|
||||
// 2D vision rotary (row/col per patch), computed once and applied
|
||||
// in every block's attention. Legacy escape hatch skips it.
|
||||
let rotary = if vision_legacy_rope() {
|
||||
None
|
||||
} else {
|
||||
Some(self.rotary.cos_sin(gh, gw, &self.device, self.dtype)?)
|
||||
};
|
||||
let rotary_ref = rotary.as_ref();
|
||||
|
||||
for (i, block) in self.blocks.iter().enumerate() {
|
||||
x = block
|
||||
.forward(&x)
|
||||
.forward(&x, rotary_ref)
|
||||
.with_context(|| format!("vision block {i}"))?;
|
||||
}
|
||||
|
||||
@@ -516,9 +704,11 @@ mod tests {
|
||||
spatial_merge_size: cfg.spatial_merge_size,
|
||||
};
|
||||
|
||||
let rotary = VisionRotaryEmbedding::new(cfg.hidden_size / cfg.num_heads);
|
||||
VisionTower {
|
||||
patch_embed,
|
||||
pos_embed,
|
||||
rotary,
|
||||
blocks,
|
||||
merger,
|
||||
config: cfg.clone(),
|
||||
@@ -548,6 +738,51 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn interpolated_pos_embed_reduces_to_sequential_at_native_grid() {
|
||||
// When the patch grid equals the pos_embed grid (gh=gw=ngrid),
|
||||
// linspace(0,ngrid-1,ngrid) is the integer ladder, so every patch
|
||||
// lands exactly on a grid node (dh=dw=0, corner-0 weight 1) and
|
||||
// the bilinear result is the raw pos_embed rows in row-major
|
||||
// order — i.e. identical to the legacy sequential lookup.
|
||||
let cfg = tiny_config();
|
||||
let tower = tiny_tower(&cfg);
|
||||
let ngrid = (cfg.num_position_embeddings as f64).sqrt() as usize; // 8
|
||||
let interp = tower.interpolated_pos_embed(ngrid, ngrid).unwrap();
|
||||
let seq = tower
|
||||
.pos_embed
|
||||
.forward(&Tensor::arange(0u32, (ngrid * ngrid) as u32, &Device::Cpu).unwrap())
|
||||
.unwrap();
|
||||
let a: Vec<f32> = interp.flatten_all().unwrap().to_vec1().unwrap();
|
||||
let b: Vec<f32> = seq.flatten_all().unwrap().to_vec1().unwrap();
|
||||
assert_eq!(a.len(), b.len());
|
||||
for (x, y) in a.iter().zip(b.iter()) {
|
||||
assert!((x - y).abs() < 1e-5, "interp {x} vs seq {y}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn vision_rotary_row_col_structure() {
|
||||
// head_dim 8 → rotary dim 4 → inv_freq over [0,2] → 2 freqs/axis.
|
||||
let rot = VisionRotaryEmbedding::new(8);
|
||||
assert_eq!(rot.inv_freq.len(), 2);
|
||||
let (cos, sin) = rot.cos_sin(2, 2, &Device::Cpu, DType::F32).unwrap();
|
||||
assert_eq!(cos.dims(), &[4, 4]); // 4 patches, head_dim/2 = 4 cols
|
||||
|
||||
// Patch (0,0): all freqs 0 → cos 1, sin 0.
|
||||
let s0: Vec<f32> = sin.i(0).unwrap().to_vec1().unwrap();
|
||||
assert!(s0.iter().all(|&s| s.abs() < 1e-6));
|
||||
|
||||
// Patch index 2 = grid (1,0): row=1 drives the first half, col=0
|
||||
// leaves the second half at zero.
|
||||
let s2: Vec<f32> = sin.i(2).unwrap().to_vec1().unwrap();
|
||||
assert!(s2[0].abs() > 1e-6, "row half must be non-zero");
|
||||
assert!(
|
||||
s2[2].abs() < 1e-6 && s2[3].abs() < 1e-6,
|
||||
"col half must be zero"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lm_token_count_matches_grid() {
|
||||
let cfg = tiny_config();
|
||||
|
||||
@@ -60,6 +60,17 @@ pub struct CandleHarness {
|
||||
/// can still load on CPU for tests, just without worker threads).
|
||||
#[allow(dead_code)]
|
||||
device_workers: Arc<RwLock<HashMap<u32, Arc<super::device_worker::DeviceWorkerHandle>>>>,
|
||||
/// Auto-recovery (#17): model ids whose poisoned context is being
|
||||
/// rebuilt via unload+reload. Insert is the single-flight gate (one
|
||||
/// recovery per model in flight); membership also lets the request
|
||||
/// path answer "recovering, retry shortly" during the reload gap
|
||||
/// rather than a bare "not loaded".
|
||||
recovering: Arc<RwLock<std::collections::HashSet<String>>>,
|
||||
/// Sender to the background recovery task. The request path enqueues
|
||||
/// a poisoned model id here; the task (holding a `Weak<Self>`) runs
|
||||
/// the unload→reload→health-gate. Unbounded + tiny (model ids), and
|
||||
/// the `recovering` set dedupes, so it can't back up.
|
||||
recovery_tx: tokio::sync::mpsc::UnboundedSender<String>,
|
||||
}
|
||||
|
||||
/// One entry in the harness's loaded-model registry. Single-GPU loads
|
||||
@@ -86,6 +97,15 @@ impl LoadedHandle {
|
||||
}
|
||||
}
|
||||
|
||||
/// The spec this model was loaded from (for auto-recovery #17).
|
||||
pub fn spec(&self) -> &ModelSpec {
|
||||
match self {
|
||||
LoadedHandle::Single(m) => &m.spec,
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => &m.spec,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn devices(&self) -> Vec<u32> {
|
||||
match self {
|
||||
LoadedHandle::Single(m) => m.devices.clone(),
|
||||
@@ -106,18 +126,18 @@ impl LoadedHandle {
|
||||
}
|
||||
}
|
||||
|
||||
/// Modalities the loaded model supports. Stage B7. TP models are
|
||||
/// always text-only today — TP-vision is tracked under issue #12.
|
||||
/// Modalities the loaded model supports. Stage B7 (single-GPU) +
|
||||
/// TP-vision (#12) — both single-GPU and TP loads advertise
|
||||
/// `"vision"` when a replicated vision tower materialised.
|
||||
pub fn capabilities(&self) -> Vec<String> {
|
||||
let mut caps = vec!["text".to_string()];
|
||||
match self {
|
||||
LoadedHandle::Single(m) => {
|
||||
if m.has_vision {
|
||||
caps.push("vision".to_string());
|
||||
}
|
||||
}
|
||||
let has_vision = match self {
|
||||
LoadedHandle::Single(m) => m.has_vision,
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(_) => {}
|
||||
LoadedHandle::Tp(m) => m.has_vision,
|
||||
};
|
||||
if has_vision {
|
||||
caps.push("vision".to_string());
|
||||
}
|
||||
caps
|
||||
}
|
||||
@@ -210,13 +230,15 @@ pub struct LoadedModel {
|
||||
/// targets and the worker forward uses it to locate splice
|
||||
/// positions in the LM input embeddings.
|
||||
pub image_token_id: Option<u32>,
|
||||
/// LM-side tokens this model's vision tower emits per image at
|
||||
/// the Stage B fixed resolution (448×448 → 196 for Qwen3.6).
|
||||
/// `None` for text-only models. Set at load time so the
|
||||
/// hot path doesn't recompute it per request. Stage B fixed
|
||||
/// resolution → constant; dynamic resolution per #14 makes it
|
||||
/// per-image.
|
||||
pub lm_tokens_per_image: Option<usize>,
|
||||
/// `patch_size × spatial_merge_size` — divides a resized pixel
|
||||
/// dimension into LM-grid units. Per-image LM token count is
|
||||
/// `(h/factor) × (w/factor)` (#14 dynamic resolution). `None` for
|
||||
/// text-only models. Set at load time.
|
||||
pub image_grid_factor: Option<usize>,
|
||||
/// The spec this model was loaded from — retained so auto-recovery
|
||||
/// (#17) can `unload_model` + `load_model(spec)` a poisoned model
|
||||
/// without an operator reconstructing it.
|
||||
pub spec: ModelSpec,
|
||||
}
|
||||
|
||||
impl LoadedModel {
|
||||
@@ -281,6 +303,19 @@ pub struct TpLoadedModel {
|
||||
pub tool_call_tokens: Option<ToolCallTokenPair>,
|
||||
/// Same shape as [`LoadedModel::chat_template`].
|
||||
pub chat_template: Option<String>,
|
||||
/// Vision capability flag (TP-vision). `true` iff every rank
|
||||
/// materialised a replicated vision tower. Mirrors
|
||||
/// [`LoadedModel::has_vision`]; drives capability advertising and
|
||||
/// the TP vision dispatch.
|
||||
pub has_vision: bool,
|
||||
/// `<|image_pad|>` token id — same as [`LoadedModel::image_token_id`].
|
||||
pub image_token_id: Option<u32>,
|
||||
/// Pixel→LM-grid divisor — same as
|
||||
/// [`LoadedModel::image_grid_factor`].
|
||||
pub image_grid_factor: Option<usize>,
|
||||
/// Loading spec, retained for auto-recovery (#17) — see
|
||||
/// [`LoadedModel::spec`].
|
||||
pub spec: ModelSpec,
|
||||
}
|
||||
|
||||
#[cfg(feature = "cuda")]
|
||||
@@ -384,10 +419,11 @@ impl ModelArch {
|
||||
offset: usize,
|
||||
image_embeds: &Tensor,
|
||||
image_token_id: u32,
|
||||
grids: &[(usize, usize)],
|
||||
) -> Result<Tensor> {
|
||||
let raw = match self {
|
||||
ModelArch::Qwen3_5Dense(m) => {
|
||||
m.forward_with_vision(input, offset, image_embeds, image_token_id)?
|
||||
m.forward_with_vision(input, offset, image_embeds, image_token_id, grids)?
|
||||
}
|
||||
other => anyhow::bail!(
|
||||
"forward_with_vision: architecture {} has no vision tower",
|
||||
@@ -397,6 +433,20 @@ impl ModelArch {
|
||||
squeeze_to_vocab(&raw)
|
||||
}
|
||||
|
||||
/// `patch_size × spatial_merge_size` for the loaded vision tower —
|
||||
/// divides a resized pixel dim into LM-grid units (an image of
|
||||
/// resized `(h, w)` yields the LM grid `(h/factor, w/factor)`).
|
||||
/// `None` for architectures/checkpoints without a vision tower.
|
||||
pub fn vision_grid_factor(&self) -> Option<usize> {
|
||||
match self {
|
||||
ModelArch::Qwen3_5Dense(m) => m.vision().map(|v| {
|
||||
let c = v.config();
|
||||
c.patch_size * c.spatial_merge_size
|
||||
}),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Encode a preprocessed image into LM-side token embeddings via
|
||||
/// the loaded vision tower. Stage A5.
|
||||
///
|
||||
@@ -769,6 +819,46 @@ fn poisoned_error(model_id: &str) -> InferenceError {
|
||||
))
|
||||
}
|
||||
|
||||
/// Reported while auto-recovery (#17) is rebuilding a poisoned model's
|
||||
/// context. Unlike [`poisoned_error`] this is a *transient* state — the
|
||||
/// model is being reloaded automatically; the client should retry.
|
||||
fn recovering_error(model_id: &str) -> InferenceError {
|
||||
InferenceError::Other(anyhow::anyhow!(
|
||||
"model '{model_id}' is recovering (its device context was poisoned \
|
||||
by an earlier failure and is being automatically rebuilt); retry \
|
||||
shortly"
|
||||
))
|
||||
}
|
||||
|
||||
/// Verification hook for #17 auto-recovery. When `NEURON_DEBUG_POISON`
|
||||
/// names a model, the **first** request for it (process-wide) returns
|
||||
/// true, so the request path can trigger recovery as if a device fault
|
||||
/// had occurred — exercising the unload→reload→healthy cycle without
|
||||
/// corrupting the GPU. One-shot (a `swap` latch) so it can't loop the
|
||||
/// model through endless recoveries. No-op unless the env var is set.
|
||||
fn debug_poison_armed(model_id: &str) -> bool {
|
||||
static FIRED: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
|
||||
let armed = std::env::var("NEURON_DEBUG_POISON").ok().as_deref() == Some(model_id);
|
||||
armed && !FIRED.swap(true, Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Background auto-recovery task (#17). Drains poisoned model ids and
|
||||
/// rebuilds each via [`CandleHarness::recover_one`]. Holds a `Weak` so a
|
||||
/// shutting-down harness lets the task exit; processes one id at a time,
|
||||
/// which (with the `recovering` set deduping enqueues) keeps recovery
|
||||
/// single-flight per model.
|
||||
async fn recovery_loop(
|
||||
weak: std::sync::Weak<CandleHarness>,
|
||||
mut rx: tokio::sync::mpsc::UnboundedReceiver<String>,
|
||||
) {
|
||||
while let Some(model_id) = rx.recv().await {
|
||||
let Some(this) = weak.upgrade() else {
|
||||
break;
|
||||
};
|
||||
this.recover_one(&model_id).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Free/total VRAM on the candle `Device` in MiB. Returns `(0, 0)` if
|
||||
/// the query fails or the device is the CPU fallback so logging never
|
||||
/// crashes the request path. Mirrors the existing helper in
|
||||
@@ -861,6 +951,45 @@ fn min_free_vram_mb() -> u64 {
|
||||
/// prefill. Called from every chat_completion entry point right after
|
||||
/// the VRAM query. A `prompt_len == 0` is accepted (some clients send
|
||||
/// empty inputs to probe the endpoint); the prefill loop handles it.
|
||||
/// Rough MiB of VRAM a vision prefill needs per 1000 prompt tokens
|
||||
/// (accumulating KV cache + per-chunk activation headroom). Tunable;
|
||||
/// the default is deliberately permissive so the guard rejects only
|
||||
/// clearly-too-large requests, not ones the chunked prefill handles.
|
||||
fn vision_prefill_mb_per_1k_tokens() -> u64 {
|
||||
env_u64("NEURON_VISION_PREFILL_MB_PER_1K_TOKENS", 500)
|
||||
}
|
||||
|
||||
/// Fixed VRAM overhead (MiB) a vision prefill reserves on top of the
|
||||
/// per-token estimate — image encode buffers + one chunk's activations.
|
||||
fn vision_prefill_base_mb() -> u64 {
|
||||
env_u64("NEURON_VISION_PREFILL_BASE_MB", 2000)
|
||||
}
|
||||
|
||||
/// Pre-flight check specific to vision prefills. Even with the chunked
|
||||
/// prefill bounding per-step activation, the accumulating KV cache for
|
||||
/// a long prompt can exhaust VRAM mid-forward — and on the TP path a
|
||||
/// mid-forward OOM strands the NCCL collective (one rank dies, the other
|
||||
/// hangs on the all-reduce, holding the pool lock). Reject up front with
|
||||
/// a clean `InsufficientVram` when the estimated footprint exceeds free
|
||||
/// VRAM, so a doomed request fails fast instead of hanging the daemon.
|
||||
///
|
||||
/// Heuristic and tunable (`NEURON_VISION_PREFILL_*`); the default errs
|
||||
/// permissive. Skipped on the CPU sentinel (`vram_free_mb == 0`).
|
||||
fn validate_vision_prefill(prompt_len: usize, vram_free_mb: u64) -> Result<(), InferenceError> {
|
||||
if vram_free_mb == 0 {
|
||||
return Ok(());
|
||||
}
|
||||
let required_mb = vision_prefill_base_mb()
|
||||
+ (prompt_len as u64).saturating_mul(vision_prefill_mb_per_1k_tokens()) / 1000;
|
||||
if required_mb > vram_free_mb {
|
||||
return Err(InferenceError::InsufficientVram {
|
||||
free_mb: vram_free_mb,
|
||||
required_mb,
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_request(prompt_len: usize, vram_free_mb: u64) -> Result<(), InferenceError> {
|
||||
let max = max_prompt_tokens();
|
||||
if prompt_len > max {
|
||||
@@ -1084,7 +1213,7 @@ impl CandleHarness {
|
||||
/// Construct a new harness for `bind_url` using `config`. Resolves
|
||||
/// every configured source's auth env var and cache dir up front so
|
||||
/// the hot load path (`hf_api_for`) is a pure HashMap lookup.
|
||||
pub fn new(bind_url: String, config: &crate::config::CandleHarnessConfig) -> Self {
|
||||
pub fn new(bind_url: String, config: &crate::config::CandleHarnessConfig) -> Arc<Self> {
|
||||
let raw_sources = config.effective_sources();
|
||||
let default_source = config.effective_default_source().to_string();
|
||||
let mut sources = HashMap::with_capacity(raw_sources.len());
|
||||
@@ -1134,13 +1263,25 @@ impl CandleHarness {
|
||||
bare model ids will fail to resolve until this is fixed"
|
||||
);
|
||||
}
|
||||
Self {
|
||||
let (recovery_tx, recovery_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
|
||||
let this = Arc::new(Self {
|
||||
models: Arc::new(RwLock::new(HashMap::new())),
|
||||
sources,
|
||||
default_source,
|
||||
bind_url,
|
||||
device_workers: Arc::new(RwLock::new(HashMap::new())),
|
||||
recovering: Arc::new(RwLock::new(std::collections::HashSet::new())),
|
||||
recovery_tx,
|
||||
});
|
||||
// Background auto-recovery task (#17). Holds a `Weak` so it can't
|
||||
// keep the harness alive. Spawned only when a tokio runtime is
|
||||
// present — sync unit tests that build a harness without one
|
||||
// simply skip it (they don't exercise recovery).
|
||||
if tokio::runtime::Handle::try_current().is_ok() {
|
||||
let weak = Arc::downgrade(&this);
|
||||
tokio::spawn(recovery_loop(weak, recovery_rx));
|
||||
}
|
||||
this
|
||||
}
|
||||
|
||||
/// Scheme to substitute for bare `org/name` model ids. Mirrors the
|
||||
@@ -1565,7 +1706,17 @@ impl CandleHarness {
|
||||
let models = self.models.read().await;
|
||||
models.get(&request.model).cloned()
|
||||
};
|
||||
let handle = handle.ok_or_else(|| InferenceError::ModelNotLoaded(request.model.clone()))?;
|
||||
let handle = match handle {
|
||||
Some(h) => h,
|
||||
// Absent from the registry: distinguish a genuinely unloaded
|
||||
// model from one whose slot is briefly gone mid auto-recovery
|
||||
// (#17), so the client gets a transient "retry shortly" instead
|
||||
// of a misleading "not loaded".
|
||||
None if self.is_recovering(&request.model).await => {
|
||||
return Err(recovering_error(&request.model));
|
||||
}
|
||||
None => return Err(InferenceError::ModelNotLoaded(request.model.clone())),
|
||||
};
|
||||
// The match is technically infallible without `cuda` (only Single
|
||||
// exists), but the cfg-gated Tp arm makes this the right shape
|
||||
// under both feature flags.
|
||||
@@ -1595,7 +1746,12 @@ impl CandleHarness {
|
||||
if loaded.poisoned.load(Ordering::Acquire) {
|
||||
let _g = span.enter();
|
||||
tracing::warn!("chat_completion: refusing request, model poisoned");
|
||||
return Err(poisoned_error(&model_id));
|
||||
return Err(self.trigger_recovery(&model_id).await);
|
||||
}
|
||||
if debug_poison_armed(&model_id) {
|
||||
let _g = span.enter();
|
||||
tracing::warn!("NEURON_DEBUG_POISON: forcing auto-recovery (#17 verification)");
|
||||
return Err(self.trigger_recovery(&model_id).await);
|
||||
}
|
||||
|
||||
// Serialise concurrent requests against this model. Holds for
|
||||
@@ -1634,11 +1790,11 @@ impl CandleHarness {
|
||||
.ok_or_else(|| InferenceError::VisionUnsupported {
|
||||
model_id: request.model.clone(),
|
||||
})?;
|
||||
let patches_per_image = loaded
|
||||
.lm_tokens_per_image
|
||||
.ok_or_else(|| InferenceError::VisionUnsupported {
|
||||
let factor = loaded.image_grid_factor.ok_or_else(|| {
|
||||
InferenceError::VisionUnsupported {
|
||||
model_id: request.model.clone(),
|
||||
})?;
|
||||
}
|
||||
})?;
|
||||
let profile = super::preprocess::PreprocessProfile::qwen3_6();
|
||||
let images = extract_images_from_request(&request, &profile).map_err(|e| {
|
||||
InferenceError::Other(anyhow::anyhow!("extract_images: {e}"))
|
||||
@@ -1650,7 +1806,12 @@ impl CandleHarness {
|
||||
"request has image content but extractor produced zero images"
|
||||
)));
|
||||
}
|
||||
let per_image_counts: Vec<usize> = vec![patches_per_image; images.len()];
|
||||
// Per-image LM token count from each image's resized grid
|
||||
// (#14 dynamic resolution; was a constant 196).
|
||||
let per_image_counts: Vec<usize> = images
|
||||
.iter()
|
||||
.map(|im| (im.h / factor) * (im.w / factor))
|
||||
.collect();
|
||||
prompt_tokens =
|
||||
expand_image_pad_tokens(&prompt_tokens, image_token_id, &per_image_counts)
|
||||
.map_err(InferenceError::Other)?;
|
||||
@@ -1684,6 +1845,12 @@ impl CandleHarness {
|
||||
);
|
||||
|
||||
validate_request(prompt_len, vram_free_mb)?;
|
||||
if vision_route.is_some() {
|
||||
validate_vision_prefill(prompt_len, vram_free_mb)?;
|
||||
}
|
||||
if vision_route.is_some() {
|
||||
validate_vision_prefill(prompt_len, vram_free_mb)?;
|
||||
}
|
||||
|
||||
// Routing: CUDA loads go through the per-device worker
|
||||
// thread (introduced in Phase 1; forward/clear added in
|
||||
@@ -1963,7 +2130,17 @@ impl CandleHarness {
|
||||
let models = self.models.read().await;
|
||||
models.get(&request.model).cloned()
|
||||
};
|
||||
let handle = handle.ok_or_else(|| InferenceError::ModelNotLoaded(request.model.clone()))?;
|
||||
let handle = match handle {
|
||||
Some(h) => h,
|
||||
// Absent from the registry: distinguish a genuinely unloaded
|
||||
// model from one whose slot is briefly gone mid auto-recovery
|
||||
// (#17), so the client gets a transient "retry shortly" instead
|
||||
// of a misleading "not loaded".
|
||||
None if self.is_recovering(&request.model).await => {
|
||||
return Err(recovering_error(&request.model));
|
||||
}
|
||||
None => return Err(InferenceError::ModelNotLoaded(request.model.clone())),
|
||||
};
|
||||
// The match is technically infallible without `cuda` (only Single
|
||||
// exists), but the cfg-gated Tp arm makes this the right shape
|
||||
// under both feature flags.
|
||||
@@ -1981,7 +2158,55 @@ impl CandleHarness {
|
||||
.tokenizer
|
||||
.encode(prompt.as_str(), true)
|
||||
.map_err(|e| InferenceError::Other(anyhow::anyhow!("tokenize: {e}")))?;
|
||||
let prompt_tokens: Vec<u32> = encoding.get_ids().to_vec();
|
||||
let mut prompt_tokens: Vec<u32> = encoding.get_ids().to_vec();
|
||||
|
||||
// Stage C1: vision routing for the streaming path. Mirrors the
|
||||
// non-streaming `chat_completion` block — detect image content,
|
||||
// reject it against text-only models, preprocess each image and
|
||||
// expand its `<|image_pad|>` sentinel to the per-image patch
|
||||
// count, then carry the payload through to a single-shot
|
||||
// image-spliced prefill. Non-image requests skip all of this.
|
||||
// Returning early here (before the `Start` event below) keeps a
|
||||
// rejected vision request from opening a half-formed SSE stream.
|
||||
let vision_route: Option<(Vec<super::device_worker::jobs::ImageInput>, u32)> =
|
||||
if request_has_images(&request) {
|
||||
if !loaded.has_vision {
|
||||
return Err(InferenceError::VisionUnsupported {
|
||||
model_id: request.model.clone(),
|
||||
});
|
||||
}
|
||||
let image_token_id =
|
||||
loaded
|
||||
.image_token_id
|
||||
.ok_or_else(|| InferenceError::VisionUnsupported {
|
||||
model_id: request.model.clone(),
|
||||
})?;
|
||||
let factor =
|
||||
loaded
|
||||
.image_grid_factor
|
||||
.ok_or_else(|| InferenceError::VisionUnsupported {
|
||||
model_id: request.model.clone(),
|
||||
})?;
|
||||
let profile = super::preprocess::PreprocessProfile::qwen3_6();
|
||||
let images = extract_images_from_request(&request, &profile)
|
||||
.map_err(|e| InferenceError::Other(anyhow::anyhow!("extract_images: {e}")))?;
|
||||
if images.is_empty() {
|
||||
return Err(InferenceError::Other(anyhow::anyhow!(
|
||||
"request has image content but extractor produced zero images"
|
||||
)));
|
||||
}
|
||||
// Per-image LM token count from each image's resized grid (#14).
|
||||
let per_image_counts: Vec<usize> = images
|
||||
.iter()
|
||||
.map(|im| (im.h / factor) * (im.w / factor))
|
||||
.collect();
|
||||
prompt_tokens =
|
||||
expand_image_pad_tokens(&prompt_tokens, image_token_id, &per_image_counts)
|
||||
.map_err(InferenceError::Other)?;
|
||||
Some((images, image_token_id))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let temperature = request.temperature.unwrap_or(0.7);
|
||||
let top_p = request.top_p;
|
||||
@@ -2008,7 +2233,7 @@ impl CandleHarness {
|
||||
// Refuse if the model is already poisoned. No point opening
|
||||
// an SSE stream just to send the Start event and then bail.
|
||||
if loaded.poisoned.load(Ordering::Acquire) {
|
||||
return Err(poisoned_error(&model_id));
|
||||
return Err(self.trigger_recovery(&model_id).await);
|
||||
}
|
||||
|
||||
// Start event: tells the wire projector to emit its
|
||||
@@ -2048,11 +2273,15 @@ impl CandleHarness {
|
||||
?eos_id,
|
||||
vram_free_mb,
|
||||
vram_total_mb,
|
||||
vision = vision_route.is_some(),
|
||||
"chat_completion (stream): starting"
|
||||
);
|
||||
}
|
||||
|
||||
validate_request(prompt_len, vram_free_mb)?;
|
||||
if vision_route.is_some() {
|
||||
validate_vision_prefill(prompt_len, vram_free_mb)?;
|
||||
}
|
||||
|
||||
// Routing parallel to the non-streaming chat_completion: CUDA
|
||||
// goes through the worker (async task), CPU keeps the
|
||||
@@ -2078,6 +2307,7 @@ impl CandleHarness {
|
||||
handle,
|
||||
tokenizer,
|
||||
prompt_tokens,
|
||||
vision_route,
|
||||
max_new,
|
||||
temperature,
|
||||
top_p,
|
||||
@@ -2221,6 +2451,69 @@ pub struct InferenceStream {
|
||||
pub reasoning_markers: Option<ReasoningTokenPair>,
|
||||
}
|
||||
|
||||
/// Auto-recovery (#17) — rebuild a poisoned model's device context
|
||||
/// automatically instead of leaving it bricked until a human reloads.
|
||||
impl CandleHarness {
|
||||
/// True while `model_id` is being auto-recovered (its slot is briefly
|
||||
/// absent from the registry during the reload).
|
||||
pub async fn is_recovering(&self, model_id: &str) -> bool {
|
||||
self.recovering.read().await.contains(model_id)
|
||||
}
|
||||
|
||||
/// Single-flight trigger from the request path: enqueue a rebuild for a
|
||||
/// poisoned model (only the first caller per model enqueues) and return
|
||||
/// the transient "recovering" error to hand back to the client.
|
||||
async fn trigger_recovery(&self, model_id: &str) -> InferenceError {
|
||||
let newly = self.recovering.write().await.insert(model_id.to_string());
|
||||
if newly {
|
||||
tracing::warn!(model = %model_id, "auto-recovery: poisoned, enqueueing rebuild");
|
||||
if self.recovery_tx.send(model_id.to_string()).is_err() {
|
||||
// Background task gone (harness shutting down). Drop the
|
||||
// marker and fall back to the manual-reload message.
|
||||
self.recovering.write().await.remove(model_id);
|
||||
tracing::error!(model = %model_id, "auto-recovery: task unavailable");
|
||||
return poisoned_error(model_id);
|
||||
}
|
||||
}
|
||||
recovering_error(model_id)
|
||||
}
|
||||
|
||||
/// Rebuild a poisoned model: `unload_model` (drops it → cudarc aborts
|
||||
/// NCCL + releases the context) then `load_model` from the retained
|
||||
/// spec. A successful reload re-runs NCCL init + sanity inside the load
|
||||
/// path, so it returns a fresh, healthy model; a failed reload leaves
|
||||
/// the model unloaded (recoverable by the next load), never poisoned
|
||||
/// forever. Runs on the background task — never inline on the request
|
||||
/// path (would deadlock on the `models` write lock).
|
||||
async fn recover_one(&self, model_id: &str) {
|
||||
let spec = {
|
||||
let models = self.models.read().await;
|
||||
models.get(model_id).map(|h| h.spec().clone())
|
||||
};
|
||||
let Some(spec) = spec else {
|
||||
self.recovering.write().await.remove(model_id);
|
||||
return;
|
||||
};
|
||||
tracing::warn!(model = %model_id, "auto-recovery: unload+reload starting");
|
||||
if let Err(e) = self.unload_model(model_id).await {
|
||||
tracing::error!(
|
||||
model = %model_id,
|
||||
error = %format!("{e:#}"),
|
||||
"auto-recovery: unload failed (continuing to reload)"
|
||||
);
|
||||
}
|
||||
match self.load_model(&spec).await {
|
||||
Ok(()) => tracing::info!(model = %model_id, "auto-recovery: reloaded; model healthy"),
|
||||
Err(e) => tracing::error!(
|
||||
model = %model_id,
|
||||
error = %format!("{e:#}"),
|
||||
"auto-recovery: reload failed; model left unloaded"
|
||||
),
|
||||
}
|
||||
self.recovering.write().await.remove(model_id);
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Harness for CandleHarness {
|
||||
fn name(&self) -> &str {
|
||||
@@ -2423,7 +2716,8 @@ impl Harness for CandleHarness {
|
||||
chat_template,
|
||||
has_vision: vision_meta.has_vision,
|
||||
image_token_id: vision_meta.image_token_id,
|
||||
lm_tokens_per_image: vision_meta.lm_tokens_per_image,
|
||||
image_grid_factor: vision_meta.image_grid_factor,
|
||||
spec: spec.clone(),
|
||||
});
|
||||
|
||||
let mut models = self.models.write().await;
|
||||
@@ -2630,6 +2924,20 @@ impl CandleHarness {
|
||||
);
|
||||
}
|
||||
|
||||
// Vision metadata from the same config.json the shards loaded
|
||||
// from. The TP model builder (Stage 1) materialises a replicated
|
||||
// vision tower on every rank when `vision_config` is present, so
|
||||
// `has_vision` here is consistent with what each rank loaded.
|
||||
let vision_meta = VisionMeta::from_config_path(&config_path);
|
||||
if vision_meta.has_vision {
|
||||
tracing::info!(
|
||||
model = %spec.model_id,
|
||||
image_token_id = ?vision_meta.image_token_id,
|
||||
image_grid_factor = ?vision_meta.image_grid_factor,
|
||||
"TP load: vision tower present, advertising vision capability"
|
||||
);
|
||||
}
|
||||
|
||||
let tp_loaded = StdArc::new(TpLoadedModel {
|
||||
model_id: spec.model_id.clone(),
|
||||
tokenizer,
|
||||
@@ -2645,6 +2953,10 @@ impl CandleHarness {
|
||||
reasoning_tokens,
|
||||
tool_call_tokens,
|
||||
chat_template,
|
||||
has_vision: vision_meta.has_vision,
|
||||
image_token_id: vision_meta.image_token_id,
|
||||
image_grid_factor: vision_meta.image_grid_factor,
|
||||
spec: spec.clone(),
|
||||
});
|
||||
|
||||
let mut models = self.models.write().await;
|
||||
@@ -2691,7 +3003,24 @@ impl CandleHarness {
|
||||
if tp.poisoned.load(Ordering::Acquire) {
|
||||
let _g = span.enter();
|
||||
tracing::warn!("TP chat_completion: refusing request, model poisoned");
|
||||
return Err(poisoned_error(&model_id));
|
||||
return Err(self.trigger_recovery(&model_id).await);
|
||||
}
|
||||
if debug_poison_armed(&model_id) {
|
||||
let _g = span.enter();
|
||||
tracing::warn!("NEURON_DEBUG_POISON: forcing auto-recovery (#17 verification)");
|
||||
return Err(self.trigger_recovery(&model_id).await);
|
||||
}
|
||||
|
||||
// Reject image-bearing requests against a TP model with no
|
||||
// vision tower, cleanly (`vision_unsupported`) rather than
|
||||
// silently dropping the image. Vision-capable TP loads fall
|
||||
// through to the image-aware prefill in chat_completion_tp_inner.
|
||||
if request_has_images(&request) && !tp.has_vision {
|
||||
let _g = span.enter();
|
||||
tracing::warn!(
|
||||
"TP chat_completion: rejecting image request, model has no vision tower"
|
||||
);
|
||||
return Err(InferenceError::VisionUnsupported { model_id });
|
||||
}
|
||||
|
||||
let tp_for_marker = Arc::clone(&tp);
|
||||
@@ -2768,7 +3097,19 @@ impl CandleHarness {
|
||||
request: ChatCompletionRequest,
|
||||
) -> Result<InferenceStream, InferenceError> {
|
||||
if tp.poisoned.load(Ordering::Acquire) {
|
||||
return Err(poisoned_error(&request.model));
|
||||
return Err(self.trigger_recovery(&request.model).await);
|
||||
}
|
||||
|
||||
// Reject image requests against a non-vision TP model before
|
||||
// opening the SSE stream. Vision-capable TP loads fall through
|
||||
// to the image-aware prefill in the orchestration task below.
|
||||
if request_has_images(&request) && !tp.has_vision {
|
||||
tracing::warn!(
|
||||
"TP chat_completion (stream): rejecting image request, model has no vision tower"
|
||||
);
|
||||
return Err(InferenceError::VisionUnsupported {
|
||||
model_id: request.model.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
let prompt = build_prompt_for_request(tp.chat_template.as_deref(), &request);
|
||||
@@ -2776,7 +3117,58 @@ impl CandleHarness {
|
||||
.tokenizer
|
||||
.encode(prompt.as_str(), true)
|
||||
.map_err(|e| InferenceError::Other(anyhow::anyhow!("tokenize: {e}")))?;
|
||||
let prompt_tokens: Vec<u32> = encoding.get_ids().to_vec();
|
||||
let mut prompt_tokens: Vec<u32> = encoding.get_ids().to_vec();
|
||||
|
||||
// TP-vision (streaming): same detection + pad expansion as the
|
||||
// non-streaming path. The resulting `vision_route` moves into
|
||||
// the orchestration task, which runs a single-shot image prefill
|
||||
// when present. Returning early here keeps a rejected request
|
||||
// from opening the SSE stream.
|
||||
let vision_route: Option<(Vec<String>, u32)> = if request_has_images(&request) {
|
||||
if !tp.has_vision {
|
||||
return Err(InferenceError::VisionUnsupported {
|
||||
model_id: request.model.clone(),
|
||||
});
|
||||
}
|
||||
let image_token_id =
|
||||
tp.image_token_id
|
||||
.ok_or_else(|| InferenceError::VisionUnsupported {
|
||||
model_id: request.model.clone(),
|
||||
})?;
|
||||
let factor = tp
|
||||
.image_grid_factor
|
||||
.ok_or_else(|| InferenceError::VisionUnsupported {
|
||||
model_id: request.model.clone(),
|
||||
})?;
|
||||
let data_uris = extract_image_data_uris(&request);
|
||||
if data_uris.is_empty() {
|
||||
return Err(InferenceError::Other(anyhow::anyhow!(
|
||||
"request has image content but extractor produced zero data URIs"
|
||||
)));
|
||||
}
|
||||
// Per-image LM token count from each image's resized grid (#14).
|
||||
// Decode header + smart_resize only; the workers re-derive the
|
||||
// same dims when they preprocess for the replicated tower.
|
||||
let profile = super::preprocess::PreprocessProfile::qwen3_6();
|
||||
let per_image_counts: Vec<usize> = data_uris
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, uri)| {
|
||||
let (h, w) =
|
||||
super::preprocess::resized_dims_for_uri(uri, &profile).map_err(|e| {
|
||||
InferenceError::Other(anyhow::anyhow!("resized_dims image #{i}: {e}"))
|
||||
})?;
|
||||
Ok::<usize, InferenceError>((h as usize / factor) * (w as usize / factor))
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
prompt_tokens =
|
||||
expand_image_pad_tokens(&prompt_tokens, image_token_id, &per_image_counts)
|
||||
.map_err(InferenceError::Other)?;
|
||||
Some((data_uris, image_token_id))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let prompt_len = prompt_tokens.len();
|
||||
|
||||
let temperature = request.temperature.unwrap_or(0.7);
|
||||
@@ -2844,6 +3236,9 @@ impl CandleHarness {
|
||||
);
|
||||
|
||||
validate_request(prompt_len, vram_free_mb)?;
|
||||
if vision_route.is_some() {
|
||||
validate_vision_prefill(prompt_len, vram_free_mb)?;
|
||||
}
|
||||
|
||||
let tp_for_task = Arc::clone(&tp);
|
||||
tokio::spawn(
|
||||
@@ -2890,14 +3285,29 @@ impl CandleHarness {
|
||||
// chunk fans out to every rank with a growing
|
||||
// offset; only the final chunk's logits are kept
|
||||
// for the first sample.
|
||||
let logits_vec = match chunked_prefill_tp(
|
||||
&mut pool,
|
||||
&model_id,
|
||||
leader_handle,
|
||||
&prompt_tokens,
|
||||
)
|
||||
.await
|
||||
{
|
||||
// Vision requests do a chunked image prefill (encode
|
||||
// once, splice per chunk); text requests chunk it the
|
||||
// same way. `vision_route` was moved into this task
|
||||
// from the synchronous setup above.
|
||||
let prefill_result = match &vision_route {
|
||||
Some((data_uris, image_token_id)) => {
|
||||
pool.generate_step_with_images(
|
||||
&model_id,
|
||||
leader_handle,
|
||||
prompt_tokens.clone(),
|
||||
0,
|
||||
*image_token_id,
|
||||
data_uris.clone(),
|
||||
prefill_chunk_tokens(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
None => {
|
||||
chunked_prefill_tp(&mut pool, &model_id, leader_handle, &prompt_tokens)
|
||||
.await
|
||||
}
|
||||
};
|
||||
let logits_vec = match prefill_result {
|
||||
Ok(l) => l,
|
||||
Err(e) => {
|
||||
failure = Some(format!("prefill: {e:#}"));
|
||||
@@ -3240,7 +3650,55 @@ async fn chat_completion_tp_inner(
|
||||
.tokenizer
|
||||
.encode(prompt.as_str(), true)
|
||||
.map_err(|e| InferenceError::Other(anyhow::anyhow!("tokenize: {e}")))?;
|
||||
let prompt_tokens: Vec<u32> = encoding.get_ids().to_vec();
|
||||
let mut prompt_tokens: Vec<u32> = encoding.get_ids().to_vec();
|
||||
|
||||
// TP-vision: when the request carries images (and the model has a
|
||||
// replicated tower — enforced by the caller's guard), expand each
|
||||
// `<|image_pad|>` sentinel to the per-image patch count and carry
|
||||
// the source data URIs through to the single-shot image prefill.
|
||||
// Mirrors the single-GPU `chat_completion` vision_route block.
|
||||
let vision_route: Option<(Vec<String>, u32)> = if request_has_images(&request) {
|
||||
if !tp.has_vision {
|
||||
return Err(InferenceError::VisionUnsupported {
|
||||
model_id: request.model.clone(),
|
||||
});
|
||||
}
|
||||
let image_token_id =
|
||||
tp.image_token_id
|
||||
.ok_or_else(|| InferenceError::VisionUnsupported {
|
||||
model_id: request.model.clone(),
|
||||
})?;
|
||||
let factor = tp
|
||||
.image_grid_factor
|
||||
.ok_or_else(|| InferenceError::VisionUnsupported {
|
||||
model_id: request.model.clone(),
|
||||
})?;
|
||||
let data_uris = extract_image_data_uris(&request);
|
||||
if data_uris.is_empty() {
|
||||
return Err(InferenceError::Other(anyhow::anyhow!(
|
||||
"request has image content but extractor produced zero data URIs"
|
||||
)));
|
||||
}
|
||||
// Per-image LM token count from each image's resized grid (#14).
|
||||
let profile = super::preprocess::PreprocessProfile::qwen3_6();
|
||||
let per_image_counts: Vec<usize> = data_uris
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, uri)| {
|
||||
let (h, w) =
|
||||
super::preprocess::resized_dims_for_uri(uri, &profile).map_err(|e| {
|
||||
InferenceError::Other(anyhow::anyhow!("resized_dims image #{i}: {e}"))
|
||||
})?;
|
||||
Ok::<usize, InferenceError>((h as usize / factor) * (w as usize / factor))
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
prompt_tokens = expand_image_pad_tokens(&prompt_tokens, image_token_id, &per_image_counts)
|
||||
.map_err(InferenceError::Other)?;
|
||||
Some((data_uris, image_token_id))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let prompt_len = prompt_tokens.len();
|
||||
|
||||
let temperature = request.temperature.unwrap_or(0.7);
|
||||
@@ -3267,6 +3725,9 @@ async fn chat_completion_tp_inner(
|
||||
);
|
||||
|
||||
validate_request(prompt_len, vram_free_mb)?;
|
||||
if vision_route.is_some() {
|
||||
validate_vision_prefill(prompt_len, vram_free_mb)?;
|
||||
}
|
||||
|
||||
// Acquire the pool lock for the duration of the request. After
|
||||
// Phase 3 the leader's TpLeaderModel lives in the device worker
|
||||
@@ -3310,9 +3771,26 @@ async fn chat_completion_tp_inner(
|
||||
// spread across multiple `generate_step` calls with monotonically
|
||||
// growing offsets.
|
||||
let prefill_start = std::time::Instant::now();
|
||||
let logits_vec = chunked_prefill_tp(&mut pool, &model_id, leader_handle, &prompt_tokens)
|
||||
.await
|
||||
.map_err(InferenceError::Other)?;
|
||||
// Vision requests do a chunked image prefill (every rank encodes its
|
||||
// replicated tower once, then splices per chunk); text requests
|
||||
// chunk the prefill the same way.
|
||||
let logits_vec = match &vision_route {
|
||||
Some((data_uris, image_token_id)) => pool
|
||||
.generate_step_with_images(
|
||||
&model_id,
|
||||
leader_handle,
|
||||
prompt_tokens.clone(),
|
||||
0,
|
||||
*image_token_id,
|
||||
data_uris.clone(),
|
||||
prefill_chunk_tokens(),
|
||||
)
|
||||
.await
|
||||
.map_err(InferenceError::Other)?,
|
||||
None => chunked_prefill_tp(&mut pool, &model_id, leader_handle, &prompt_tokens)
|
||||
.await
|
||||
.map_err(InferenceError::Other)?,
|
||||
};
|
||||
let (post_prefill_vram_free_mb, _) = tp.query_vram().await;
|
||||
tracing::info!(
|
||||
model = %model_id,
|
||||
@@ -3662,10 +4140,12 @@ fn build_prompt_for_request(
|
||||
struct VisionMeta {
|
||||
has_vision: bool,
|
||||
image_token_id: Option<u32>,
|
||||
/// LM-side tokens this model's vision tower emits per image at
|
||||
/// the Stage B fixed `PreprocessProfile::qwen3_6()` resolution
|
||||
/// (448×448). Equal to `(H/patch_size/spatial_merge_size)²`.
|
||||
lm_tokens_per_image: Option<usize>,
|
||||
/// `patch_size × spatial_merge_size` — the divisor that turns a
|
||||
/// resized pixel dimension into an LM-grid dimension. An image of
|
||||
/// resized `(h, w)` emits `(h/factor) × (w/factor)` LM tokens (#14
|
||||
/// dynamic resolution; was a constant 196 at the old fixed 448²).
|
||||
/// `None` for text-only models.
|
||||
image_grid_factor: Option<usize>,
|
||||
}
|
||||
|
||||
impl VisionMeta {
|
||||
@@ -3694,22 +4174,18 @@ impl VisionMeta {
|
||||
.get("image_token_id")
|
||||
.and_then(|x| x.as_u64())
|
||||
.map(|n| n as u32);
|
||||
// Compute LM tokens per image at the Stage B fixed resolution
|
||||
// (PreprocessProfile::qwen3_6() → 448×448). One LM token per
|
||||
// spatial-merge group of patches.
|
||||
let target_h = super::preprocess::PreprocessProfile::qwen3_6().target_height as usize;
|
||||
let target_w = super::preprocess::PreprocessProfile::qwen3_6().target_width as usize;
|
||||
let lm_tokens_per_image = if patch_size > 0 && spatial_merge_size > 0 {
|
||||
let gh = target_h / patch_size / spatial_merge_size;
|
||||
let gw = target_w / patch_size / spatial_merge_size;
|
||||
Some(gh * gw)
|
||||
// The pixel→LM-grid divisor. An image resized to (h, w) emits
|
||||
// (h/factor) × (w/factor) LM tokens — computed per image at
|
||||
// request time now that resolution is dynamic (#14).
|
||||
let image_grid_factor = if patch_size > 0 && spatial_merge_size > 0 {
|
||||
Some(patch_size * spatial_merge_size)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
Self {
|
||||
has_vision: true,
|
||||
image_token_id,
|
||||
lm_tokens_per_image,
|
||||
image_grid_factor,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -3756,13 +4232,13 @@ fn extract_images_from_request(
|
||||
.and_then(|v| v.get("url"))
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| anyhow::anyhow!("image_url part missing url field"))?;
|
||||
let pixels = super::preprocess::preprocess_data_uri(url, profile)
|
||||
let (pixels, h, w) = super::preprocess::preprocess_data_uri(url, profile)
|
||||
.with_context(|| format!("preprocess image #{}", out.len()))?;
|
||||
out.push(super::device_worker::jobs::ImageInput {
|
||||
pixels,
|
||||
c: 3,
|
||||
h: profile.target_height as usize,
|
||||
w: profile.target_width as usize,
|
||||
h: h as usize,
|
||||
w: w as usize,
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -3770,6 +4246,37 @@ fn extract_images_from_request(
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Collect the raw `image_url.url` strings (data URIs) from a chat
|
||||
/// request, in prompt order. The TP vision path (Stage C / TP-vision)
|
||||
/// ships these verbatim to every rank, which each preprocess + encode
|
||||
/// identically — so unlike `extract_images_from_request` (which
|
||||
/// preprocesses on the leader for the single-GPU worker job) this
|
||||
/// keeps the source form for replicated per-rank encoding.
|
||||
///
|
||||
/// Cuda-gated: the only callers are the TP entry points, which compile
|
||||
/// only under the `cuda` feature.
|
||||
#[cfg(feature = "cuda")]
|
||||
fn extract_image_data_uris(request: &ChatCompletionRequest) -> Vec<String> {
|
||||
let mut out = Vec::new();
|
||||
for msg in &request.messages {
|
||||
if let MessageContent::Parts(parts) = &msg.content {
|
||||
for part in parts {
|
||||
if part.get("type").and_then(|v| v.as_str()) != Some("image_url") {
|
||||
continue;
|
||||
}
|
||||
if let Some(url) = part
|
||||
.get("image_url")
|
||||
.and_then(|v| v.get("url"))
|
||||
.and_then(|v| v.as_str())
|
||||
{
|
||||
out.push(url.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Expand each occurrence of `image_token_id` in `input_ids` into
|
||||
/// `patches_per_image[i]` copies (one expansion per image, in order).
|
||||
/// Stage B4 helper.
|
||||
@@ -4046,6 +4553,17 @@ async fn run_inference_via_worker(
|
||||
/// forward step through `worker.forward_logits()`. Same per-step
|
||||
/// CPU-side sampling discipline — no device tensor escapes the
|
||||
/// worker thread.
|
||||
///
|
||||
/// `images` carries the Stage C vision payload. When `Some`, prefill
|
||||
/// is a single-shot `forward_logits_with_images` that splices image
|
||||
/// embeddings at `image_token_id` positions (same contract as the
|
||||
/// non-streaming [`run_inference_with_images_via_worker`]); image
|
||||
/// embeddings are prefill-only, so every decode step below takes the
|
||||
/// plain `forward_logits` path regardless. When `None`, prefill is
|
||||
/// chunked (`chunked_prefill_via_worker`) to bound activation memory
|
||||
/// — the original text-only behaviour, unchanged. The decode loop and
|
||||
/// the `route_token!` reasoning/tool-call state machine are shared
|
||||
/// across both prefill shapes, so there's exactly one copy to maintain.
|
||||
#[cfg(feature = "cuda")]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn stream_inference_via_worker(
|
||||
@@ -4053,6 +4571,7 @@ async fn stream_inference_via_worker(
|
||||
handle: super::device_worker::ArchHandle,
|
||||
tokenizer: Tokenizer,
|
||||
prompt_tokens: Vec<u32>,
|
||||
images: Option<(Vec<super::device_worker::jobs::ImageInput>, u32)>,
|
||||
max_new: usize,
|
||||
temperature: f64,
|
||||
top_p: Option<f64>,
|
||||
@@ -4098,11 +4617,19 @@ async fn stream_inference_via_worker(
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("clear_kv_cache: {e}"))?;
|
||||
|
||||
// Chunked prefill (see `chunked_prefill_via_worker`). The owning
|
||||
// `prompt_tokens: Vec<u32>` is borrowed for the loop's duration;
|
||||
// we still need `prompt_len` (already extracted above) for the
|
||||
// decode-step offset arithmetic.
|
||||
let logits_vec = chunked_prefill_via_worker(&*worker, handle, &prompt_tokens).await?;
|
||||
// Prefill. Vision-bearing requests (`images = Some`) do a
|
||||
// single-shot prefill that splices the image embeddings; text-only
|
||||
// requests use chunked prefill (see `chunked_prefill_via_worker`)
|
||||
// to bound activation memory. Either way the owning
|
||||
// `prompt_tokens: Vec<u32>` outlives this step; we use `prompt_len`
|
||||
// (already extracted above) for the decode-step offset arithmetic.
|
||||
let logits_vec = match images {
|
||||
Some((imgs, image_token_id)) => worker
|
||||
.forward_logits_with_images(handle, prompt_tokens.clone(), 0, imgs, image_token_id)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("forward_logits_with_images: {e}"))?,
|
||||
None => chunked_prefill_via_worker(&*worker, handle, &prompt_tokens).await?,
|
||||
};
|
||||
let logits = Tensor::new(logits_vec.as_slice(), &Device::Cpu)?;
|
||||
let mut next_token = match sample_with_penalty(&logits, &all_tokens, &mut logits_processor) {
|
||||
Ok(t) => t,
|
||||
@@ -4699,4 +5226,62 @@ mod tests {
|
||||
let out = expand_image_pad_tokens(&input, pad, &[]).unwrap();
|
||||
assert_eq!(out, input);
|
||||
}
|
||||
|
||||
/// `request_has_images` is the gate that routes both the
|
||||
/// non-streaming (`chat_completion`) and streaming
|
||||
/// (`inference_stream`, Stage C1) paths to the vision-aware
|
||||
/// prefill. Exercise the three shapes it must distinguish: plain
|
||||
/// text, a text-only content-parts array, and a parts array
|
||||
/// carrying an `image_url`.
|
||||
#[test]
|
||||
fn request_has_images_detects_image_url_parts() {
|
||||
let text_only: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
}))
|
||||
.unwrap();
|
||||
assert!(!request_has_images(&text_only));
|
||||
|
||||
let parts_text_only: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": [
|
||||
{"type": "text", "text": "hello"}
|
||||
]}],
|
||||
}))
|
||||
.unwrap();
|
||||
assert!(!request_has_images(&parts_text_only));
|
||||
|
||||
let with_image: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": [
|
||||
{"type": "text", "text": "what is this?"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,AAA="}}
|
||||
]}],
|
||||
}))
|
||||
.unwrap();
|
||||
assert!(request_has_images(&with_image));
|
||||
}
|
||||
|
||||
/// The vision pre-flight guard rejects a prefill whose estimated
|
||||
/// footprint exceeds free VRAM (so a doomed request fails clean
|
||||
/// instead of OOM-hanging the TP collective), passes one that fits,
|
||||
/// and is skipped on the CPU sentinel.
|
||||
#[test]
|
||||
fn vision_prefill_guard_behaviour() {
|
||||
// CPU sentinel (vram_free_mb == 0) is always allowed.
|
||||
assert!(validate_vision_prefill(10_000_000, 0).is_ok());
|
||||
|
||||
// A clearly-oversized prompt against tiny free VRAM is rejected
|
||||
// for any non-degenerate config (default: 2000 base + 500/1k).
|
||||
assert!(matches!(
|
||||
validate_vision_prefill(10_000_000, 50),
|
||||
Err(InferenceError::InsufficientVram { .. })
|
||||
));
|
||||
|
||||
// With defaults, the agent-0-sized 12,960-token prompt that
|
||||
// OOM'd single-shot fits the estimate at ~12 GB free (2000 +
|
||||
// 12960*500/1000 = 8480 MiB) — the chunked prefill handles it,
|
||||
// so the guard must NOT reject it.
|
||||
assert!(validate_vision_prefill(12_960, 12_445).is_ok());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -43,7 +43,7 @@
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use cortex_core::openai::{ChatMessage, MessageContent};
|
||||
use minijinja::Environment;
|
||||
use minijinja::{Environment, Error as MjError, ErrorKind as MjErrorKind, Value as MjValue};
|
||||
use serde_json::Value;
|
||||
use std::path::Path;
|
||||
|
||||
@@ -65,12 +65,55 @@ pub fn chat_templates_enabled() -> bool {
|
||||
}
|
||||
}
|
||||
|
||||
/// Convenience: probe for `tokenizer_config.json` in the same
|
||||
/// directory the tokenizer was loaded from. Both files come from
|
||||
/// the same HuggingFace snapshot in the hf-hub cache, so the
|
||||
/// sibling path is reliable.
|
||||
/// Probe for the model's chat template in the same directory the
|
||||
/// tokenizer was loaded from, following HuggingFace `transformers`
|
||||
/// precedence: a standalone `chat_template.jinja` (then
|
||||
/// `chat_template.json`) wins over the `chat_template` field in
|
||||
/// `tokenizer_config.json`.
|
||||
///
|
||||
/// This matters for multimodal models: Qwen3-VL / Qwen3.6 ship their
|
||||
/// vision-aware template (the one that emits
|
||||
/// `<|vision_start|><|image_pad|><|vision_end|>` per image) **only** in
|
||||
/// `chat_template.jinja`, and may not ship a `tokenizer_config.json` at
|
||||
/// all. Reading `tokenizer_config.json` alone returned `None`, which
|
||||
/// dropped image content into the text-only `format_qwen3_prompt`
|
||||
/// fallback — so image requests rendered zero `<|image_pad|>` tokens
|
||||
/// and the vision path bailed on the count mismatch.
|
||||
pub fn load_chat_template_alongside(tokenizer_json_path: &Path) -> Option<String> {
|
||||
let parent = tokenizer_json_path.parent()?;
|
||||
|
||||
// 1. Standalone Jinja file — raw template text, highest priority.
|
||||
let jinja_path = parent.join("chat_template.jinja");
|
||||
match std::fs::read_to_string(&jinja_path) {
|
||||
Ok(text) if !text.trim().is_empty() => {
|
||||
tracing::info!(
|
||||
path = %jinja_path.display(),
|
||||
"chat_template: loaded standalone chat_template.jinja"
|
||||
);
|
||||
return Some(text);
|
||||
}
|
||||
Ok(_) => {
|
||||
tracing::warn!(
|
||||
path = %jinja_path.display(),
|
||||
"chat_template: chat_template.jinja present but empty; trying other sources"
|
||||
);
|
||||
}
|
||||
Err(_) => {} // absent — fall through, common case
|
||||
}
|
||||
|
||||
// 2. Standalone JSON file — `{"chat_template": "..."}` form.
|
||||
let json_path = parent.join("chat_template.json");
|
||||
if json_path.exists()
|
||||
&& let Some(t) = load_chat_template_from(&json_path)
|
||||
{
|
||||
tracing::info!(
|
||||
path = %json_path.display(),
|
||||
"chat_template: loaded standalone chat_template.json"
|
||||
);
|
||||
return Some(t);
|
||||
}
|
||||
|
||||
// 3. The `chat_template` field inside tokenizer_config.json.
|
||||
let config_path = parent.join("tokenizer_config.json");
|
||||
load_chat_template_from(&config_path)
|
||||
}
|
||||
@@ -148,6 +191,25 @@ pub fn render_chat_template(
|
||||
kwargs: &Value,
|
||||
) -> Result<String> {
|
||||
let mut env = Environment::new();
|
||||
|
||||
// HF chat templates are authored against Python's Jinja2 with its
|
||||
// string semantics. Bridge the two so real model templates render:
|
||||
//
|
||||
// - `pycompat::unknown_method_callback` supplies Python str/list/dict
|
||||
// methods minijinja lacks natively (`startswith`, `endswith`,
|
||||
// `split`, `rstrip`, `lstrip`, …) — the Qwen3.6 template uses
|
||||
// several in its think-block and tool-response handling.
|
||||
// - `raise_exception` is the global HF templates call to reject
|
||||
// malformed inputs (e.g. an image in a system message). Map it to
|
||||
// a render error so the caller falls back / surfaces it.
|
||||
env.set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback);
|
||||
env.add_function(
|
||||
"raise_exception",
|
||||
|msg: String| -> Result<MjValue, MjError> {
|
||||
Err(MjError::new(MjErrorKind::InvalidOperation, msg))
|
||||
},
|
||||
);
|
||||
|
||||
// Compile the template against a fixed name so error messages
|
||||
// surface "chat_template" rather than `<template>`.
|
||||
env.add_template("chat_template", template)
|
||||
@@ -210,6 +272,114 @@ mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
/// Reproduces the Qwen3.6 vision template's image-insertion
|
||||
/// condition against the OpenAI `image_url` content-part shape our
|
||||
/// renderer forwards. Confirms minijinja's `'image_url' in item`
|
||||
/// matches a serde_json object that carries that key — i.e. the
|
||||
/// template *can* emit `<|image_pad|>` for our parts.
|
||||
#[test]
|
||||
fn image_url_part_renders_image_pad() {
|
||||
// Condition copied from doc/vision-qwen3_6-spec.md (lines 8-18
|
||||
// of the real chat_template.jinja).
|
||||
let template = "{%- for message in messages -%}\
|
||||
{%- if message.content is string -%}\
|
||||
{{ message.content }}\
|
||||
{%- else -%}\
|
||||
{%- for item in message.content -%}\
|
||||
{%- if 'image' in item or 'image_url' in item or item.type == 'image' -%}\
|
||||
<|vision_start|><|image_pad|><|vision_end|>\
|
||||
{%- elif item.type == 'text' -%}\
|
||||
{{ item.text }}\
|
||||
{%- endif -%}\
|
||||
{%- endfor -%}\
|
||||
{%- endif -%}\
|
||||
{%- endfor -%}";
|
||||
let messages = vec![ChatMessage {
|
||||
role: "user".into(),
|
||||
content: MessageContent::Parts(vec![
|
||||
json!({"type": "text", "text": "what is this?"}),
|
||||
json!({"type": "image_url", "image_url": {"url": "data:image/png;base64,AAA="}}),
|
||||
]),
|
||||
extra: Value::Object(Default::default()),
|
||||
}];
|
||||
let out = render_chat_template(template, &messages, &Value::Null, &Value::Null)
|
||||
.expect("render should succeed");
|
||||
assert!(
|
||||
out.contains("<|image_pad|>"),
|
||||
"expected the image_url part to emit <|image_pad|>; rendered: {out:?}"
|
||||
);
|
||||
}
|
||||
|
||||
/// `chat_template.jinja` must win over `tokenizer_config.json`'s
|
||||
/// `chat_template` field — the transformers precedence Qwen3.6
|
||||
/// relies on (its vision template ships only in the `.jinja` file).
|
||||
#[test]
|
||||
fn standalone_jinja_template_takes_precedence() {
|
||||
let dir = std::env::temp_dir().join(format!(
|
||||
"neuron_ct_precedence_{}_{}",
|
||||
std::process::id(),
|
||||
line!()
|
||||
));
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
std::fs::write(dir.join("chat_template.jinja"), "FROM_JINJA").unwrap();
|
||||
std::fs::write(
|
||||
dir.join("tokenizer_config.json"),
|
||||
r#"{"chat_template": "FROM_CONFIG"}"#,
|
||||
)
|
||||
.unwrap();
|
||||
// tokenizer_json_path is the sibling the loader takes a parent of.
|
||||
let got = load_chat_template_alongside(&dir.join("tokenizer.json"));
|
||||
std::fs::remove_dir_all(&dir).ok();
|
||||
assert_eq!(got.as_deref(), Some("FROM_JINJA"));
|
||||
}
|
||||
|
||||
/// With no standalone file, fall back to the tokenizer_config.json
|
||||
/// field — the text-only path stays unchanged.
|
||||
#[test]
|
||||
fn falls_back_to_tokenizer_config_when_no_standalone() {
|
||||
let dir = std::env::temp_dir().join(format!(
|
||||
"neuron_ct_fallback_{}_{}",
|
||||
std::process::id(),
|
||||
line!()
|
||||
));
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
std::fs::write(
|
||||
dir.join("tokenizer_config.json"),
|
||||
r#"{"chat_template": "FROM_CONFIG"}"#,
|
||||
)
|
||||
.unwrap();
|
||||
let got = load_chat_template_alongside(&dir.join("tokenizer.json"));
|
||||
std::fs::remove_dir_all(&dir).ok();
|
||||
assert_eq!(got.as_deref(), Some("FROM_CONFIG"));
|
||||
}
|
||||
|
||||
/// The *actual* Qwen3.6-27B `chat_template.jinja` (verbatim from
|
||||
/// beast's HF cache) must render in minijinja and emit exactly one
|
||||
/// `<|image_pad|>` for a text+image user turn. This is the real
|
||||
/// end-to-end check the unit tests above only approximate — it
|
||||
/// catches any minijinja incompatibility (namespace, macros,
|
||||
/// reverse slice, string methods) before it reaches production.
|
||||
#[test]
|
||||
fn real_qwen3_6_template_renders_one_image_pad() {
|
||||
let template = include_str!("testdata/qwen3_6_chat_template.jinja");
|
||||
let messages = vec![ChatMessage {
|
||||
role: "user".into(),
|
||||
content: MessageContent::Parts(vec![
|
||||
json!({"type": "text", "text": "what is this?"}),
|
||||
json!({"type": "image_url", "image_url": {"url": "data:image/png;base64,AAA="}}),
|
||||
]),
|
||||
extra: Value::Object(Default::default()),
|
||||
}];
|
||||
let out = render_chat_template(template, &messages, &Value::Null, &Value::Null)
|
||||
.expect("real Qwen3.6 template should render in minijinja");
|
||||
let pads = out.matches("<|image_pad|>").count();
|
||||
assert_eq!(
|
||||
pads, 1,
|
||||
"expected exactly one <|image_pad|>; rendered:\n{out}"
|
||||
);
|
||||
assert!(out.contains("<|vision_start|>") && out.contains("<|vision_end|>"));
|
||||
}
|
||||
|
||||
fn user_msg(text: &str) -> ChatMessage {
|
||||
ChatMessage {
|
||||
role: "user".into(),
|
||||
|
||||
@@ -201,6 +201,16 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
|
||||
let _ = reply.send(resp);
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
Job::GetLeaderComm { reply } => {
|
||||
// Clone the leader's Arc<Comm> out for the async-side
|
||||
// watchdog. `None` before NcclInit. (#17 Stage 2)
|
||||
let comm = state
|
||||
.nccl
|
||||
.comm()
|
||||
.map(crate::harness::tp::nccl_state::SendComm);
|
||||
let _ = reply.send(comm);
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
Job::TpLoadShard {
|
||||
model_id,
|
||||
config_json,
|
||||
@@ -262,6 +272,27 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
|
||||
let result = tp_forward_logits(&mut state, handle, &tokens, offset);
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
Job::TpForwardLogitsWithImages {
|
||||
handle,
|
||||
tokens,
|
||||
offset,
|
||||
image_token_id,
|
||||
image_data_uris,
|
||||
chunk_size,
|
||||
reply,
|
||||
} => {
|
||||
let result = tp_forward_logits_with_images(
|
||||
&mut state,
|
||||
handle,
|
||||
&tokens,
|
||||
offset,
|
||||
image_token_id,
|
||||
&image_data_uris,
|
||||
chunk_size,
|
||||
);
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
// Handled by the matches!() check above; reaching here
|
||||
// means a Shutdown slipped past which is a bug.
|
||||
Job::Shutdown => unreachable!("Shutdown should break above"),
|
||||
@@ -734,6 +765,61 @@ fn tp_forward_logits(
|
||||
Ok(values)
|
||||
}
|
||||
|
||||
/// 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
|
||||
/// device, encodes + splices + forwards via
|
||||
/// `TpLeaderModel::forward_with_images`, and copies the `[vocab]`
|
||||
/// logits to CPU. Mirrors the single-GPU `forward_logits_with_images`
|
||||
/// but on the TP leader's replicated tower.
|
||||
#[cfg(feature = "cuda")]
|
||||
fn tp_forward_logits_with_images(
|
||||
state: &mut DeviceWorkerState,
|
||||
handle: TpHandle,
|
||||
tokens: &[u32],
|
||||
offset: usize,
|
||||
image_token_id: u32,
|
||||
image_data_uris: &[String],
|
||||
chunk_size: usize,
|
||||
) -> anyhow::Result<Vec<f32>> {
|
||||
use crate::harness::preprocess::{PreprocessProfile, preprocess_data_uri};
|
||||
use candle_core::{DType, Tensor};
|
||||
|
||||
if image_data_uris.is_empty() {
|
||||
anyhow::bail!("TpForwardLogitsWithImages dispatched with zero images");
|
||||
}
|
||||
|
||||
// Preprocess every image into a device-resident (C, H, W) tensor at
|
||||
// its native-aspect resized dims (#14). Same `smart_resize` + decode
|
||||
// path the subprocess workers run, so the encoded embeddings — and
|
||||
// the per-image grids derived from these dims — match across ranks
|
||||
// bit-for-bit.
|
||||
let profile = PreprocessProfile::qwen3_6();
|
||||
let mut pixels: Vec<Tensor> = Vec::with_capacity(image_data_uris.len());
|
||||
for (idx, uri) in image_data_uris.iter().enumerate() {
|
||||
let (px, h, w) = preprocess_data_uri(uri, &profile)
|
||||
.with_context(|| format!("preprocess image[{idx}] (TP leader)"))?;
|
||||
let t = Tensor::from_vec(px, (3, h as usize, w as usize), &state.device)?;
|
||||
pixels.push(t);
|
||||
}
|
||||
|
||||
let model = state.tp_models.get_mut(&handle).ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"TpForwardLogitsWithImages: no model for handle {}",
|
||||
handle.0
|
||||
)
|
||||
})?;
|
||||
|
||||
// Chunked prefill (encode once, splice per chunk) — bounded
|
||||
// activation, in lockstep with the subprocess ranks.
|
||||
let logits =
|
||||
model.prefill_with_images_chunked(tokens, offset, &pixels, image_token_id, chunk_size)?;
|
||||
let logits = logits.squeeze(0)?.squeeze(0)?;
|
||||
let logits = logits.to_dtype(DType::F32)?.flatten_all()?;
|
||||
let values = logits.to_vec1::<f32>()?;
|
||||
Ok(values)
|
||||
}
|
||||
|
||||
/// Forward step + copy the `[vocab]` logits to a CPU `Vec<f32>` ready
|
||||
/// for sampling on the async caller. The model's `device()` (CUDA or
|
||||
/// CPU) determines where the kernel runs; this fn doesn't care.
|
||||
@@ -799,9 +885,17 @@ fn forward_logits_with_images(
|
||||
anyhow::anyhow!("ForwardLogitsWithImages: no model for handle {}", handle.0)
|
||||
})?;
|
||||
|
||||
// pixel→LM-grid divisor (patch×merge) for this tower; each image's
|
||||
// LM grid is (h/factor, w/factor) (#14 dynamic resolution).
|
||||
let factor = arch.vision_grid_factor().ok_or_else(|| {
|
||||
anyhow::anyhow!("ForwardLogitsWithImages: loaded model has no vision tower")
|
||||
})?;
|
||||
|
||||
// Encode every image on the worker's device, collecting per-image
|
||||
// post-merger embeddings as device-resident tensors.
|
||||
// post-merger embeddings as device-resident tensors plus their LM
|
||||
// grids (for the interleaved-M-RoPE position ids).
|
||||
let mut per_image: Vec<Tensor> = Vec::with_capacity(images.len());
|
||||
let mut grids: Vec<(usize, usize)> = Vec::with_capacity(images.len());
|
||||
for (idx, img) in images.into_iter().enumerate() {
|
||||
anyhow::ensure!(
|
||||
img.pixels.len() == img.c * img.h * img.w,
|
||||
@@ -811,6 +905,7 @@ fn forward_logits_with_images(
|
||||
img.h,
|
||||
img.w,
|
||||
);
|
||||
grids.push((img.h / factor, img.w / factor));
|
||||
let image = Tensor::from_vec(img.pixels, (img.c, img.h, img.w), &state.device)?;
|
||||
let embed = arch
|
||||
.encode_image(&image)
|
||||
@@ -823,7 +918,7 @@ fn forward_logits_with_images(
|
||||
let image_embeds = Tensor::cat(&per_image.iter().collect::<Vec<_>>(), 0)?;
|
||||
|
||||
let input = Tensor::new(tokens, &state.device)?.unsqueeze(0)?;
|
||||
let logits = arch.forward_with_vision(&input, offset, &image_embeds, image_token_id)?;
|
||||
let logits = arch.forward_with_vision(&input, offset, &image_embeds, image_token_id, &grids)?;
|
||||
let values = logits
|
||||
.to_dtype(DType::F32)?
|
||||
.flatten_all()?
|
||||
@@ -919,6 +1014,10 @@ fn drain_poisoned(job: Job, device_index: u32) {
|
||||
message: format!("device worker {device_index} poisoned"),
|
||||
});
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
Job::GetLeaderComm { reply } => {
|
||||
let _ = reply.send(None);
|
||||
}
|
||||
Job::NcclSanity { reply } => {
|
||||
let _ = reply.send(crate::harness::tp::rpc::WorkerResponse::Error {
|
||||
kind: "device_worker_poisoned".into(),
|
||||
@@ -941,6 +1040,10 @@ fn drain_poisoned(job: Job, device_index: u32) {
|
||||
Job::TpForwardLogits { reply, .. } => {
|
||||
let _ = reply.send(Err(err()));
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
Job::TpForwardLogitsWithImages { reply, .. } => {
|
||||
let _ = reply.send(Err(err()));
|
||||
}
|
||||
Job::Shutdown => {
|
||||
// Filtered by the matches!() guard in run(); reaching
|
||||
// here would be a logic error.
|
||||
|
||||
@@ -36,8 +36,13 @@ pub struct TpHandle(pub u64);
|
||||
/// `Clone` so the vision-aware dispatch in `chat_completion` can
|
||||
/// match `&vision_route` (carrying borrowed images) and still hand
|
||||
/// owned `Vec<ImageInput>` to the worker job. The clone cost is one
|
||||
/// pixel-buffer memcpy per image — fine at fixed-resolution sizes
|
||||
/// (3 × 448 × 448 × 4 bytes = ~2.4 MiB per image).
|
||||
/// pixel-buffer memcpy per image — now variable with dynamic resolution
|
||||
/// (#14): `3 × h × w × 4` bytes, up to ~6.3 MiB at the default 1024²
|
||||
/// `max_pixels` budget.
|
||||
///
|
||||
/// `h`/`w` are the **resized** dims (factor-aligned), so the per-image LM
|
||||
/// grid is `(h/factor, w/factor)` — derived downstream for the splice
|
||||
/// and the interleaved-M-RoPE position ids.
|
||||
#[derive(Clone)]
|
||||
pub struct ImageInput {
|
||||
pub pixels: Vec<f32>,
|
||||
@@ -187,6 +192,17 @@ pub enum Job {
|
||||
NcclSanity {
|
||||
reply: oneshot::Sender<crate::harness::tp::rpc::WorkerResponse>,
|
||||
},
|
||||
/// Hand a clonable handle to the leader's NCCL `Comm` back to the
|
||||
/// async side, so the TP step watchdog can call `ncclCommAbort` on
|
||||
/// it from a *different* thread to unblock a wedged collective
|
||||
/// (#17 Stage 2). Fetched once at init while the worker thread is
|
||||
/// still responsive — a thread already wedged in a collective can't
|
||||
/// service this job, which is exactly why the handle is cached
|
||||
/// up front. Replies `None` before `NcclInit` has run.
|
||||
#[cfg(feature = "cuda")]
|
||||
GetLeaderComm {
|
||||
reply: oneshot::Sender<Option<crate::harness::tp::nccl_state::SendComm>>,
|
||||
},
|
||||
/// Load the leader's TP shard on the worker thread. The dispatch
|
||||
/// handler reads `state.nccl.comm()` directly (no cross-thread
|
||||
/// `Arc<Comm>` transfer, no `SendComm` wrapper) and builds the
|
||||
@@ -231,6 +247,24 @@ pub enum Job {
|
||||
offset: usize,
|
||||
reply: oneshot::Sender<Result<Vec<f32>>>,
|
||||
},
|
||||
/// 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
|
||||
/// the leader's replicated tower, splices at `image_token_id`, and
|
||||
/// returns CPU-side `[vocab]` logits. Image tensors never escape the
|
||||
/// worker thread. Caller fans out `GenerateStepWithImages` to the
|
||||
/// subprocess ranks and drains them; only the leader forward moves
|
||||
/// here.
|
||||
#[cfg(feature = "cuda")]
|
||||
TpForwardLogitsWithImages {
|
||||
handle: TpHandle,
|
||||
tokens: Vec<u32>,
|
||||
offset: usize,
|
||||
image_token_id: u32,
|
||||
image_data_uris: Vec<String>,
|
||||
chunk_size: usize,
|
||||
reply: oneshot::Sender<Result<Vec<f32>>>,
|
||||
},
|
||||
/// Tell the worker to break its dispatch loop and exit. Any jobs
|
||||
/// queued after this in the channel reply `Err` to their oneshot
|
||||
/// senders (the senders are dropped on the worker's exit, which
|
||||
|
||||
@@ -161,6 +161,27 @@ impl DeviceWorkerHandle {
|
||||
}
|
||||
}
|
||||
|
||||
/// Fetch a clonable handle to the leader's NCCL `Comm` (#17 Stage 2).
|
||||
/// The TP step watchdog caches this at init so it can call
|
||||
/// `ncclCommAbort` from the async thread to unblock a wedged
|
||||
/// collective. Returns `None` if uninitialised, poisoned, or gone —
|
||||
/// the caller treats a missing handle as "can't abort" and logs it.
|
||||
#[cfg(feature = "cuda")]
|
||||
pub async fn get_leader_comm(&self) -> Option<crate::harness::tp::nccl_state::SendComm> {
|
||||
if self.poisoned.load(Ordering::Acquire) {
|
||||
return None;
|
||||
}
|
||||
let (reply_tx, reply_rx) = oneshot::channel();
|
||||
if self
|
||||
.tx
|
||||
.send(Job::GetLeaderComm { reply: reply_tx })
|
||||
.is_err()
|
||||
{
|
||||
return None;
|
||||
}
|
||||
reply_rx.await.ok().flatten()
|
||||
}
|
||||
|
||||
/// Load a GGUF (pre-quantized) single-GPU model on the worker
|
||||
/// thread. The hf-hub resolution happens on the async caller; the
|
||||
/// resolved local `gguf_path` plus the spec's model_id are sent
|
||||
@@ -572,6 +593,50 @@ impl DeviceWorkerHandle {
|
||||
}
|
||||
}
|
||||
|
||||
/// Image-bearing TP leader forward (single-shot vision prefill).
|
||||
/// Routes `Job::TpForwardLogitsWithImages` onto the worker thread;
|
||||
/// the handler preprocesses + encodes + splices + forwards and
|
||||
/// returns CPU-side `[vocab]` logits. The `WorkerPool` fans the
|
||||
/// matching `GenerateStepWithImages` out to subprocess ranks so the
|
||||
/// row-parallel collectives complete.
|
||||
#[cfg(feature = "cuda")]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn tp_forward_logits_with_images(
|
||||
&self,
|
||||
handle: TpHandle,
|
||||
tokens: Vec<u32>,
|
||||
offset: usize,
|
||||
image_token_id: u32,
|
||||
image_data_uris: Vec<String>,
|
||||
chunk_size: usize,
|
||||
) -> Result<Vec<f32>, WorkerError> {
|
||||
if self.poisoned.load(Ordering::Acquire) {
|
||||
return Err(WorkerError::Poisoned {
|
||||
device_index: self.device_index,
|
||||
});
|
||||
}
|
||||
let (reply_tx, reply_rx) = oneshot::channel();
|
||||
self.tx
|
||||
.send(Job::TpForwardLogitsWithImages {
|
||||
handle,
|
||||
tokens,
|
||||
offset,
|
||||
image_token_id,
|
||||
image_data_uris,
|
||||
chunk_size,
|
||||
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,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Send `Job::Shutdown` and join the thread. Idempotent — calling
|
||||
/// twice is a no-op the second time.
|
||||
pub fn shutdown(&self) -> anyhow::Result<()> {
|
||||
|
||||
@@ -114,10 +114,8 @@ impl HarnessRegistry {
|
||||
for config in configs {
|
||||
match config.name.as_str() {
|
||||
"candle" => {
|
||||
let harness = Arc::new(candle::CandleHarness::new(
|
||||
bind_url.to_string(),
|
||||
&settings.candle,
|
||||
));
|
||||
let harness =
|
||||
candle::CandleHarness::new(bind_url.to_string(), &settings.candle);
|
||||
registry.candle = Some(Arc::clone(&harness));
|
||||
registry.harnesses.insert("candle".into(), harness);
|
||||
}
|
||||
|
||||
@@ -2,11 +2,11 @@
|
||||
//!
|
||||
//! Decodes `data:image/...;base64,...` URIs from OpenAI-style
|
||||
//! `image_url` content parts into the patch tensors a candle vision
|
||||
//! tower expects. Stage A ships **fixed resolution** — every image
|
||||
//! is resized to the same target dimensions (default 448×448 for
|
||||
//! Qwen3.6, configurable per-call) so the patch count is constant
|
||||
//! per image. Variable resolution per [Qwen2VL convention] is tracked
|
||||
//! as issue #14.
|
||||
//! tower expects. Resolution is **dynamic** (#14): each image is
|
||||
//! resized to its native aspect via Qwen `smart_resize` — a
|
||||
//! factor-aligned `(h, w)` whose pixel count lands in the profile's
|
||||
//! `[min_pixels, max_pixels]` budget — so the LM token count varies per
|
||||
//! image (`(h/factor) × (w/factor)`).
|
||||
//!
|
||||
//! Spec reference: `doc/vision-qwen3_6-spec.md` — preprocessor
|
||||
//! section.
|
||||
@@ -21,7 +21,7 @@
|
||||
//! Pipeline (per image):
|
||||
//! 1. data: URI → base64 decode → bytes
|
||||
//! 2. bytes → image::DynamicImage (PNG/JPEG/WebP/etc)
|
||||
//! 3. resize_exact to target H×W (pixel space)
|
||||
//! 3. smart_resize to a native-aspect, factor-aligned H×W (pixel space)
|
||||
//! 4. RGB→f32, normalise per mean/std
|
||||
//! 5. layout to (C, H, W) tensor
|
||||
//!
|
||||
@@ -34,39 +34,126 @@ use base64::Engine;
|
||||
use image::DynamicImage;
|
||||
use image::imageops::FilterType;
|
||||
|
||||
/// Preprocessing target. Captures the resize dimensions and the
|
||||
/// channel-wise normalisation constants from the model's
|
||||
/// `preprocessor_config.json`. Stage A ships a single `qwen3_6()`
|
||||
/// constructor for fixed-resolution Qwen3.6 preprocessing; other
|
||||
/// models can ship their own profile when added.
|
||||
/// Preprocessing target. Captures the resize policy (Qwen `smart_resize`
|
||||
/// factor + pixel budget) and the channel-wise normalisation constants
|
||||
/// from the model's `preprocessor_config.json`. Images are resized to
|
||||
/// their **native aspect** — a factor-aligned `(h, w)` whose pixel count
|
||||
/// lands in `[min_pixels, max_pixels]` — not a fixed square (#14).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PreprocessProfile {
|
||||
pub target_height: u32,
|
||||
pub target_width: u32,
|
||||
/// Both output dims are multiples of this. For Qwen3.6 it is
|
||||
/// `patch_size(16) × spatial_merge_size(2) = 32`, so the post-merge
|
||||
/// LM grid is exactly `(h/factor, w/factor)`.
|
||||
pub factor: u32,
|
||||
/// Lower pixel bound — tiny images are upscaled to at least this.
|
||||
pub min_pixels: u32,
|
||||
/// Upper pixel bound — large images are downscaled to at most this.
|
||||
/// Caps per-image LM tokens (`max_pixels / factor²`) and the
|
||||
/// O(patches²) ViT attention cost.
|
||||
pub max_pixels: u32,
|
||||
pub image_mean: [f32; 3],
|
||||
pub image_std: [f32; 3],
|
||||
}
|
||||
|
||||
/// The Qwen3.6 vision tower rejects any image whose **patch** count
|
||||
/// exceeds its learned pos-embed budget (`num_position_embeddings =
|
||||
/// 2304 = 48²`; see `vision.rs`). At `patch_size = 16` that is
|
||||
/// `2304 × 16² = 589_824` source pixels. `max_pixels` is hard-capped to
|
||||
/// this so `smart_resize` can never produce an over-budget grid — a
|
||||
/// per-rank "patch count exceeds pos_embed budget" error mid-TP-forward
|
||||
/// would otherwise poison the device context. The pos-embed grid is the
|
||||
/// resolution Qwen3.6 was trained at, so this cap is principled, not just
|
||||
/// defensive.
|
||||
const QWEN3_6_MAX_PIXELS_CAP: u32 = 2304 * 16 * 16; // 589_824 → ≤ 2304 patches → ≤ 576 LM tokens
|
||||
|
||||
/// Default pixel budget for Qwen3.6: `256²` (64 LM tokens) up to the
|
||||
/// pos-embed cap (576 LM tokens). Generous for documents/OCR, bounded
|
||||
/// for serving. Operators lower it with `NEURON_VISION_MIN_PIXELS` /
|
||||
/// `NEURON_VISION_MAX_PIXELS` (the upper bound is still clamped to the
|
||||
/// cap above — raising it past the budget would poison the model).
|
||||
const QWEN3_6_MIN_PIXELS: u32 = 65_536;
|
||||
|
||||
fn env_pixels(name: &str, default: u32) -> u32 {
|
||||
std::env::var(name)
|
||||
.ok()
|
||||
.and_then(|v| v.trim().parse::<u32>().ok())
|
||||
.unwrap_or(default)
|
||||
}
|
||||
|
||||
impl PreprocessProfile {
|
||||
/// Stage A profile for Qwen3.6. Resize to 448×448, normalise to
|
||||
/// `[-1, 1]` via mean=std=0.5. Fits within the model's
|
||||
/// `num_position_embeddings=2304` budget at 28×28 = 784 patches
|
||||
/// before merging.
|
||||
/// Profile for Qwen3.6. Native-aspect `smart_resize` (factor 32),
|
||||
/// normalise to `[-1, 1]` via mean=std=0.5. Pixel budget defaults to
|
||||
/// [`QWEN3_6_MIN_PIXELS`]…[`QWEN3_6_MAX_PIXELS_CAP`], overridable via
|
||||
/// `NEURON_VISION_MIN_PIXELS` / `NEURON_VISION_MAX_PIXELS`. Clamped
|
||||
/// sane: `factor² ≤ min ≤ max`, and `max ≤` the pos-embed cap (so the
|
||||
/// vision tower never rejects a resized image and poisons the context).
|
||||
pub fn qwen3_6() -> Self {
|
||||
let factor = 32u32;
|
||||
let f2 = factor * factor;
|
||||
let min_pixels = env_pixels("NEURON_VISION_MIN_PIXELS", QWEN3_6_MIN_PIXELS)
|
||||
.max(f2)
|
||||
.min(QWEN3_6_MAX_PIXELS_CAP);
|
||||
let max_pixels = env_pixels("NEURON_VISION_MAX_PIXELS", QWEN3_6_MAX_PIXELS_CAP)
|
||||
.min(QWEN3_6_MAX_PIXELS_CAP)
|
||||
.max(min_pixels);
|
||||
Self {
|
||||
target_height: 448,
|
||||
target_width: 448,
|
||||
factor,
|
||||
min_pixels,
|
||||
max_pixels,
|
||||
image_mean: [0.5, 0.5, 0.5],
|
||||
image_std: [0.5, 0.5, 0.5],
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-channel CHW tensor length: 3 * H * W.
|
||||
pub fn pixels_chw(&self) -> usize {
|
||||
3 * (self.target_height as usize) * (self.target_width as usize)
|
||||
/// The factor-aligned `(h, w)` this profile would resize a source
|
||||
/// `src_h × src_w` image to. Pure integer policy — no pixel work.
|
||||
pub fn resized_dims(&self, src_h: u32, src_w: u32) -> Result<(u32, u32)> {
|
||||
smart_resize(src_h, src_w, self.factor, self.min_pixels, self.max_pixels)
|
||||
}
|
||||
}
|
||||
|
||||
/// Qwen `smart_resize`: the smallest `factor`-aligned `(h_bar, w_bar)`
|
||||
/// that preserves aspect ratio as closely as possible while keeping the
|
||||
/// pixel count within `[min_pixels, max_pixels]`. Direct port of the
|
||||
/// canonical Qwen2-VL / Qwen3-VL image-processor function (so neuron's
|
||||
/// grid matches what the model was trained on).
|
||||
///
|
||||
/// Returns `(height, width)`. Errors if the aspect ratio exceeds 200:1
|
||||
/// (degenerate input — a 1-pixel-tall strip), matching upstream.
|
||||
pub fn smart_resize(
|
||||
height: u32,
|
||||
width: u32,
|
||||
factor: u32,
|
||||
min_pixels: u32,
|
||||
max_pixels: u32,
|
||||
) -> Result<(u32, u32)> {
|
||||
let h = height.max(1) as f64;
|
||||
let w = width.max(1) as f64;
|
||||
let ratio = h.max(w) / h.min(w);
|
||||
if ratio > 200.0 {
|
||||
anyhow::bail!(
|
||||
"image aspect ratio {ratio:.1}:1 exceeds the 200:1 limit ({height}×{width}); \
|
||||
refusing to resize"
|
||||
);
|
||||
}
|
||||
let f = factor as f64;
|
||||
let (minp, maxp) = (min_pixels as f64, max_pixels as f64);
|
||||
// round-to-nearest-factor (may be 0 for sub-factor inputs; the
|
||||
// min-pixels branch below grows it back up).
|
||||
let mut h_bar = (h / f).round() * f;
|
||||
let mut w_bar = (w / f).round() * f;
|
||||
if h_bar * w_bar > maxp {
|
||||
let beta = (h * w / maxp).sqrt();
|
||||
h_bar = f.max((h / beta / f).floor() * f);
|
||||
w_bar = f.max((w / beta / f).floor() * f);
|
||||
} else if h_bar * w_bar < minp {
|
||||
let beta = (minp / (h * w)).sqrt();
|
||||
h_bar = (h * beta / f).ceil() * f;
|
||||
w_bar = (w * beta / f).ceil() * f;
|
||||
}
|
||||
Ok((h_bar as u32, w_bar as u32))
|
||||
}
|
||||
|
||||
/// Decode a `data:image/...;base64,...` URI into an in-memory image.
|
||||
///
|
||||
/// Accepts the OpenAI Chat Completions `image_url` shape — a string
|
||||
@@ -106,16 +193,13 @@ pub fn decode_data_uri(uri: &str) -> Result<DynamicImage> {
|
||||
/// faster on CPU. Quality difference is marginal for downstream
|
||||
/// vision-encoder consumption. The numerical-validation issue (#15)
|
||||
/// will quantify any discrepancy.
|
||||
pub fn preprocess(img: &DynamicImage, profile: &PreprocessProfile) -> Vec<f32> {
|
||||
pub fn preprocess(img: &DynamicImage, profile: &PreprocessProfile) -> Result<(Vec<f32>, u32, u32)> {
|
||||
let (h_bar, w_bar) = profile.resized_dims(img.height(), img.width())?;
|
||||
let rgb = img
|
||||
.resize_exact(
|
||||
profile.target_width,
|
||||
profile.target_height,
|
||||
FilterType::Triangle,
|
||||
)
|
||||
.resize_exact(w_bar, h_bar, FilterType::Triangle)
|
||||
.to_rgb8();
|
||||
let h = profile.target_height as usize;
|
||||
let w = profile.target_width as usize;
|
||||
let h = h_bar as usize;
|
||||
let w = w_bar as usize;
|
||||
let mut out = vec![0.0_f32; 3 * h * w];
|
||||
// Row-major (C, H, W). Candle's Conv2d expects NCHW, so this is
|
||||
// the natural layout — the caller stacks `n` of these along the
|
||||
@@ -131,16 +215,27 @@ pub fn preprocess(img: &DynamicImage, profile: &PreprocessProfile) -> Vec<f32> {
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
Ok((out, h_bar, w_bar))
|
||||
}
|
||||
|
||||
/// Combined helper: decode + preprocess in one call. Most call
|
||||
/// sites just want the final tensor; the two-step path exists for
|
||||
/// callers (tests, future video preprocessing) that need the
|
||||
/// Combined helper: decode + preprocess in one call. Returns the
|
||||
/// `(3, h, w)` row-major pixels plus the resized `(h, w)` — the caller
|
||||
/// needs the dims to build the tensor and to derive the LM token grid
|
||||
/// `(h/factor, w/factor)`. Most call sites use this; the two-step path
|
||||
/// exists for callers (tests, future video preprocessing) that need the
|
||||
/// intermediate `DynamicImage`.
|
||||
pub fn preprocess_data_uri(uri: &str, profile: &PreprocessProfile) -> Result<Vec<f32>> {
|
||||
pub fn preprocess_data_uri(uri: &str, profile: &PreprocessProfile) -> Result<(Vec<f32>, u32, u32)> {
|
||||
let img = decode_data_uri(uri)?;
|
||||
Ok(preprocess(&img, profile))
|
||||
preprocess(&img, profile)
|
||||
}
|
||||
|
||||
/// Resized `(h, w)` for a data-URI image **without** running the pixel
|
||||
/// normalisation — decode header + `smart_resize` only. Lets a caller
|
||||
/// that just needs the LM token count (e.g. the TP leader expanding the
|
||||
/// prompt) avoid materialising the full pixel tensor twice.
|
||||
pub fn resized_dims_for_uri(uri: &str, profile: &PreprocessProfile) -> Result<(u32, u32)> {
|
||||
let img = decode_data_uri(uri)?;
|
||||
profile.resized_dims(img.height(), img.width())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -205,13 +300,17 @@ mod tests {
|
||||
// decoding so this test isolates the resize+normalise path.
|
||||
let img: ImageBuffer<Rgb<u8>, Vec<u8>> = ImageBuffer::from_pixel(2, 2, Rgb([255, 0, 0]));
|
||||
let dyn_img = DynamicImage::ImageRgb8(img);
|
||||
let out = preprocess(&dyn_img, &profile);
|
||||
let (out, h_bar, w_bar) = preprocess(&dyn_img, &profile).expect("preprocess");
|
||||
|
||||
assert_eq!(out.len(), profile.pixels_chw());
|
||||
let h = h_bar as usize;
|
||||
let w = w_bar as usize;
|
||||
assert_eq!(out.len(), 3 * h * w);
|
||||
// Dims are factor-aligned and at least the min-pixel floor.
|
||||
assert_eq!(h_bar % profile.factor, 0);
|
||||
assert_eq!(w_bar % profile.factor, 0);
|
||||
assert!(h * w >= profile.min_pixels as usize);
|
||||
// After mean=0.5, std=0.5: red channel (255/255=1.0) → (1.0 - 0.5)/0.5 = 1.0
|
||||
// green/blue (0.0) → (0.0 - 0.5)/0.5 = -1.0
|
||||
let h = profile.target_height as usize;
|
||||
let w = profile.target_width as usize;
|
||||
assert!(
|
||||
(out[0] - 1.0).abs() < 1e-5,
|
||||
"R[0] should be 1.0, got {}",
|
||||
@@ -229,9 +328,12 @@ mod tests {
|
||||
#[test]
|
||||
fn preprocess_data_uri_end_to_end() {
|
||||
let profile = PreprocessProfile::qwen3_6();
|
||||
let out = preprocess_data_uri(&red_png_uri(), &profile).expect("e2e preprocess");
|
||||
assert_eq!(out.len(), profile.pixels_chw());
|
||||
let (out, h, w) = preprocess_data_uri(&red_png_uri(), &profile).expect("e2e preprocess");
|
||||
assert_eq!(out.len(), 3 * h as usize * w as usize);
|
||||
assert!(out.iter().all(|v| v.is_finite()));
|
||||
// resized_dims_for_uri agrees with the full preprocess.
|
||||
let (h2, w2) = resized_dims_for_uri(&red_png_uri(), &profile).expect("dims");
|
||||
assert_eq!((h, w), (h2, w2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -240,10 +342,10 @@ mod tests {
|
||||
// 1x1 grayscale = 200 → after conversion to RGB, all three
|
||||
// channels equal 200, normalised → (200/255 - 0.5)/0.5 ≈ 0.569
|
||||
let gray = DynamicImage::ImageLuma8(ImageBuffer::from_pixel(1, 1, image::Luma([200])));
|
||||
let out = preprocess(&gray, &profile);
|
||||
let (out, h_bar, w_bar) = preprocess(&gray, &profile).expect("preprocess");
|
||||
let expected = ((200.0 / 255.0) - 0.5) / 0.5;
|
||||
let h = profile.target_height as usize;
|
||||
let w = profile.target_width as usize;
|
||||
let h = h_bar as usize;
|
||||
let w = w_bar as usize;
|
||||
for c in 0..3 {
|
||||
let v = out[c * h * w];
|
||||
assert!(
|
||||
@@ -252,4 +354,88 @@ mod tests {
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn smart_resize_keeps_factor_aligned_square_in_budget() {
|
||||
// 448×448 sits inside [65536, 1048576] and is factor-aligned →
|
||||
// unchanged. (Regression guard for the old fixed-res sweet spot.)
|
||||
let (h, w) = smart_resize(448, 448, 32, 65_536, 1_048_576).unwrap();
|
||||
assert_eq!((h, w), (448, 448));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn smart_resize_preserves_aspect_and_caps_at_max() {
|
||||
// 3000×4000 (landscape) → downscaled under max_pixels, aspect kept.
|
||||
let (h, w) = smart_resize(3000, 4000, 32, 65_536, 1_048_576).unwrap();
|
||||
assert_eq!(h % 32, 0);
|
||||
assert_eq!(w % 32, 0);
|
||||
assert!(
|
||||
(h as u64) * (w as u64) <= 1_048_576,
|
||||
"must respect max_pixels"
|
||||
);
|
||||
assert!(w > h, "landscape orientation preserved");
|
||||
// aspect ≈ 4000/3000 = 1.333; allow a factor-rounding tolerance.
|
||||
let ar = w as f64 / h as f64;
|
||||
assert!((ar - 4.0 / 3.0).abs() < 0.15, "aspect ~4:3, got {ar:.3}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn smart_resize_floors_tiny_image_at_min() {
|
||||
// 16×16 → upscaled to at least min_pixels, factor-aligned.
|
||||
let (h, w) = smart_resize(16, 16, 32, 65_536, 1_048_576).unwrap();
|
||||
assert_eq!(h % 32, 0);
|
||||
assert_eq!(w % 32, 0);
|
||||
assert!((h as u64) * (w as u64) >= 65_536, "must respect min_pixels");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn smart_resize_tall_nonsquare_stays_nonsquare() {
|
||||
// A tall screenshot keeps portrait orientation.
|
||||
let (h, w) = smart_resize(2000, 500, 32, 65_536, 1_048_576).unwrap();
|
||||
assert!(h > w, "portrait orientation preserved");
|
||||
assert_eq!(h % 32, 0);
|
||||
assert_eq!(w % 32, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn smart_resize_rejects_extreme_aspect() {
|
||||
let err = smart_resize(1, 500, 32, 65_536, 1_048_576).unwrap_err();
|
||||
assert!(format!("{err:#}").contains("200:1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn qwen3_6_never_exceeds_pos_embed_patch_budget() {
|
||||
// The pos-embed cap must hold for huge, tall, wide, and extreme
|
||||
// images — exceeding 2304 patches errors mid-tower and poisons
|
||||
// the device context, so this invariant is load-bearing.
|
||||
let p = PreprocessProfile::qwen3_6();
|
||||
for (sh, sw) in [
|
||||
(8000u32, 6000u32),
|
||||
(808, 1600),
|
||||
(4000, 400),
|
||||
(1, 199),
|
||||
(16, 16),
|
||||
] {
|
||||
let (h, w) = p.resized_dims(sh, sw).unwrap();
|
||||
let patches = (h / 16) * (w / 16);
|
||||
assert!(
|
||||
patches <= 2304,
|
||||
"{sh}x{sw} → {h}x{w} = {patches} patches exceeds the 2304 budget"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn qwen3_6_default_budget_bounds_lm_tokens() {
|
||||
// A huge source image caps at max_pixels → the per-image LM token
|
||||
// count stays within budget (so it can't blow NEURON_MAX_PROMPT_TOKENS).
|
||||
let p = PreprocessProfile::qwen3_6();
|
||||
let (h, w) = p.resized_dims(8000, 6000).unwrap();
|
||||
let lm_tokens = (h / p.factor) * (w / p.factor);
|
||||
let budget = p.max_pixels / (p.factor * p.factor);
|
||||
assert!(
|
||||
lm_tokens <= budget,
|
||||
"max-res image LM tokens {lm_tokens} must stay within budget {budget}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
154
crates/neuron/src/harness/testdata/qwen3_6_chat_template.jinja
vendored
Normal file
154
crates/neuron/src/harness/testdata/qwen3_6_chat_template.jinja
vendored
Normal file
@@ -0,0 +1,154 @@
|
||||
{%- set image_count = namespace(value=0) %}
|
||||
{%- set video_count = namespace(value=0) %}
|
||||
{%- macro render_content(content, do_vision_count, is_system_content=false) %}
|
||||
{%- if content is string %}
|
||||
{{- content }}
|
||||
{%- elif content is iterable and content is not mapping %}
|
||||
{%- for item in content %}
|
||||
{%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
|
||||
{%- if is_system_content %}
|
||||
{{- raise_exception('System message cannot contain images.') }}
|
||||
{%- endif %}
|
||||
{%- if do_vision_count %}
|
||||
{%- set image_count.value = image_count.value + 1 %}
|
||||
{%- endif %}
|
||||
{%- if add_vision_id %}
|
||||
{{- 'Picture ' ~ image_count.value ~ ': ' }}
|
||||
{%- endif %}
|
||||
{{- '<|vision_start|><|image_pad|><|vision_end|>' }}
|
||||
{%- elif 'video' in item or item.type == 'video' %}
|
||||
{%- if is_system_content %}
|
||||
{{- raise_exception('System message cannot contain videos.') }}
|
||||
{%- endif %}
|
||||
{%- if do_vision_count %}
|
||||
{%- set video_count.value = video_count.value + 1 %}
|
||||
{%- endif %}
|
||||
{%- if add_vision_id %}
|
||||
{{- 'Video ' ~ video_count.value ~ ': ' }}
|
||||
{%- endif %}
|
||||
{{- '<|vision_start|><|video_pad|><|vision_end|>' }}
|
||||
{%- elif 'text' in item %}
|
||||
{{- item.text }}
|
||||
{%- else %}
|
||||
{{- raise_exception('Unexpected item type in content.') }}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- elif content is none or content is undefined %}
|
||||
{{- '' }}
|
||||
{%- else %}
|
||||
{{- raise_exception('Unexpected content type.') }}
|
||||
{%- endif %}
|
||||
{%- endmacro %}
|
||||
{%- if not messages %}
|
||||
{{- raise_exception('No messages provided.') }}
|
||||
{%- endif %}
|
||||
{%- if tools and tools is iterable and tools is not mapping %}
|
||||
{{- '<|im_start|>system\n' }}
|
||||
{{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
|
||||
{%- for tool in tools %}
|
||||
{{- "\n" }}
|
||||
{{- tool | tojson }}
|
||||
{%- endfor %}
|
||||
{{- "\n</tools>" }}
|
||||
{{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
|
||||
{%- if messages[0].role == 'system' %}
|
||||
{%- set content = render_content(messages[0].content, false, true)|trim %}
|
||||
{%- if content %}
|
||||
{{- '\n\n' + content }}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
{{- '<|im_end|>\n' }}
|
||||
{%- else %}
|
||||
{%- if messages[0].role == 'system' %}
|
||||
{%- set content = render_content(messages[0].content, false, true)|trim %}
|
||||
{{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
|
||||
{%- for message in messages[::-1] %}
|
||||
{%- set index = (messages|length - 1) - loop.index0 %}
|
||||
{%- if ns.multi_step_tool and message.role == "user" %}
|
||||
{%- set content = render_content(message.content, false)|trim %}
|
||||
{%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
|
||||
{%- set ns.multi_step_tool = false %}
|
||||
{%- set ns.last_query_index = index %}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- if ns.multi_step_tool %}
|
||||
{{- raise_exception('No user query found in messages.') }}
|
||||
{%- endif %}
|
||||
{%- for message in messages %}
|
||||
{%- set content = render_content(message.content, true)|trim %}
|
||||
{%- if message.role == "system" %}
|
||||
{%- if not loop.first %}
|
||||
{{- raise_exception('System message must be at the beginning.') }}
|
||||
{%- endif %}
|
||||
{%- elif message.role == "user" %}
|
||||
{{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
|
||||
{%- elif message.role == "assistant" %}
|
||||
{%- set reasoning_content = '' %}
|
||||
{%- if message.reasoning_content is string %}
|
||||
{%- set reasoning_content = message.reasoning_content %}
|
||||
{%- else %}
|
||||
{%- if '</think>' in content %}
|
||||
{%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
|
||||
{%- set content = content.split('</think>')[-1].lstrip('\n') %}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
{%- set reasoning_content = reasoning_content|trim %}
|
||||
{%- if (preserve_thinking is defined and preserve_thinking is true) or (loop.index0 > ns.last_query_index) %}
|
||||
{{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
|
||||
{%- else %}
|
||||
{{- '<|im_start|>' + message.role + '\n' + content }}
|
||||
{%- endif %}
|
||||
{%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
|
||||
{%- for tool_call in message.tool_calls %}
|
||||
{%- if tool_call.function is defined %}
|
||||
{%- set tool_call = tool_call.function %}
|
||||
{%- endif %}
|
||||
{%- if loop.first %}
|
||||
{%- if content|trim %}
|
||||
{{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
|
||||
{%- else %}
|
||||
{{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
|
||||
{%- endif %}
|
||||
{%- else %}
|
||||
{{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
|
||||
{%- endif %}
|
||||
{%- if tool_call.arguments is defined %}
|
||||
{%- for args_name, args_value in tool_call.arguments|items %}
|
||||
{{- '<parameter=' + args_name + '>\n' }}
|
||||
{%- set args_value = args_value | string if args_value is string else args_value | tojson | safe %}
|
||||
{{- args_value }}
|
||||
{{- '\n</parameter>\n' }}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
{{- '</function>\n</tool_call>' }}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
{{- '<|im_end|>\n' }}
|
||||
{%- elif message.role == "tool" %}
|
||||
{%- if loop.previtem and loop.previtem.role != "tool" %}
|
||||
{{- '<|im_start|>user' }}
|
||||
{%- endif %}
|
||||
{{- '\n<tool_response>\n' }}
|
||||
{{- content }}
|
||||
{{- '\n</tool_response>' }}
|
||||
{%- if not loop.last and loop.nextitem.role != "tool" %}
|
||||
{{- '<|im_end|>\n' }}
|
||||
{%- elif loop.last %}
|
||||
{{- '<|im_end|>\n' }}
|
||||
{%- endif %}
|
||||
{%- else %}
|
||||
{{- raise_exception('Unexpected message role.') }}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- if add_generation_prompt %}
|
||||
{{- '<|im_start|>assistant\n' }}
|
||||
{%- if enable_thinking is defined and enable_thinking is false %}
|
||||
{{- '<think>\n\n</think>\n\n' }}
|
||||
{%- else %}
|
||||
{{- '<think>\n' }}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
@@ -62,6 +62,30 @@ impl TpLeaderModel {
|
||||
}
|
||||
}
|
||||
|
||||
/// Chunked image prefill on rank 0. Only the vision-capable
|
||||
/// `qwen3_5` arch supports it; the dense `qwen3` arch has no tower.
|
||||
pub fn prefill_with_images_chunked(
|
||||
&mut self,
|
||||
tokens: &[u32],
|
||||
base_offset: usize,
|
||||
image_pixels: &[candle_core::Tensor],
|
||||
image_token_id: u32,
|
||||
chunk_size: usize,
|
||||
) -> candle_core::Result<candle_core::Tensor> {
|
||||
match self {
|
||||
TpLeaderModel::Qwen3_5(m) => m.prefill_with_images_chunked(
|
||||
tokens,
|
||||
base_offset,
|
||||
image_pixels,
|
||||
image_token_id,
|
||||
chunk_size,
|
||||
),
|
||||
TpLeaderModel::Qwen3(_) => {
|
||||
candle_core::bail!("prefill_with_images_chunked: qwen3 (dense) has no vision tower")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn clear_kv_cache(&mut self) {
|
||||
match self {
|
||||
TpLeaderModel::Qwen3(m) => m.clear_kv_cache(),
|
||||
@@ -221,9 +245,67 @@ pub struct WorkerPool {
|
||||
/// Phase 4 the load itself moves onto the worker and that bridge
|
||||
/// goes away.
|
||||
pub(crate) leader_worker: std::sync::Arc<super::device_worker::DeviceWorkerHandle>,
|
||||
/// Cached handle to the leader's NCCL `Comm`, fetched at `init_nccl`
|
||||
/// while the worker thread is responsive. The TP step watchdog uses
|
||||
/// it to `ncclCommAbort` a wedged collective from the async thread —
|
||||
/// the one NCCL op allowed concurrently with an in-flight collective,
|
||||
/// and the only way to unblock the in-process leader thread so
|
||||
/// recovery's `unload` doesn't itself hang (#17 Stage 2). `None` if
|
||||
/// init couldn't cache it; the watchdog then logs that it can't abort.
|
||||
#[cfg(feature = "cuda")]
|
||||
leader_comm: Option<nccl_state::SendComm>,
|
||||
}
|
||||
|
||||
/// Per-step deadline for a TP forward (#17 Stage 2). A healthy decode
|
||||
/// step or chunked prefill completes in well under a second; a wedged
|
||||
/// NCCL collective never returns. Generous default so no legitimate step
|
||||
/// trips it; overridable via `NEURON_TP_STEP_TIMEOUT_S` (seconds).
|
||||
#[cfg(feature = "cuda")]
|
||||
fn tp_step_timeout() -> std::time::Duration {
|
||||
let secs = std::env::var("NEURON_TP_STEP_TIMEOUT_S")
|
||||
.ok()
|
||||
.and_then(|v| v.trim().parse::<u64>().ok())
|
||||
.filter(|&s| s > 0)
|
||||
.unwrap_or(120);
|
||||
std::time::Duration::from_secs(secs)
|
||||
}
|
||||
|
||||
impl WorkerPool {
|
||||
/// Abort the leader's NCCL comm to unblock a collective the watchdog
|
||||
/// found wedged (#17 Stage 2). Logs the whole sequence loudly so a
|
||||
/// real-world hang leaves a greppable forensic trail
|
||||
/// (`tp watchdog:` / `ncclCommAbort`). Calling abort from this async
|
||||
/// thread while the worker thread is blocked inside the collective is
|
||||
/// the one concurrent NCCL op the library sanctions — it is how a
|
||||
/// stuck/failed collective is unblocked.
|
||||
#[cfg(feature = "cuda")]
|
||||
fn watchdog_abort_leader_comm(&self, model_id: &str, secs: u64) {
|
||||
tracing::error!(
|
||||
model = %model_id,
|
||||
timeout_s = secs,
|
||||
"tp watchdog: leader forward exceeded deadline — NCCL collective wedged; \
|
||||
aborting comm to unblock the leader thread for auto-recovery"
|
||||
);
|
||||
match &self.leader_comm {
|
||||
Some(c) => match c.0.abort() {
|
||||
Ok(()) => tracing::error!(
|
||||
model = %model_id,
|
||||
"tp watchdog: ncclCommAbort succeeded — wedged collective unblocked; \
|
||||
failing the step so the model auto-recovers (unload+reload)"
|
||||
),
|
||||
Err(e) => tracing::error!(
|
||||
model = %model_id, error = ?e,
|
||||
"tp watchdog: ncclCommAbort failed — recovery may stall until a process restart"
|
||||
),
|
||||
},
|
||||
None => tracing::error!(
|
||||
model = %model_id,
|
||||
"tp watchdog: no cached leader comm handle — cannot abort; recovery will rely \
|
||||
on a process restart"
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
/// Spawn `world_size - 1` worker subprocesses. Rank 0 is the
|
||||
/// leader (in-process) and is *not* spawned here — the leader
|
||||
/// holds rank 0's NCCL Comm and shard in its own address space.
|
||||
@@ -300,6 +382,8 @@ impl WorkerPool {
|
||||
workers,
|
||||
exe,
|
||||
leader_worker,
|
||||
#[cfg(feature = "cuda")]
|
||||
leader_comm: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -380,6 +464,23 @@ impl WorkerPool {
|
||||
world_size = self.world_size,
|
||||
"NCCL communicator established across all ranks"
|
||||
);
|
||||
|
||||
// Cache the leader's Comm handle now, while the worker thread is
|
||||
// responsive, so the TP step watchdog can abort a wedged
|
||||
// collective later (it can't fetch it then — the thread is stuck).
|
||||
// (#17 Stage 2.)
|
||||
#[cfg(feature = "cuda")]
|
||||
{
|
||||
self.leader_comm = self.leader_worker.get_leader_comm().await;
|
||||
if self.leader_comm.is_some() {
|
||||
tracing::debug!("cached leader NCCL comm handle for the TP step watchdog");
|
||||
} else {
|
||||
tracing::warn!(
|
||||
"could not cache leader NCCL comm handle; the TP step watchdog will be \
|
||||
unable to abort a wedged collective (a hang would need a process restart)"
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -604,10 +705,27 @@ impl WorkerPool {
|
||||
// that's the invariant the whole refactor exists to
|
||||
// preserve.
|
||||
let leader_start = std::time::Instant::now();
|
||||
let leader_result = self
|
||||
let timeout = tp_step_timeout();
|
||||
let leader_fut = self
|
||||
.leader_worker
|
||||
.tp_forward_logits(leader_handle, tokens, offset)
|
||||
.await;
|
||||
.tp_forward_logits(leader_handle, tokens, offset);
|
||||
let leader_result = match tokio::time::timeout(timeout, leader_fut).await {
|
||||
Ok(r) => r,
|
||||
Err(_elapsed) => {
|
||||
// Watchdog (#17 Stage 2): the NCCL collective is wedged.
|
||||
// Abort the leader comm to unblock its thread, then fail
|
||||
// the step WITHOUT draining (the subprocess workers are
|
||||
// wedged too; recovery's unload kills them). The error
|
||||
// poisons the model → auto-recovery, which no longer hangs
|
||||
// because the leader thread is now responsive.
|
||||
self.watchdog_abort_leader_comm(model_id, timeout.as_secs());
|
||||
anyhow::bail!(
|
||||
"tp watchdog: leader forward exceeded {}s deadline; aborted wedged NCCL \
|
||||
comm — model will auto-recover",
|
||||
timeout.as_secs()
|
||||
);
|
||||
}
|
||||
};
|
||||
let leader_ok = leader_result.is_ok();
|
||||
let leader_ms = leader_start.elapsed().as_millis();
|
||||
// Surface the leader's own error at WARN before draining
|
||||
@@ -687,6 +805,146 @@ impl WorkerPool {
|
||||
}
|
||||
}
|
||||
|
||||
/// 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:
|
||||
///
|
||||
/// - subprocess workers get `GenerateStepWithImages` (carrying the
|
||||
/// source `image_data_uris`); each preprocesses + encodes through
|
||||
/// its replicated tower and splices locally;
|
||||
/// - the leader runs the same encode + splice + forward on its
|
||||
/// device worker thread via `tp_forward_logits_with_images`.
|
||||
///
|
||||
/// The row-parallel `AllReduce`s synchronise the ranks exactly as in
|
||||
/// the text path. Because the tower is replicated and the preprocess
|
||||
/// is deterministic, every rank's spliced hidden state matches — no
|
||||
/// embedding broadcast. Only used for prefill; decode reuses
|
||||
/// `generate_step`.
|
||||
#[cfg(feature = "cuda")]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn generate_step_with_images(
|
||||
&mut self,
|
||||
model_id: &str,
|
||||
leader_handle: super::device_worker::TpHandle,
|
||||
tokens: Vec<u32>,
|
||||
offset: usize,
|
||||
image_token_id: u32,
|
||||
image_data_uris: Vec<String>,
|
||||
chunk_size: usize,
|
||||
) -> Result<Vec<f32>> {
|
||||
let step_start = std::time::Instant::now();
|
||||
let tokens_len = tokens.len();
|
||||
tracing::debug!(
|
||||
model = %model_id,
|
||||
tokens = tokens_len,
|
||||
offset,
|
||||
images = image_data_uris.len(),
|
||||
chunk_size,
|
||||
"WorkerPool::generate_step_with_images: fan-out"
|
||||
);
|
||||
|
||||
// 1. Fan-out the image-bearing prefill to subprocess workers.
|
||||
for w in &mut self.workers {
|
||||
w.send_only(&WorkerRequest::GenerateStepWithImages {
|
||||
model_id: model_id.to_string(),
|
||||
tokens: tokens.clone(),
|
||||
offset,
|
||||
image_token_id,
|
||||
image_data_uris: image_data_uris.clone(),
|
||||
chunk_size,
|
||||
})
|
||||
.await?;
|
||||
}
|
||||
|
||||
// 2. Leader's image forward on its device worker thread. The
|
||||
// AllReduce CustomOps block until every worker issues the
|
||||
// matching collective; CPU-side logits keep the device tensor
|
||||
// from escaping the worker thread.
|
||||
let leader_start = std::time::Instant::now();
|
||||
let timeout = tp_step_timeout();
|
||||
let leader_fut = self.leader_worker.tp_forward_logits_with_images(
|
||||
leader_handle,
|
||||
tokens,
|
||||
offset,
|
||||
image_token_id,
|
||||
image_data_uris,
|
||||
chunk_size,
|
||||
);
|
||||
let leader_result = match tokio::time::timeout(timeout, leader_fut).await {
|
||||
Ok(r) => r,
|
||||
Err(_elapsed) => {
|
||||
// Watchdog (#17 Stage 2) — see generate_step. Vision
|
||||
// prefill is still well under the deadline on healthy
|
||||
// hardware; a timeout means a wedged collective.
|
||||
self.watchdog_abort_leader_comm(model_id, timeout.as_secs());
|
||||
anyhow::bail!(
|
||||
"tp watchdog: leader image forward exceeded {}s deadline; aborted wedged \
|
||||
NCCL comm — model will auto-recover",
|
||||
timeout.as_secs()
|
||||
);
|
||||
}
|
||||
};
|
||||
let leader_ok = leader_result.is_ok();
|
||||
let leader_ms = leader_start.elapsed().as_millis();
|
||||
if !leader_ok {
|
||||
let detail = leader_result
|
||||
.as_ref()
|
||||
.err()
|
||||
.map(|e| format!("{e:#}"))
|
||||
.unwrap_or_default();
|
||||
tracing::warn!(
|
||||
model = %model_id,
|
||||
tokens = tokens_len,
|
||||
offset,
|
||||
leader_ms,
|
||||
error = %detail,
|
||||
"WorkerPool::generate_step_with_images: leader forward failed"
|
||||
);
|
||||
}
|
||||
|
||||
// 3. ALWAYS drain worker responses, regardless of the leader's
|
||||
// outcome, so stale GenerateStepOk replies don't poison the
|
||||
// next request's recv (same invariant as generate_step).
|
||||
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;
|
||||
tracing::debug!(
|
||||
model = %model_id,
|
||||
leader_ms,
|
||||
leader_ok,
|
||||
errors = worker_errors.len(),
|
||||
total_ms = step_start.elapsed().as_millis(),
|
||||
"WorkerPool::generate_step_with_images: workers drained"
|
||||
);
|
||||
|
||||
match leader_result {
|
||||
Ok(values) => {
|
||||
if worker_errors.is_empty() {
|
||||
Ok(values)
|
||||
} else {
|
||||
anyhow::bail!(
|
||||
"GenerateStepWithImages: leader succeeded but workers failed: {}",
|
||||
worker_errors.join("; ")
|
||||
)
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
if worker_errors.is_empty() {
|
||||
Err(anyhow::Error::new(e)
|
||||
.context("GenerateStepWithImages: leader forward failed"))
|
||||
} else {
|
||||
Err(anyhow::Error::new(e).context(format!(
|
||||
"GenerateStepWithImages: leader forward failed and workers also failed: {}",
|
||||
worker_errors.join("; ")
|
||||
)))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Reset the KV cache for `model_id` on every rank. Called at the
|
||||
/// start of every inference so a fresh request doesn't attend over
|
||||
/// the previous one's tokens.
|
||||
|
||||
@@ -119,40 +119,25 @@ mod cuda_impl {
|
||||
}
|
||||
}
|
||||
|
||||
/// `Arc<Comm>` doesn't impl `Send` because `Comm` wraps a raw
|
||||
/// `ncclComm_t` pointer. The NCCL contract is "operations against a
|
||||
/// given comm must be serialised", not "the handle must stay on the
|
||||
/// thread that created it" — so it's safe to move an `Arc<Comm>`
|
||||
/// across threads as long as no concurrent ops are issued. The
|
||||
/// pool's outer Mutex serialises us into `spawn_blocking`, so this
|
||||
/// wrapper at the move boundary is the only thing missing.
|
||||
/// Thin newtype over `Arc<Comm>`, kept for call-site clarity — it marks
|
||||
/// the points where a comm handle is intentionally moved across threads
|
||||
/// (e.g. cached async-side for the TP step watchdog's `ncclCommAbort`).
|
||||
///
|
||||
/// `Sync` is also marked safe because the `Arc<Comm>` clones held
|
||||
/// by the row-parallel layers are only used from the
|
||||
/// `spawn_blocking` thread driving the forward pass; concurrent
|
||||
/// access from another thread would still be a bug.
|
||||
/// `Send`/`Sync` are provided upstream by `cudarc`'s `Comm` (which
|
||||
/// asserts the NCCL thread-safety invariant, including aborting from a
|
||||
/// different thread than one inside a collective), so this type derives
|
||||
/// them automatically — no manual `unsafe impl` here.
|
||||
pub struct SendComm(pub Arc<Comm>);
|
||||
|
||||
// SAFETY: see the doc-comment above; the invariant is enforced at
|
||||
// the call site (pool Mutex + single spawn_blocking thread), not at
|
||||
// the type level.
|
||||
unsafe impl Send for SendComm {}
|
||||
unsafe impl Sync for SendComm {}
|
||||
|
||||
impl SendComm {
|
||||
pub fn into_inner(self) -> Arc<Comm> {
|
||||
self.0
|
||||
}
|
||||
}
|
||||
|
||||
// SAFETY: `cudarc::nccl::Comm` contains a raw `ncclComm_t` pointer
|
||||
// (libnccl-allocated state). NCCL requires that operations against
|
||||
// one Comm be issued one at a time; we serialise access by storing
|
||||
// NcclState behind a Mutex in `WorkerPool`. The Comm itself is
|
||||
// move-safe — NCCL doesn't track the calling OS thread, only the
|
||||
// stream the operations are dispatched against.
|
||||
unsafe impl Send for NcclState {}
|
||||
unsafe impl Sync for NcclState {}
|
||||
// `NcclState`'s `Send`/`Sync` are auto-derived: its `Arc<Comm>` and
|
||||
// `Arc<CudaContext>` fields are now `Send`/`Sync` (cudarc asserts the
|
||||
// comm thread-safety invariant), so no manual `unsafe impl` is needed.
|
||||
|
||||
/// Generate a fresh NCCL `Id` and return it hex-encoded. Used by
|
||||
/// the leader to mint the shared communicator id which is then
|
||||
|
||||
@@ -88,6 +88,33 @@ pub enum WorkerRequest {
|
||||
offset: usize,
|
||||
},
|
||||
|
||||
/// Like `GenerateStep` but the prefill carries image content. Every
|
||||
/// rank preprocesses the same `image_data_uris` through its
|
||||
/// *replicated* vision tower, splices the resulting patch embeddings
|
||||
/// at `image_token_id` positions, and runs the forward — the
|
||||
/// row-parallel `AllReduce`s still synchronise every rank. Because
|
||||
/// the tower is replicated and `preprocess_data_uri` is
|
||||
/// deterministic, the spliced hidden state is identical on every
|
||||
/// rank, so no embedding broadcast is needed. Sent only for the
|
||||
/// (single-shot) image-bearing prefill; decode steps use plain
|
||||
/// `GenerateStep`. Worker replies with the same `GenerateStepOk`.
|
||||
GenerateStepWithImages {
|
||||
model_id: String,
|
||||
tokens: Vec<u32>,
|
||||
offset: usize,
|
||||
/// `<|image_pad|>` sentinel id (248056 for Qwen3.6); splice
|
||||
/// target in the expanded token stream.
|
||||
image_token_id: u32,
|
||||
/// Source image data URIs (`data:image/...;base64,...`), one per
|
||||
/// image in prompt order. Each rank decodes + preprocesses these
|
||||
/// identically; tens of KB each, so cheap over the stdin pipe.
|
||||
image_data_uris: Vec<String>,
|
||||
/// Prefill chunk size (tokens). Sent explicitly so every rank
|
||||
/// walks the prompt in identical windows and the per-chunk
|
||||
/// row-parallel collectives stay paired across ranks.
|
||||
chunk_size: usize,
|
||||
},
|
||||
|
||||
/// Reset the KV cache for this model on this rank. Sent at the
|
||||
/// start of every inference so a fresh request doesn't accidentally
|
||||
/// attend over the previous one's tokens.
|
||||
@@ -191,6 +218,33 @@ mod tests {
|
||||
assert_eq!(wire, r#"{"op":"init","comm_id":"deadbeef"}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_generate_step_with_images_round_trip() {
|
||||
let req = WorkerRequest::GenerateStepWithImages {
|
||||
model_id: "Qwen/Qwen3.6-27B".into(),
|
||||
tokens: vec![1, 2, 248056, 3],
|
||||
offset: 0,
|
||||
image_token_id: 248056,
|
||||
image_data_uris: vec!["data:image/png;base64,AAA=".into()],
|
||||
chunk_size: 512,
|
||||
};
|
||||
let wire = serde_json::to_string(&req).unwrap();
|
||||
assert!(wire.contains(r#""op":"generate_step_with_images""#));
|
||||
match roundtrip(&req) {
|
||||
WorkerRequest::GenerateStepWithImages {
|
||||
tokens,
|
||||
image_token_id,
|
||||
image_data_uris,
|
||||
..
|
||||
} => {
|
||||
assert_eq!(tokens, vec![1, 2, 248056, 3]);
|
||||
assert_eq!(image_token_id, 248056);
|
||||
assert_eq!(image_data_uris.len(), 1);
|
||||
}
|
||||
other => panic!("expected GenerateStepWithImages, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_shutdown_round_trip() {
|
||||
assert_eq!(
|
||||
|
||||
@@ -46,6 +46,8 @@ use super::tp_linear::{ColumnParallelLinear, RowParallelLinear};
|
||||
use crate::harness::arch::qwen3_5::linear_attn::repeat_interleave;
|
||||
use crate::harness::arch::qwen3_5::rmsnorm::{Qwen3_5RmsNorm, Qwen3_5RmsNormGated, l2norm};
|
||||
use crate::harness::arch::qwen3_5::rope::RotaryEmbedding;
|
||||
use crate::harness::arch::qwen3_5::splice_runs;
|
||||
use crate::harness::arch::qwen3_5::vision::VisionTower;
|
||||
pub use crate::harness::arch::qwen3_5::{Config, TextConfig};
|
||||
|
||||
// ─── linear-attention (Gated DeltaNet) ──────────────────────────────
|
||||
@@ -524,7 +526,8 @@ impl TpQwen3_5Attention {
|
||||
&mut self,
|
||||
x: &Tensor,
|
||||
attn_mask: Option<&Tensor>,
|
||||
offset: usize,
|
||||
cos: &Tensor,
|
||||
sin: &Tensor,
|
||||
) -> candle_core::Result<Tensor> {
|
||||
let (b, l, _) = x.dims3()?;
|
||||
|
||||
@@ -557,7 +560,7 @@ impl TpQwen3_5Attention {
|
||||
.transpose(1, 2)?
|
||||
.contiguous()?;
|
||||
|
||||
let (q, k) = self.rotary.apply(&q, &k, offset)?;
|
||||
let (q, k) = self.rotary.apply_cos_sin(&q, &k, cos, sin)?;
|
||||
let (k, v) = self.kv_cache.append(&k, &v)?;
|
||||
let k = repeat_kv(k, self.num_kv_groups)?.contiguous()?;
|
||||
let v = repeat_kv(v, self.num_kv_groups)?.contiguous()?;
|
||||
@@ -805,11 +808,12 @@ impl TpQwen3_5DecoderLayer {
|
||||
&mut self,
|
||||
x: &Tensor,
|
||||
attn_mask: Option<&Tensor>,
|
||||
offset: usize,
|
||||
cos: &Tensor,
|
||||
sin: &Tensor,
|
||||
) -> candle_core::Result<Tensor> {
|
||||
let h = self.input_layernorm.forward(x)?;
|
||||
let attn_out = match &mut self.attention {
|
||||
TpAttentionKind::Full(attn) => attn.forward(&h, attn_mask, offset)?,
|
||||
TpAttentionKind::Full(attn) => attn.forward(&h, attn_mask, cos, sin)?,
|
||||
TpAttentionKind::Linear(net) => net.forward(&h)?,
|
||||
};
|
||||
let x = (x + attn_out)?;
|
||||
@@ -832,6 +836,15 @@ pub struct TpQwen3_5Model {
|
||||
embed_tokens: Embedding,
|
||||
layers: Vec<TpQwen3_5DecoderLayer>,
|
||||
norm: Qwen3_5RmsNorm,
|
||||
/// Replicated rotary, shared with every full-attention layer. The
|
||||
/// model builds the per-forward cos/sin (interleaved M-RoPE for image
|
||||
/// tokens, plain for text) once and the layers apply it. Identical on
|
||||
/// every rank, so per-rank position ids stay consistent.
|
||||
rotary: Arc<RotaryEmbedding>,
|
||||
/// `offset + rope_delta` is the text-axis decode position; set from
|
||||
/// `get_rope_index` during a vision prefill, reset in `clear_kv_cache`.
|
||||
/// See `Qwen3_5Model::rope_delta`.
|
||||
rope_delta: i64,
|
||||
device: Device,
|
||||
dtype: DType,
|
||||
}
|
||||
@@ -898,6 +911,8 @@ impl TpQwen3_5Model {
|
||||
embed_tokens,
|
||||
layers,
|
||||
norm,
|
||||
rotary,
|
||||
rope_delta: 0,
|
||||
device,
|
||||
dtype,
|
||||
})
|
||||
@@ -954,6 +969,8 @@ impl TpQwen3_5Model {
|
||||
embed_tokens,
|
||||
layers,
|
||||
norm,
|
||||
rotary,
|
||||
rope_delta: 0,
|
||||
device,
|
||||
dtype,
|
||||
})
|
||||
@@ -967,6 +984,14 @@ impl TpQwen3_5Model {
|
||||
for l in &mut self.layers {
|
||||
l.clear_kv_cache();
|
||||
}
|
||||
self.rope_delta = 0;
|
||||
}
|
||||
|
||||
/// Set the decode `rope_delta` computed by `get_rope_index` during a
|
||||
/// vision prefill, so decode after the image resumes text positions
|
||||
/// from the image-compressed counter.
|
||||
pub fn set_rope_delta(&mut self, delta: i64) {
|
||||
self.rope_delta = delta;
|
||||
}
|
||||
|
||||
fn causal_mask(&self, b: usize, tgt: usize, offset: usize) -> candle_core::Result<Tensor> {
|
||||
@@ -978,15 +1003,88 @@ impl TpQwen3_5Model {
|
||||
}
|
||||
|
||||
pub fn forward(&mut self, input: &Tensor, offset: usize) -> candle_core::Result<Tensor> {
|
||||
self.forward_inner(input, offset, None, None, None)
|
||||
}
|
||||
|
||||
/// Forward for a vision-prefill chunk: optional image-embedding
|
||||
/// splice plus explicit interleaved-M-RoPE `position_ids` (the
|
||||
/// chunk's slice of the full prompt's 3D positions). Used by
|
||||
/// `TpQwen3_5ForCausalLM::prefill_with_images_chunked`, which
|
||||
/// computes the positions once over the whole prompt and slices them
|
||||
/// per chunk so every rank steps in lockstep.
|
||||
pub fn forward_with_positions(
|
||||
&mut self,
|
||||
input: &Tensor,
|
||||
offset: usize,
|
||||
position_ids: &Tensor,
|
||||
image_embeds: Option<&Tensor>,
|
||||
image_token_id: Option<u32>,
|
||||
) -> candle_core::Result<Tensor> {
|
||||
self.forward_inner(
|
||||
input,
|
||||
offset,
|
||||
image_embeds,
|
||||
image_token_id,
|
||||
Some(position_ids),
|
||||
)
|
||||
}
|
||||
|
||||
/// Shared forward. Splices image embeddings at `image_token_id`
|
||||
/// positions when present, then builds the rotary cos/sin — from the
|
||||
/// explicit `position_ids` (interleaved M-RoPE, vision) when given,
|
||||
/// else plain positions at `offset + rope_delta` (text / decode) —
|
||||
/// and runs the sharded decoder stack. The TP replicated-hidden-state
|
||||
/// invariant holds because every rank encodes the same pixels and
|
||||
/// computes the same positions.
|
||||
fn forward_inner(
|
||||
&mut self,
|
||||
input: &Tensor,
|
||||
offset: usize,
|
||||
image_embeds: Option<&Tensor>,
|
||||
image_token_id: Option<u32>,
|
||||
position_ids: Option<&Tensor>,
|
||||
) -> candle_core::Result<Tensor> {
|
||||
let (b, l) = input.dims2()?;
|
||||
let mut h = self.embed_tokens.forward(input)?;
|
||||
|
||||
if let (Some(img), Some(tok_id)) = (image_embeds, image_token_id) {
|
||||
let ids: Vec<u32> = input.flatten_all()?.to_vec1()?;
|
||||
let mut positions: Vec<u32> = Vec::with_capacity(img.dim(0)?);
|
||||
for (idx, id) in ids.iter().enumerate() {
|
||||
if *id == tok_id {
|
||||
positions.push(idx as u32);
|
||||
}
|
||||
}
|
||||
let n_img_tokens = img.dim(0)?;
|
||||
if positions.len() != n_img_tokens {
|
||||
candle_core::bail!(
|
||||
"TP forward: chunk has {} image-token positions but image_embeds carries \
|
||||
{} tokens — patch-count expansion / chunk slicing mismatch",
|
||||
positions.len(),
|
||||
n_img_tokens,
|
||||
);
|
||||
}
|
||||
if !positions.is_empty() {
|
||||
let img = img.to_dtype(self.dtype)?;
|
||||
h = splice_runs(&h, &img, &positions)?;
|
||||
}
|
||||
}
|
||||
|
||||
let (cos, sin) = match position_ids {
|
||||
Some(pos) => self.rotary.mrope_cos_sin(pos)?,
|
||||
None => {
|
||||
let base = (offset as i64 + self.rope_delta).max(0) as usize;
|
||||
self.rotary.plain_cos_sin(base, l)?
|
||||
}
|
||||
};
|
||||
|
||||
let causal = if l == 1 {
|
||||
None
|
||||
} else {
|
||||
Some(self.causal_mask(b, l, offset)?)
|
||||
};
|
||||
for layer in &mut self.layers {
|
||||
h = layer.forward(&h, causal.as_ref(), offset)?;
|
||||
h = layer.forward(&h, causal.as_ref(), &cos, &sin)?;
|
||||
}
|
||||
self.norm.forward(&h)
|
||||
}
|
||||
@@ -995,6 +1093,41 @@ impl TpQwen3_5Model {
|
||||
pub struct TpQwen3_5ForCausalLM {
|
||||
base: TpQwen3_5Model,
|
||||
lm_head: super::tp_linear::MaybeQuantLinear,
|
||||
/// Replicated vision tower (TP-vision). Loaded on every rank from
|
||||
/// the full, unsharded `model.visual.*` weights; `None` for
|
||||
/// text-only checkpoints. Each rank encodes the same image
|
||||
/// independently — no sharding, no broadcast — which keeps the
|
||||
/// spliced input embeddings identical across ranks (the
|
||||
/// replicated-hidden-state invariant the sharded layers rely on).
|
||||
vision: Option<VisionTower>,
|
||||
/// `<|image_pad|>` sentinel id (mirrors `Config::image_token_id`);
|
||||
/// the splice target for `forward_with_vision`.
|
||||
image_token_id: Option<u32>,
|
||||
}
|
||||
|
||||
/// Load the replicated vision tower from the unsharded `model.visual.*`
|
||||
/// weights when the config carries a `vision_config` block. Shared by
|
||||
/// the cuda and non-cuda `load` variants. `vb.pp("model.visual")`
|
||||
/// resolves against the same full safetensors every rank mmaps; plain
|
||||
/// `.get()` on a `ShardedVarBuilder` returns the full (replicated)
|
||||
/// tensor, so this loads identically regardless of `world_size`.
|
||||
fn load_replicated_vision_tower(
|
||||
config: &Config,
|
||||
vb: &ShardedVarBuilder,
|
||||
) -> Result<Option<VisionTower>> {
|
||||
match config.vision_config.clone() {
|
||||
Some(vcfg) => {
|
||||
tracing::info!(
|
||||
depth = vcfg.depth,
|
||||
hidden_size = vcfg.hidden_size,
|
||||
"loading qwen3_5 vision tower (TP replicated)"
|
||||
);
|
||||
let tower = VisionTower::load(vcfg, vb.pp("model.visual"))
|
||||
.context("load qwen3_5 vision tower (model.visual.*) [TP replicated]")?;
|
||||
Ok(Some(tower))
|
||||
}
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
impl TpQwen3_5ForCausalLM {
|
||||
@@ -1012,7 +1145,14 @@ impl TpQwen3_5ForCausalLM {
|
||||
let cfg = &config.text_config;
|
||||
let base = TpQwen3_5Model::load(cfg, vb, mmap, rank, world_size, comm, quant)?;
|
||||
let lm_head = build_lm_head(cfg, vb, &base, quant)?;
|
||||
let model = Self { base, lm_head };
|
||||
let vision = load_replicated_vision_tower(&config, vb)?;
|
||||
let image_token_id = config.image_token_id;
|
||||
let model = Self {
|
||||
base,
|
||||
lm_head,
|
||||
vision,
|
||||
image_token_id,
|
||||
};
|
||||
log_construction_complete(cfg, rank, world_size, quant, model.device());
|
||||
Ok(model)
|
||||
}
|
||||
@@ -1029,17 +1169,198 @@ impl TpQwen3_5ForCausalLM {
|
||||
let cfg = &config.text_config;
|
||||
let base = TpQwen3_5Model::load(cfg, vb, mmap, rank, world_size, quant)?;
|
||||
let lm_head = build_lm_head(cfg, vb, &base, quant)?;
|
||||
let model = Self { base, lm_head };
|
||||
let vision = load_replicated_vision_tower(&config, vb)?;
|
||||
let image_token_id = config.image_token_id;
|
||||
let model = Self {
|
||||
base,
|
||||
lm_head,
|
||||
vision,
|
||||
image_token_id,
|
||||
};
|
||||
log_construction_complete(cfg, rank, world_size, quant, model.device());
|
||||
Ok(model)
|
||||
}
|
||||
|
||||
/// True when this TP load materialised a replicated vision tower.
|
||||
/// Drives capability advertising and the Stage 3 vision dispatch.
|
||||
pub fn has_vision(&self) -> bool {
|
||||
self.vision.is_some()
|
||||
}
|
||||
|
||||
/// `<|image_pad|>` sentinel id, when known.
|
||||
pub fn image_token_id(&self) -> Option<u32> {
|
||||
self.image_token_id
|
||||
}
|
||||
|
||||
/// Encode one preprocessed `(C, H, W)` image into LM-side patch
|
||||
/// embeddings `(N_lm, hidden)` via this rank's replicated tower.
|
||||
/// Errors when loaded without a vision tower.
|
||||
pub fn encode_image(&self, image: &Tensor) -> Result<Tensor> {
|
||||
self.vision
|
||||
.as_ref()
|
||||
.ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"encode_image: this TP Qwen3.6 load has no vision tower \
|
||||
(config.json::vision_config absent or weights missing)"
|
||||
)
|
||||
})?
|
||||
.forward(image)
|
||||
}
|
||||
|
||||
pub fn forward(&mut self, input: &Tensor, offset: usize) -> candle_core::Result<Tensor> {
|
||||
let (_, l) = input.dims2()?;
|
||||
let hidden = self.base.forward(input, offset)?;
|
||||
hidden.i((.., l - 1.., ..))?.apply(&self.lm_head)
|
||||
}
|
||||
|
||||
/// Forward for a vision-prefill chunk (optional image splice +
|
||||
/// explicit interleaved-M-RoPE `position_ids`). Mirrors `forward`
|
||||
/// but routes through `TpQwen3_5Model::forward_with_positions`.
|
||||
pub fn forward_with_positions(
|
||||
&mut self,
|
||||
input: &Tensor,
|
||||
offset: usize,
|
||||
position_ids: &Tensor,
|
||||
image_embeds: Option<&Tensor>,
|
||||
image_token_id: Option<u32>,
|
||||
) -> candle_core::Result<Tensor> {
|
||||
let (_, l) = input.dims2()?;
|
||||
let hidden = self.base.forward_with_positions(
|
||||
input,
|
||||
offset,
|
||||
position_ids,
|
||||
image_embeds,
|
||||
image_token_id,
|
||||
)?;
|
||||
hidden.i((.., l - 1.., ..))?.apply(&self.lm_head)
|
||||
}
|
||||
|
||||
/// End-to-end image prefill on one rank: encode each preprocessed
|
||||
/// `(C, H, W)` pixel tensor through this rank's replicated tower,
|
||||
/// concatenate the per-image embeddings along the patch axis, and
|
||||
/// forward with the splice. Shared by the leader (`TpLeaderModel`)
|
||||
/// and the subprocess worker (`WorkerModel`) so every rank runs the
|
||||
/// identical encode → splice → forward and keeps the replicated
|
||||
/// hidden state in lockstep. Returns last-position logits
|
||||
/// `(B, 1, vocab)`, same contract as `forward`.
|
||||
/// Encode every preprocessed `(C,H,W)` image once through this
|
||||
/// rank's replicated tower and concatenate along the patch axis →
|
||||
/// `(sum_patches, hidden)`. Done once per prefill, not per chunk.
|
||||
fn encode_images_concat(&self, image_pixels: &[Tensor]) -> candle_core::Result<Tensor> {
|
||||
let mut per_image = Vec::with_capacity(image_pixels.len());
|
||||
for (idx, img) in image_pixels.iter().enumerate() {
|
||||
let embed = self
|
||||
.encode_image(img)
|
||||
.map_err(|e| candle_core::Error::Msg(format!("encode image[{idx}]: {e:#}")))?;
|
||||
per_image.push(embed);
|
||||
}
|
||||
Tensor::cat(&per_image.iter().collect::<Vec<_>>(), 0)
|
||||
}
|
||||
|
||||
/// Chunked image prefill on one rank. Encodes the image(s) once,
|
||||
/// then walks the (pre-expanded) prompt in `chunk_size`-token
|
||||
/// windows — exactly like the text `chunked_prefill_tp` — splicing
|
||||
/// the patch embeddings into whichever chunk(s) carry `<|image_pad|>`
|
||||
/// positions. Activation memory is bounded by the chunk, not the
|
||||
/// full prompt, so a long vision context no longer single-shot-OOMs.
|
||||
///
|
||||
/// Every rank runs the identical chunk sequence (same `tokens.len()`
|
||||
/// and `chunk_size`), so the row-parallel `AllReduce`s pair up
|
||||
/// chunk-by-chunk across ranks with no extra synchronisation. The KV
|
||||
/// cache accumulates across chunks via the growing offset; only the
|
||||
/// final chunk's last-position logits are returned (intermediate
|
||||
/// chunks just populate the cache, same as the text path).
|
||||
pub fn prefill_with_images_chunked(
|
||||
&mut self,
|
||||
tokens: &[u32],
|
||||
base_offset: usize,
|
||||
image_pixels: &[Tensor],
|
||||
image_token_id: u32,
|
||||
chunk_size: usize,
|
||||
) -> candle_core::Result<Tensor> {
|
||||
if image_pixels.is_empty() {
|
||||
candle_core::bail!("prefill_with_images_chunked: called with zero images");
|
||||
}
|
||||
if tokens.is_empty() {
|
||||
candle_core::bail!("prefill_with_images_chunked: empty prompt");
|
||||
}
|
||||
let chunk_size = chunk_size.max(1);
|
||||
let device = self.device().clone();
|
||||
let image_embeds = self.encode_images_concat(image_pixels)?;
|
||||
|
||||
// Each image's LM grid (lm_gh, lm_gw) = (h/factor, w/factor),
|
||||
// factor = patch×merge. Recomputed per rank from this rank's own
|
||||
// pixel tensors — deterministic, so every rank's grids (and hence
|
||||
// M-RoPE positions) match without crossing the RPC (#14).
|
||||
let factor = self
|
||||
.vision
|
||||
.as_ref()
|
||||
.map(|v| {
|
||||
let c = v.config();
|
||||
c.patch_size * c.spatial_merge_size
|
||||
})
|
||||
.ok_or_else(|| {
|
||||
candle_core::Error::Msg(
|
||||
"prefill_with_images_chunked: loaded without a vision tower".into(),
|
||||
)
|
||||
})?;
|
||||
let grids: Vec<(usize, usize)> = image_pixels
|
||||
.iter()
|
||||
.map(|t| {
|
||||
let (_, h, w) = t.dims3()?;
|
||||
Ok::<(usize, usize), candle_core::Error>((h / factor, w / factor))
|
||||
})
|
||||
.collect::<candle_core::Result<Vec<_>>>()?;
|
||||
|
||||
// Interleaved-M-RoPE 3D position ids for the whole prompt,
|
||||
// computed once and sliced per chunk so every rank assigns image
|
||||
// tokens their grid coordinates (and text after an image resumes
|
||||
// from the compressed counter). `rope_delta` is stored on the base
|
||||
// model for the decode that follows this prefill. Every chunk —
|
||||
// text or image — uses the M-RoPE slice, because each image shifts
|
||||
// the positions of the text around it.
|
||||
let (text, height, width, delta) =
|
||||
crate::harness::arch::qwen3_5::rope::get_rope_index(tokens, image_token_id, &grids)
|
||||
.map_err(|e| candle_core::Error::Msg(format!("get_rope_index: {e}")))?;
|
||||
self.base.set_rope_delta(delta);
|
||||
let full_pos = crate::harness::arch::qwen3_5::rope::mrope_position_tensor(
|
||||
&text, &height, &width, &device,
|
||||
)?;
|
||||
|
||||
let mut last_logits: Option<Tensor> = None;
|
||||
// Rows of `image_embeds` already spliced by earlier chunks. The
|
||||
// `<|image_pad|>` run is contiguous, so chunks consume embedding
|
||||
// rows in order.
|
||||
let mut img_off = 0usize;
|
||||
let mut start = 0usize;
|
||||
while start < tokens.len() {
|
||||
let end = (start + chunk_size).min(tokens.len());
|
||||
let chunk = &tokens[start..end];
|
||||
let input = Tensor::new(chunk, &device)?.unsqueeze(0)?;
|
||||
let pos_slice = full_pos.narrow(1, start, end - start)?;
|
||||
let n_here = chunk.iter().filter(|&&t| t == image_token_id).count();
|
||||
let logits = if n_here == 0 {
|
||||
self.forward_with_positions(&input, base_offset + start, &pos_slice, None, None)?
|
||||
} else {
|
||||
// Splice the next `n_here` patch rows at this chunk's
|
||||
// local image-pad positions.
|
||||
let rows = image_embeds.narrow(0, img_off, n_here)?;
|
||||
img_off += n_here;
|
||||
self.forward_with_positions(
|
||||
&input,
|
||||
base_offset + start,
|
||||
&pos_slice,
|
||||
Some(&rows),
|
||||
Some(image_token_id),
|
||||
)?
|
||||
};
|
||||
last_logits = Some(logits);
|
||||
start = end;
|
||||
}
|
||||
last_logits
|
||||
.ok_or_else(|| candle_core::Error::Msg("prefill_with_images_chunked: no chunks".into()))
|
||||
}
|
||||
|
||||
pub fn clear_kv_cache(&mut self) {
|
||||
self.base.clear_kv_cache();
|
||||
}
|
||||
|
||||
@@ -47,6 +47,34 @@ impl WorkerModel {
|
||||
}
|
||||
}
|
||||
|
||||
/// Chunked image prefill on this rank. Only the vision-capable
|
||||
/// `qwen3_5` arch has a replicated tower; the dense `qwen3` arch
|
||||
/// errors. The returned logits are discarded by the caller (the
|
||||
/// leader samples from its own rank-0 copy) — the value is the NCCL
|
||||
/// collectives the forward issues, chunk by chunk in lockstep with
|
||||
/// the leader.
|
||||
fn prefill_with_images_chunked(
|
||||
&mut self,
|
||||
tokens: &[u32],
|
||||
base_offset: usize,
|
||||
image_pixels: &[candle_core::Tensor],
|
||||
image_token_id: u32,
|
||||
chunk_size: usize,
|
||||
) -> candle_core::Result<candle_core::Tensor> {
|
||||
match self {
|
||||
WorkerModel::Qwen3_5(m) => m.prefill_with_images_chunked(
|
||||
tokens,
|
||||
base_offset,
|
||||
image_pixels,
|
||||
image_token_id,
|
||||
chunk_size,
|
||||
),
|
||||
WorkerModel::Qwen3(_) => {
|
||||
candle_core::bail!("prefill_with_images_chunked: qwen3 (dense) has no vision tower")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn clear_kv_cache(&mut self) {
|
||||
match self {
|
||||
WorkerModel::Qwen3(m) => m.clear_kv_cache(),
|
||||
@@ -167,6 +195,21 @@ impl WorkerState {
|
||||
tokens,
|
||||
offset,
|
||||
} => self.handle_generate_step(&model_id, tokens, offset),
|
||||
WorkerRequest::GenerateStepWithImages {
|
||||
model_id,
|
||||
tokens,
|
||||
offset,
|
||||
image_token_id,
|
||||
image_data_uris,
|
||||
chunk_size,
|
||||
} => self.handle_generate_step_with_images(
|
||||
&model_id,
|
||||
tokens,
|
||||
offset,
|
||||
image_token_id,
|
||||
image_data_uris,
|
||||
chunk_size,
|
||||
),
|
||||
WorkerRequest::ClearKvCache { model_id } => self.handle_clear_kv_cache(&model_id),
|
||||
WorkerRequest::UnloadModel { model_id } => self.handle_unload_model(&model_id),
|
||||
WorkerRequest::Shutdown => WorkerResponse::Bye,
|
||||
@@ -418,6 +461,117 @@ impl WorkerState {
|
||||
}
|
||||
}
|
||||
|
||||
/// 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
|
||||
/// splices + forwards. The logits are discarded (the leader samples
|
||||
/// from rank 0); the row-parallel `AllReduce`s are the point.
|
||||
#[cfg(feature = "cuda")]
|
||||
fn handle_generate_step_with_images(
|
||||
&mut self,
|
||||
model_id: &str,
|
||||
tokens: Vec<u32>,
|
||||
offset: usize,
|
||||
image_token_id: u32,
|
||||
image_data_uris: Vec<String>,
|
||||
chunk_size: usize,
|
||||
) -> WorkerResponse {
|
||||
use crate::harness::preprocess::{PreprocessProfile, preprocess_data_uri};
|
||||
use candle_core::Tensor;
|
||||
|
||||
if image_data_uris.is_empty() {
|
||||
return WorkerResponse::Error {
|
||||
kind: "bad_request".into(),
|
||||
message: "GenerateStepWithImages with zero images".into(),
|
||||
};
|
||||
}
|
||||
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();
|
||||
|
||||
// Preprocess each image identically to the leader so the encoded
|
||||
// embeddings — and thus the spliced hidden state and per-image
|
||||
// grids — match across ranks. Native-aspect `smart_resize` (#14);
|
||||
// deterministic, so each rank derives the same dims.
|
||||
let profile = PreprocessProfile::qwen3_6();
|
||||
let mut pixels: Vec<Tensor> = Vec::with_capacity(image_data_uris.len());
|
||||
for (idx, uri) in image_data_uris.iter().enumerate() {
|
||||
let (px, h, w) = match preprocess_data_uri(uri, &profile) {
|
||||
Ok(p) => p,
|
||||
Err(e) => {
|
||||
return WorkerResponse::Error {
|
||||
kind: "bad_request".into(),
|
||||
message: format!("preprocess image[{idx}]: {e:#}"),
|
||||
};
|
||||
}
|
||||
};
|
||||
match Tensor::from_vec(px, (3, h as usize, w as usize), &device) {
|
||||
Ok(t) => pixels.push(t),
|
||||
Err(e) => {
|
||||
return WorkerResponse::Error {
|
||||
kind: "forward_failed".into(),
|
||||
message: format!("build image[{idx}] tensor: {e}"),
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let start = std::time::Instant::now();
|
||||
tracing::debug!(
|
||||
rank = self.config.rank,
|
||||
model = %model_id,
|
||||
tokens = tokens.len(),
|
||||
offset,
|
||||
images = pixels.len(),
|
||||
chunk_size,
|
||||
"worker GenerateStepWithImages: chunked prefill starting"
|
||||
);
|
||||
// Drop the logits — the leader samples from its own rank-0 copy.
|
||||
// The chunked prefill builds its own per-chunk input tensors.
|
||||
if let Err(e) =
|
||||
model.prefill_with_images_chunked(&tokens, offset, &pixels, image_token_id, chunk_size)
|
||||
{
|
||||
tracing::warn!(
|
||||
rank = self.config.rank,
|
||||
model = %model_id,
|
||||
elapsed_ms = start.elapsed().as_millis(),
|
||||
error = %e,
|
||||
"worker GenerateStepWithImages: forward failed"
|
||||
);
|
||||
return WorkerResponse::Error {
|
||||
kind: "forward_failed".into(),
|
||||
message: format!("TP image forward: {e}"),
|
||||
};
|
||||
}
|
||||
tracing::debug!(
|
||||
rank = self.config.rank,
|
||||
model = %model_id,
|
||||
elapsed_ms = start.elapsed().as_millis(),
|
||||
"worker GenerateStepWithImages: forward done"
|
||||
);
|
||||
WorkerResponse::GenerateStepOk
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "cuda"))]
|
||||
fn handle_generate_step_with_images(
|
||||
&mut self,
|
||||
_model_id: &str,
|
||||
_tokens: Vec<u32>,
|
||||
_offset: usize,
|
||||
_image_token_id: u32,
|
||||
_image_data_uris: Vec<String>,
|
||||
_chunk_size: usize,
|
||||
) -> WorkerResponse {
|
||||
WorkerResponse::Error {
|
||||
kind: "cuda_feature_not_enabled".into(),
|
||||
message: "GenerateStepWithImages requires --features cuda".into(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "cuda")]
|
||||
fn handle_clear_kv_cache(&mut self, model_id: &str) -> WorkerResponse {
|
||||
let Some(model) = self.models.get_mut(model_id) else {
|
||||
|
||||
@@ -646,6 +646,54 @@ mod tests {
|
||||
assert_eq!(parts[1]["image_url"]["url"], "data:image/png;base64,AAA=");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multiple_images_translate_in_order_and_tolerate_detail() {
|
||||
// C2: a Responses request carrying several InputImage parts
|
||||
// (with `detail` set) must translate to a chat Parts array that
|
||||
// preserves image order and the `image_url.url` shape the chat
|
||||
// vision path (`extract_images_from_request`) walks. The
|
||||
// `detail` hint has no chat-completions analogue we forward, so
|
||||
// it's dropped — but it must not break translation.
|
||||
let req = ResponsesRequest {
|
||||
model: "m".into(),
|
||||
input: ResponsesInput::Items(vec![ResponsesInputItem::Message {
|
||||
role: "user".into(),
|
||||
content: ResponsesMessageContent::Parts(vec![
|
||||
ResponsesContentPart::InputText {
|
||||
text: "compare these".into(),
|
||||
},
|
||||
ResponsesContentPart::InputImage {
|
||||
image_url: "data:image/png;base64,FIRST".into(),
|
||||
detail: Some("high".into()),
|
||||
},
|
||||
ResponsesContentPart::InputImage {
|
||||
image_url: "data:image/png;base64,SECOND".into(),
|
||||
detail: None,
|
||||
},
|
||||
]),
|
||||
}]),
|
||||
instructions: None,
|
||||
stream: false,
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
previous_response_id: None,
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
let chat = request_to_chat(req).unwrap();
|
||||
let parts = match &chat.messages[0].content {
|
||||
MessageContent::Parts(p) => p.clone(),
|
||||
other => panic!("expected Parts, got {other:?}"),
|
||||
};
|
||||
// text + two images, in input order.
|
||||
assert_eq!(parts.len(), 3);
|
||||
assert_eq!(parts[0]["type"], "text");
|
||||
assert_eq!(parts[1]["image_url"]["url"], "data:image/png;base64,FIRST");
|
||||
assert_eq!(parts[2]["image_url"]["url"], "data:image/png;base64,SECOND");
|
||||
// `detail` is not forwarded into the chat image_url object.
|
||||
assert!(parts[1]["image_url"].get("detail").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn text_only_parts_collapse_to_string() {
|
||||
let req = ResponsesRequest {
|
||||
|
||||
303
script/deploy.sh
303
script/deploy.sh
@@ -1,303 +0,0 @@
|
||||
#!/bin/env bash
|
||||
#
|
||||
# Rolling deploy across the helexa fleet, driven by asset/manifest.yml.
|
||||
# Installs / upgrades cortex on the gateway host and the appropriate
|
||||
# helexa-neuron-<flavour> package on each neuron host, then restarts
|
||||
# their services.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
REPO_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)"
|
||||
MANIFEST="${REPO_DIR}/asset/manifest.yml"
|
||||
|
||||
if [[ ! -f "${MANIFEST}" ]]; then
|
||||
echo "fatal: manifest not found at ${MANIFEST}" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Parse the manifest with yq. NOTE: this expects the pip-installed yq
|
||||
# (a jq wrapper using jq syntax) — `pip install yq`. The Fedora rpm
|
||||
# `yq` is mikefarah/yq and uses different (yaml-native) syntax; if a
|
||||
# host has that one instead these queries will fail.
|
||||
cortex_host=$(yq -r '.cortex.host' "${MANIFEST}")
|
||||
|
||||
# Emit one TAB-separated 'host\tflavour' line per neuron.
|
||||
mapfile -t neuron_entries < <(
|
||||
yq -r '.neurons[] | .host + "\t" + .flavour' "${MANIFEST}"
|
||||
)
|
||||
|
||||
# Return the installed package's "version-release" string, or
|
||||
# "(not installed)" when rpm reports the package as absent. Capture
|
||||
# rpm's output into a variable so its "package X is not installed"
|
||||
# stdout message (rpm writes that to stdout, not stderr, when -q fails)
|
||||
# doesn't leak into the result.
|
||||
installed_nvr() {
|
||||
local host="$1" pkg="$2"
|
||||
local nvr
|
||||
if nvr=$(ssh "${host}" "rpm -q --qf '%{version}-%{release}' ${pkg} 2>/dev/null"); then
|
||||
echo "${nvr}"
|
||||
else
|
||||
echo "(not installed)"
|
||||
fi
|
||||
}
|
||||
|
||||
# Ensure the rpm.lair.cafe unstable repo is configured AND enabled on
|
||||
# the remote host.
|
||||
#
|
||||
# The upstream .repo file at https://rpm.lair.cafe/lair-cafe-unstable.repo
|
||||
# ships with `enabled=0` so a host that just fetched it won't start
|
||||
# pulling unstable packages by accident. We have to explicitly flip
|
||||
# enabled=1 via `dnf config-manager setopt`. Both addrepo and setopt
|
||||
# are idempotent.
|
||||
#
|
||||
# Non-fatal — if either step fails the subsequent `dnf install` will
|
||||
# surface a clearer diagnostic on its own.
|
||||
ensure_lair_repo() {
|
||||
local host="$1"
|
||||
if ! ssh "${host}" "test -f /etc/yum.repos.d/lair-cafe-unstable.repo" 2>/dev/null; then
|
||||
echo "[${host}] adding rpm.lair.cafe unstable repo"
|
||||
if ! ssh "${host}" sudo dnf config-manager addrepo \
|
||||
--from-repofile=https://rpm.lair.cafe/lair-cafe-unstable.repo \
|
||||
>/dev/null 2>&1; then
|
||||
echo "[${host}] WARNING: failed to add lair.cafe repo file (proceeding anyway)"
|
||||
return 0
|
||||
fi
|
||||
fi
|
||||
# The .repo file ships enabled=0; flip it on. Cheap, idempotent.
|
||||
if ! ssh "${host}" sudo dnf config-manager setopt \
|
||||
lair-cafe-unstable.enabled=1 >/dev/null 2>&1; then
|
||||
echo "[${host}] WARNING: failed to enable lair-cafe-unstable (proceeding anyway)"
|
||||
fi
|
||||
}
|
||||
|
||||
# Ensure libcudnn.so.9 is resolvable on the remote host so the
|
||||
# neuron binary (built with --features cudnn) doesn't fail at startup
|
||||
# with "cannot open shared object file: No such file or directory".
|
||||
#
|
||||
# Probes ldconfig first — if cuDNN was installed manually (.tar/.run
|
||||
# install), it'll be cached by ldconfig and we don't touch it.
|
||||
# Otherwise adds NVIDIA's RHEL9 CUDA repo (the Fedora 43 CUDA repo
|
||||
# doesn't ship cuDNN packages — only the RHEL9 one does) and installs
|
||||
# libcudnn9-cuda-13.
|
||||
ensure_cudnn_runtime() {
|
||||
local host="$1"
|
||||
if ssh "${host}" "ldconfig -p | grep -q libcudnn.so.9" 2>/dev/null; then
|
||||
return 0
|
||||
fi
|
||||
echo "[${host}] installing cuDNN runtime"
|
||||
if ! ssh "${host}" "test -f /etc/yum.repos.d/cuda-rhel9-x86_64.repo" 2>/dev/null; then
|
||||
if ! ssh "${host}" sudo dnf config-manager addrepo \
|
||||
--from-repofile=https://developer.download.nvidia.com/compute/cuda/repos/rhel9/x86_64/cuda-rhel9.repo \
|
||||
>/dev/null 2>&1; then
|
||||
echo "[${host}] WARNING: failed to add rhel9 CUDA repo (proceeding anyway)"
|
||||
fi
|
||||
fi
|
||||
if ! ssh "${host}" sudo dnf install -y libcudnn9-cuda-13 >/dev/null 2>&1; then
|
||||
echo "[${host}] WARNING: failed to install libcudnn9-cuda-13"
|
||||
echo "[${host}] neuron may fail to start; install cuDNN manually if so"
|
||||
fi
|
||||
}
|
||||
|
||||
# True when the named package needs to be installed or upgraded on the
|
||||
# remote host — either it's not present, or a newer version exists in
|
||||
# the repo. False only when the installed version is current.
|
||||
#
|
||||
# `dnf check-update <pkg>` returns 0 when the package isn't installed
|
||||
# at all (there's nothing to update), so we have to probe with rpm -q
|
||||
# first to distinguish "absent" from "current". Other dnf failures
|
||||
# collapse into "needs update" so the subsequent install step surfaces
|
||||
# the real diagnostic rather than this check swallowing it.
|
||||
needs_update() {
|
||||
local host="$1" pkg="$2"
|
||||
# Not installed → needs work.
|
||||
if ! ssh "${host}" "rpm -q ${pkg}" >/dev/null 2>&1; then
|
||||
return 0
|
||||
fi
|
||||
# Installed; ask dnf whether the repo has something newer.
|
||||
if ssh "${host}" sudo dnf check-update --refresh -q "${pkg}" >/dev/null 2>&1; then
|
||||
return 1
|
||||
else
|
||||
return 0
|
||||
fi
|
||||
}
|
||||
|
||||
# True if the named package is currently installed on the remote host.
|
||||
# Used to decide between `dnf install` (fresh) and `dnf upgrade` (stale):
|
||||
# dnf5's `install` is a no-op when the package is already present at
|
||||
# any version — it does NOT auto-upgrade to the latest available — so
|
||||
# the wrong command silently leaves the host on an old build.
|
||||
is_installed() {
|
||||
local host="$1" pkg="$2"
|
||||
ssh "${host}" "rpm -q ${pkg}" >/dev/null 2>&1
|
||||
}
|
||||
|
||||
# Install or upgrade the named package on the remote, picking the
|
||||
# right dnf verb based on the installed-or-not state. Returns 0 with
|
||||
# dnf's combined stdout/stderr captured in __DNF_OUTPUT__ on success,
|
||||
# and 1 with the same captured output on failure.
|
||||
__DNF_OUTPUT__=""
|
||||
install_or_upgrade() {
|
||||
local host="$1" pkg="$2"
|
||||
local cmd
|
||||
if is_installed "${host}" "${pkg}"; then
|
||||
cmd="upgrade"
|
||||
else
|
||||
cmd="install"
|
||||
fi
|
||||
if __DNF_OUTPUT__=$(
|
||||
ssh "${host}" sudo dnf "${cmd}" --refresh --allowerasing -y "${pkg}" 2>&1
|
||||
); then
|
||||
return 0
|
||||
else
|
||||
return 1
|
||||
fi
|
||||
}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# cortex (gateway)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
ensure_lair_repo "${cortex_host}"
|
||||
cortex_nvr=$(installed_nvr "${cortex_host}" cortex)
|
||||
if needs_update "${cortex_host}" cortex; then
|
||||
echo "[${cortex_host}] cortex update available (current: ${cortex_nvr})"
|
||||
# Stop the service only if the unit file exists — fresh installs
|
||||
# don't have it, and `systemctl stop` on a missing unit returns
|
||||
# non-zero, which would otherwise short-circuit the install branch
|
||||
# under set -e.
|
||||
if ssh "${cortex_host}" "[ ! -f /usr/lib/systemd/system/cortex.service ] || sudo systemctl stop cortex.service"; then
|
||||
echo "[${cortex_host}] stopped cortex service"
|
||||
if install_or_upgrade "${cortex_host}" cortex; then
|
||||
cortex_nvr=$(installed_nvr "${cortex_host}" cortex)
|
||||
echo "[${cortex_host}] installed/upgraded cortex to ${cortex_nvr}"
|
||||
else
|
||||
echo "[${cortex_host}] failed to install/upgrade cortex:"
|
||||
echo "${__DNF_OUTPUT__}" | sed "s/^/[${cortex_host}] /"
|
||||
fi
|
||||
else
|
||||
echo "[${cortex_host}] failed to stop cortex service"
|
||||
fi
|
||||
else
|
||||
echo "[${cortex_host}] cortex is up to date (${cortex_nvr})"
|
||||
ssh "${cortex_host}" sudo systemctl stop cortex.service || true
|
||||
fi
|
||||
|
||||
# Sync cortex.toml whether the package was upgraded or not — the config
|
||||
# can change without a package bump.
|
||||
if rsync \
|
||||
--archive \
|
||||
--compress \
|
||||
--rsync-path 'sudo rsync' \
|
||||
--chown root:root \
|
||||
--chmod 644 \
|
||||
"${REPO_DIR}/cortex.toml" \
|
||||
"${cortex_host}:/etc/cortex/cortex.toml"; then
|
||||
echo "[${cortex_host}] sync'd cortex.toml"
|
||||
else
|
||||
echo "[${cortex_host}] failed to sync cortex.toml"
|
||||
fi
|
||||
|
||||
# Sync models.toml on the same lifecycle as cortex.toml — operator-owned,
|
||||
# gitignored, drives /v1/models catalogue × topology resolution.
|
||||
if [[ -f "${REPO_DIR}/models.toml" ]]; then
|
||||
if rsync \
|
||||
--archive \
|
||||
--compress \
|
||||
--rsync-path 'sudo rsync' \
|
||||
--chown root:root \
|
||||
--chmod 644 \
|
||||
"${REPO_DIR}/models.toml" \
|
||||
"${cortex_host}:/etc/cortex/models.toml"; then
|
||||
echo "[${cortex_host}] sync'd models.toml"
|
||||
else
|
||||
echo "[${cortex_host}] failed to sync models.toml"
|
||||
fi
|
||||
else
|
||||
echo "[${cortex_host}] no local models.toml — leaving /etc/cortex/models.toml untouched"
|
||||
fi
|
||||
|
||||
ssh "${cortex_host}" sudo systemctl daemon-reload
|
||||
if ssh "${cortex_host}" systemctl is-active --quiet cortex.service; then
|
||||
echo "[${cortex_host}] cortex service is active"
|
||||
elif ssh "${cortex_host}" sudo systemctl start cortex.service; then
|
||||
echo "[${cortex_host}] started cortex service"
|
||||
else
|
||||
echo "[${cortex_host}] failed to start cortex service"
|
||||
fi
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# neuron (per-host, flavour from manifest)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
for entry in "${neuron_entries[@]}"; do
|
||||
IFS=$'\t' read -r neuron_host neuron_flavour <<< "${entry}"
|
||||
package="helexa-neuron-${neuron_flavour}"
|
||||
# First dot-component of the host keys the per-host config file
|
||||
# under asset/neuron/<short>.toml. A host listed in the manifest
|
||||
# without a corresponding config still deploys (the package's
|
||||
# default /etc/neuron/neuron.toml stays in place; no pre-warm).
|
||||
short_host="${neuron_host%%.*}"
|
||||
host_config="${REPO_DIR}/asset/neuron/${short_host}.toml"
|
||||
|
||||
ensure_lair_repo "${neuron_host}"
|
||||
ensure_cudnn_runtime "${neuron_host}"
|
||||
neuron_nvr=$(installed_nvr "${neuron_host}" "${package}")
|
||||
|
||||
# Stop the service unconditionally before any reconfig step.
|
||||
# `default_models` is read at activation, so a config change without
|
||||
# a bounce silently leaves the host on the previous pre-warm set.
|
||||
# Same shape as the cortex flow above. The `[ ! -f … ]` guard skips
|
||||
# the stop on a fresh install where the unit file isn't there yet.
|
||||
if ssh "${neuron_host}" "[ ! -f /usr/lib/systemd/system/neuron.service ] || sudo systemctl stop neuron.service"; then
|
||||
echo "[${neuron_host}] stopped neuron service"
|
||||
else
|
||||
echo "[${neuron_host}] failed to stop neuron service (continuing)"
|
||||
fi
|
||||
|
||||
if needs_update "${neuron_host}" "${package}"; then
|
||||
echo "[${neuron_host}] ${package} update available (current: ${neuron_nvr})"
|
||||
# --allowerasing lets dnf swap out a previously-installed
|
||||
# bare helexa-neuron or a different flavour without manual
|
||||
# intervention. The Conflicts: clauses in the spec ensure
|
||||
# only one flavour is ever resident.
|
||||
if install_or_upgrade "${neuron_host}" "${package}"; then
|
||||
neuron_nvr=$(installed_nvr "${neuron_host}" "${package}")
|
||||
echo "[${neuron_host}] installed/upgraded ${package} to ${neuron_nvr}"
|
||||
# Ensure firewalld allows neuron port
|
||||
ssh "${neuron_host}" "sudo firewall-cmd --query-service=helexa-neuron --quiet 2>/dev/null || sudo firewall-cmd --add-service=helexa-neuron --permanent && sudo firewall-cmd --reload" 2>/dev/null || true
|
||||
else
|
||||
echo "[${neuron_host}] failed to install ${package}:"
|
||||
echo "${__DNF_OUTPUT__}" | sed "s/^/[${neuron_host}] /"
|
||||
fi
|
||||
else
|
||||
echo "[${neuron_host}] ${package} is up to date (${neuron_nvr})"
|
||||
fi
|
||||
|
||||
# Sync per-host neuron.toml — drives default_models pre-warm so
|
||||
# `/v1/models` on the gateway exposes the host's headline model
|
||||
# immediately after the service comes back up. Missing per-host
|
||||
# config leaves the package's installed neuron.toml untouched.
|
||||
if [[ -f "${host_config}" ]]; then
|
||||
if rsync \
|
||||
--archive \
|
||||
--compress \
|
||||
--rsync-path 'sudo rsync' \
|
||||
--chown root:root \
|
||||
--chmod 644 \
|
||||
"${host_config}" \
|
||||
"${neuron_host}:/etc/neuron/neuron.toml"; then
|
||||
echo "[${neuron_host}] sync'd asset/neuron/${short_host}.toml"
|
||||
else
|
||||
echo "[${neuron_host}] failed to sync neuron.toml"
|
||||
fi
|
||||
else
|
||||
echo "[${neuron_host}] no asset/neuron/${short_host}.toml — leaving /etc/neuron/neuron.toml untouched"
|
||||
fi
|
||||
|
||||
if ssh "${neuron_host}" "sudo systemctl daemon-reload && sudo systemctl start neuron.service"; then
|
||||
echo "[${neuron_host}] started neuron service"
|
||||
else
|
||||
echo "[${neuron_host}] failed to start neuron service"
|
||||
fi
|
||||
done
|
||||
@@ -53,6 +53,11 @@ for host in "${cortex_host}" "${neuron_hosts[@]}"; do
|
||||
# own fresh lookup.
|
||||
sudo install -d -o gitea_ci -g gitea_ci -m 0700 \
|
||||
/var/lib/gitea_ci/.ssh
|
||||
# Grant journal read access so the deploy workflow can capture
|
||||
# `journalctl -u <unit> -I` after a service start without
|
||||
# needing a sudoers entry. Idempotent — usermod -aG on an
|
||||
# already-member is a no-op.
|
||||
sudo usermod -aG systemd-journal gitea_ci
|
||||
'; then
|
||||
echo " failed to provision gitea_ci — skipping ${host}"
|
||||
continue
|
||||
|
||||
Reference in New Issue
Block a user