use anyhow::{bail, ensure, Context, Result};
use mlx_native::GgmlWorkloadClass;
use std::collections::BTreeMap;
use std::path::Path;
use std::sync::Arc;
use crate::convert::cli_driver::{RetainedQwenArtifactLoad, VerifiedStoredQwenArtifact};
use crate::serve::header::LoadProgress;
use crate::serve::multi_seq_kv::SlotId;
use super::execution_config::Qwen35ExecutionConfiguration;
use super::execution_dispatch::{
with_execution_configuration, with_execution_trace_capture, Qwen35EncodedDispatchObservation,
Qwen35TraceWeightSlot,
};
use super::execution_observation::{
capture_loaded_tensor_catalog, VerifiedExecutedTensorCatalog, VerifiedLoadedTensorCatalog,
};
use super::execution_trace::{verify_encoded_dispatch_catalog, VerifiedEncodedDispatchCatalog};
use super::kv_cache::HybridKvCache;
use super::model::Qwen35Model;
use super::source_config::qwen35_config_from_authenticated_source;
use super::Qwen35Config;
const MAX_SESSION_ENCODED_DISPATCHES: usize = 8_192;
pub(crate) struct NoDwqQwen35LoadedCandidate {
loaded: RetainedQwenArtifactLoad<LoadedQwenCandidateState>,
execution: Qwen35ExecutionConfiguration,
}
struct LoadedQwenCandidateState {
model: Qwen35Model,
catalog: Arc<VerifiedLoadedTensorCatalog>,
}
pub(crate) struct NoDwqQwen35TextSession<'a> {
candidate: &'a NoDwqQwen35LoadedCandidate,
kv_cache: HybridKvCache,
executed_catalog: Arc<VerifiedExecutedTensorCatalog>,
trace_weight_slots: BTreeMap<usize, Qwen35TraceWeightSlot>,
dispatches: Vec<Qwen35EncodedDispatchObservation>,
captured_prompt: bool,
captured_decode: bool,
poisoned: bool,
}
pub(super) struct ReconciledLoadedCandidateIdentity {
configuration: Qwen35ExecutionConfiguration,
conversion_receipt_sha256: String,
loaded_catalog: Arc<VerifiedLoadedTensorCatalog>,
}
impl ReconciledLoadedCandidateIdentity {
fn new(
configuration: Qwen35ExecutionConfiguration,
conversion_receipt_sha256: String,
loaded_catalog: Arc<VerifiedLoadedTensorCatalog>,
) -> Self {
Self {
configuration,
conversion_receipt_sha256,
loaded_catalog,
}
}
pub(super) fn into_parts(
self,
) -> (
Qwen35ExecutionConfiguration,
String,
Arc<VerifiedLoadedTensorCatalog>,
) {
(
self.configuration,
self.conversion_receipt_sha256,
self.loaded_catalog,
)
}
}
impl NoDwqQwen35LoadedCandidate {
pub(crate) fn config(&self) -> &Qwen35Config {
&self.loaded.value().model.cfg
}
pub(crate) fn execution(&self) -> &Qwen35ExecutionConfiguration {
&self.execution
}
pub(crate) fn conversion_receipt_sha256(&self) -> &str {
&self.loaded.conversion().receipt().receipt_sha256
}
pub(crate) fn loaded_catalog_sha256(&self) -> &str {
self.loaded.value().catalog.catalog_sha256()
}
pub(crate) fn loaded_catalog_conversion_receipt_sha256(&self) -> &str {
self.loaded.value().catalog.conversion_receipt_sha256()
}
pub(crate) fn loaded_tensor_count(&self) -> usize {
self.loaded.value().catalog.observations().len()
}
pub(crate) fn start_text_session(
&self,
max_seq_len: u32,
) -> Result<NoDwqQwen35TextSession<'_>> {
let model = &self.loaded.value().model;
let kv_cache = with_execution_configuration(&self.execution, || {
model.ensure_gpu_cache_primed()?;
model.with_gpu_cache_mut(|device, _registry| {
HybridKvCache::new(&model.cfg, device, max_seq_len, 1)
})
})?;
let executed_catalog = model
.loaded_candidate_executed_catalog()?
.context("evidence-bearing Qwen cache omitted executed tensor observations")?;
let trace_weight_slots = model
.loaded_candidate_trace_slots()?
.context("evidence-bearing Qwen cache omitted trace weight slots")?;
Ok(NoDwqQwen35TextSession {
candidate: self,
kv_cache,
executed_catalog,
trace_weight_slots,
dispatches: Vec::new(),
captured_prompt: false,
captured_decode: false,
poisoned: false,
})
}
}
impl NoDwqQwen35TextSession<'_> {
fn preflight_capture(&self) -> Result<()> {
let maximum = self
.trace_weight_slots
.len()
.checked_mul(2)
.context("encoded dispatch observation bound overflow")?;
if maximum > MAX_SESSION_ENCODED_DISPATCHES {
bail!(
"encoded dispatch topology exceeds bounded session limit {}",
MAX_SESSION_ENCODED_DISPATCHES
);
}
Ok(())
}
pub(crate) fn executed_catalog_sha256(&self) -> &str {
self.executed_catalog.catalog_sha256()
}
pub(crate) fn executed_tensor_count(&self) -> usize {
self.executed_catalog.observations().len()
}
pub(crate) fn executed_catalog_loaded_parent_sha256(&self) -> &str {
self.executed_catalog.loaded_catalog_sha256()
}
pub(crate) fn encoded_dispatches(&self) -> &[Qwen35EncodedDispatchObservation] {
&self.dispatches
}
pub(crate) fn seal_encoded_dispatches(self) -> Result<VerifiedEncodedDispatchCatalog> {
verify_encoded_dispatch_catalog(
&self.candidate.execution,
&self.executed_catalog,
self.dispatches,
)
}
#[cfg(test)]
pub(crate) fn duplicate_observation_fails_sealing(&self) -> bool {
let mut observations = self.dispatches.clone();
let Some(first) = observations.first().cloned() else {
return false;
};
observations.push(first);
verify_encoded_dispatch_catalog(
&self.candidate.execution,
&self.executed_catalog,
observations,
)
.is_err()
}
pub(crate) fn forward(&mut self, tokens: &[u32], positions_flat: &[i32]) -> Result<Vec<f32>> {
if self.poisoned {
bail!("bounded evidence session is poisoned by an earlier failed execution");
}
if tokens.len() <= 8 {
bail!(
"bounded prompt evidence requires more than eight tokens; width-N and decode are separate regimes"
);
}
if positions_flat.len() != tokens.len().saturating_mul(4) {
bail!("bounded prompt evidence requires four MROPE positions per token");
}
if self.captured_decode {
bail!("bounded prompt evidence cannot be captured after decode");
}
if self.captured_prompt {
bail!("the bounded evidence session already captured its prompt execution");
}
self.preflight_capture()?;
self.poisoned = true;
let model = &self.candidate.loaded.value().model;
let capture = with_execution_trace_capture(
&self.candidate.execution,
self.trace_weight_slots.clone(),
GgmlWorkloadClass::Prompt,
|| model.forward_gpu(tokens, positions_flat, &mut self.kv_cache, SlotId(0)),
);
let (logits, dispatches) = capture?;
self.captured_prompt = true;
self.poisoned = false;
self.dispatches.extend(dispatches);
Ok(logits)
}
pub(crate) fn forward_greedy(&mut self, token: u32, positions_flat: [i32; 4]) -> Result<u32> {
if self.poisoned {
bail!("bounded evidence session is poisoned by an earlier failed execution");
}
if !self.captured_prompt {
bail!("bounded decode evidence requires a successfully captured prompt first");
}
if self.captured_decode {
bail!("the bounded evidence session already captured its decode execution");
}
self.preflight_capture()?;
self.poisoned = true;
let model = &self.candidate.loaded.value().model;
let capture = with_execution_trace_capture(
&self.candidate.execution,
self.trace_weight_slots.clone(),
GgmlWorkloadClass::DecodeSingle,
|| model.forward_gpu_greedy(&[token], &positions_flat, &mut self.kv_cache, SlotId(0)),
);
let (token, dispatches) = capture?;
self.captured_decode = true;
self.poisoned = false;
self.dispatches.extend(dispatches);
Ok(token)
}
}
pub(crate) fn load_no_dwq_qwen35_candidate(
artifact: VerifiedStoredQwenArtifact,
dwq_overlay_path: Option<&Path>,
progress: &mut LoadProgress,
) -> Result<NoDwqQwen35LoadedCandidate> {
let loaded = artifact
.load_and_reconcile(|gguf, source_config, conversion| {
let actual = Qwen35Model::load_config_only(gguf)?;
let expected = qwen35_config_from_authenticated_source(source_config)?;
if actual != expected {
bail!("GGUF execution configuration differs from authenticated source config");
}
let execution = Qwen35ExecutionConfiguration::resolve_no_dwq_from_environment(
&actual,
dwq_overlay_path,
)?;
execution.validate()?;
let ((model, execution), catalog) = capture_loaded_tensor_catalog(conversion, || {
let model = Qwen35Model::load_from_gguf(gguf, progress)
.context("load copied no-DWQ Qwen weights from retained GGUF")?;
if model.cfg != expected {
bail!("loaded Qwen configuration drifted from authenticated source geometry");
}
validate_loaded_global_geometry(&model, &expected)?;
Ok((model, execution))
})?;
Ok((model, execution, catalog))
})
.map_err(|error| anyhow::anyhow!(error.to_string()))?;
let execution = loaded.value().1.clone();
let loaded = loaded.try_map(|(mut model, _, catalog), conversion| {
let catalog = Arc::new(catalog);
model.bind_loaded_candidate_identity(ReconciledLoadedCandidateIdentity::new(
execution.clone(),
conversion.receipt().receipt_sha256.clone(),
Arc::clone(&catalog),
))?;
Ok::<_, anyhow::Error>(LoadedQwenCandidateState { model, catalog })
})?;
Ok(NoDwqQwen35LoadedCandidate { loaded, execution })
}
fn validate_loaded_global_geometry(model: &Qwen35Model, expected: &Qwen35Config) -> Result<()> {
ensure!(
!model.tied_word_embeddings,
"authenticated Qwen source-teacher profile requires an untied output head"
);
let expected_vocab = usize::try_from(expected.vocab_size)?;
let expected_hidden = usize::try_from(expected.hidden_size)?;
let expected_embedding_values = usize::try_from(expected.vocab_size)?
.checked_mul(expected_hidden)
.context("authenticated Qwen embedding geometry overflows usize")?;
let embedding_matches = match &model.token_embd_native {
Some(weight) => {
model.token_embd.is_empty()
&& weight.info.rows == expected_vocab
&& weight.info.cols == expected_hidden
}
None => model.token_embd.len() == expected_embedding_values,
};
let output_matches = match &model.output_weight_native {
Some(weight) => {
model.output_weight.is_empty()
&& weight.info.rows == expected_vocab
&& weight.info.cols == expected_hidden
}
None => model.output_weight.len() == expected_embedding_values,
};
if !embedding_matches || !output_matches || model.output_norm.len() != expected_hidden {
bail!(
"loaded Qwen globals differ from authenticated source geometry; alternate storage and synthetic vocab extension are not admitted"
);
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn loaded_candidate_identity_is_atomic_and_synthetic_vocab_is_rejected() {
let config = json!({
"hidden_size": 32,
"num_hidden_layers": 1,
"num_attention_heads": 1,
"num_key_value_heads": 1,
"head_dim": 32,
"linear_num_key_heads": 1,
"linear_num_value_heads": 1,
"linear_key_head_dim": 32,
"linear_value_head_dim": 32,
"linear_conv_kernel_dim": 4,
"intermediate_size": 32,
"max_position_embeddings": 128,
"vocab_size": 4
});
let projected = qwen35_config_from_authenticated_source(&config).unwrap();
let resolved = Qwen35ExecutionConfiguration::from_resolved(
mlx_native::GgmlRoutingPolicy::default(),
super::super::execution_config::Qwen35GateUpPolicy::Separate,
)
.unwrap();
let mut model = Qwen35Model::empty_from_cfg(projected.clone());
let catalog = Arc::new(
super::super::execution_observation::VerifiedLoadedTensorCatalog::for_test_empty(
"a".repeat(64),
),
);
model
.bind_loaded_candidate_identity(ReconciledLoadedCandidateIdentity::new(
resolved.clone(),
"a".repeat(64),
Arc::clone(&catalog),
))
.unwrap();
assert_eq!(
model.loaded_candidate_cache_identity(),
Some((
"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
resolved.graph_configuration_sha256(),
resolved.routing_policy_sha256(),
))
);
assert!(model
.bind_loaded_candidate_identity(ReconciledLoadedCandidateIdentity::new(
resolved,
"b".repeat(64),
catalog,
))
.is_err());
let mut synthesized = Qwen35Model::empty_from_cfg(projected.clone());
synthesized.token_embd.extend([0.0; 32]);
assert!(validate_loaded_global_geometry(&synthesized, &projected).is_err());
}
}