use std::path::PathBuf;
use std::time::{Duration, Instant};
use anyhow::{Context, Result};
use tokenizers::Tokenizer;
use crate::inference::models::qwen35::kv_cache::{HybridKvCache, HybridKvCacheSnapshot};
use crate::inference::models::qwen35::model::Qwen35Model;
use crate::serve::load_info::{
self, ArchFamily, ChatTemplateSource, LoadInfo, LoadInfoBuilder, MoeShape, TokenizerSource,
};
use crate::core::provenance::{self, Provenance};
use crate::serve::multi_seq_kv::SlotId;
use super::engine::{effective_repetition_penalty, LoadOptions, SamplingParams};
pub const QWEN35_EOS_STOP_NAMES: &[&str] = &["<|im_end|>", "<|endoftext|>"];
pub struct Qwen35LoadedModel {
pub model: Qwen35Model,
pub tokenizer: Tokenizer,
pub chat_template: String,
pub model_id: String,
pub model_path: PathBuf,
pub eos_token_ids: Vec<u32>,
pub hidden_size: usize,
pub vocab_size: usize,
pub context_length: Option<usize>,
pub quant_type: Option<String>,
pub load_duration: Duration,
pub provenance: Provenance,
pub prompt_cache: HybridPromptCache,
pub lcp_registry: crate::serve::kv_persist::lcp_registry::LcpRegistry<
crate::inference::models::qwen35::kv_cache::HybridKvCacheSnapshot,
>,
pub kv_metrics_sink:
Option<std::sync::Arc<dyn crate::serve::kv_persist::metrics::KvCacheMetricsSink>>,
pub disk_persistor: Option<
std::sync::Arc<
crate::serve::kv_persist::families::qwen35_disk_persistor::Qwen35DiskPersistor,
>,
>,
pub lcp_hydrated_for_cfg: std::collections::HashSet<String>,
pub tq_kv_active: bool,
pub persistent_kv_cache: Option<HybridKvCache>,
}
impl Qwen35LoadedModel {
pub fn load(opts: &LoadOptions) -> Result<Self> {
let load_start = Instant::now();
let model_path = &opts.model_path;
anyhow::ensure!(
model_path.exists(),
"Model not found: {}",
model_path.display()
);
let gguf = mlx_native::gguf::GgufFile::open(model_path)
.map_err(|e| anyhow::anyhow!("GGUF open: {e}"))?;
let provenance = provenance::detect(&gguf);
let tokenizer_path =
crate::serve::find_tokenizer(model_path, opts.tokenizer_path.as_deref())?;
let stderr_is_tty = std::io::IsTerminal::is_terminal(&std::io::stderr());
let verbosity = if tracing::enabled!(tracing::Level::INFO) {
1
} else {
0
};
let cfg_preview = Qwen35Model::load_config_only(&gguf)
.context("Qwen35Model::load_config_only (progress sizing)")?;
let mut progress = crate::serve::header::LoadProgress::new(
stderr_is_tty,
verbosity,
cfg_preview.num_hidden_layers as usize,
);
let mut model = Qwen35Model::load_from_gguf(&gguf, &mut progress)
.context("Qwen35Model::load_from_gguf")?;
if let Some(overlay_path) = opts.dwq_overlay_path.as_ref() {
let device = mlx_native::MlxDevice::new()
.map_err(|e| anyhow::anyhow!("Qwen35 DWQ overlay device: {e}"))?;
let stacked = model
.apply_dwq_overlay(&device, overlay_path)
.with_context(|| {
format!("Qwen35 DWQ overlay from {} failed", overlay_path.display())
})?;
tracing::info!(
count = stacked,
path = %overlay_path.display(),
"DWQ overlay applied to Qwen35LoadedModel"
);
}
let mut eos_token_ids: Vec<u32> = Vec::with_capacity(2);
if let Some(id) = gguf.metadata_u32("tokenizer.ggml.eos_token_id") {
eos_token_ids.push(id);
}
if let Some(id) = gguf.metadata_u32("tokenizer.ggml.eot_token_id") {
if !eos_token_ids.contains(&id) {
eos_token_ids.push(id);
}
}
if let Some(arr) = gguf.metadata("tokenizer.ggml.tokens") {
if let mlx_native::gguf::MetadataValue::Array(elems) = arr {
for (i, el) in elems.iter().enumerate() {
if let mlx_native::gguf::MetadataValue::String(s) = el {
if QWEN35_EOS_STOP_NAMES.contains(&s.as_str()) {
let id = i as u32;
if !eos_token_ids.contains(&id) {
eos_token_ids.push(id);
}
}
}
}
}
}
if eos_token_ids.is_empty() {
eos_token_ids.push(151_645);
}
let _eos_token: u32 = eos_token_ids[0];
tracing::info!(
count = eos_token_ids.len(),
ids = ?eos_token_ids,
"Qwen35 EOS token set resolved (iter-267 multi-source)"
);
let _tokenizer_path = tokenizer_path;
let mut tokenizer =
crate::inference::models::qwen35::tokenizer::build_tokenizer_from_gguf(&gguf)
.map_err(|e| anyhow::anyhow!("GGUF-driven tokenizer build failed: {e}"))?;
tokenizer
.with_truncation(None)
.map_err(|e| anyhow::anyhow!("Failed to disable tokenizer truncation: {e}"))?;
let chat_template = gguf
.metadata_string("tokenizer.chat_template")
.map(|s| s.to_string())
.unwrap_or_default();
let model_id = gguf
.metadata_string("general.name")
.map(|s| s.to_string())
.unwrap_or_else(|| {
model_path
.file_stem()
.map(|s| s.to_string_lossy().into_owned())
.unwrap_or_else(|| "qwen35-model".to_string())
});
let hidden_size = model.cfg.hidden_size as usize;
let vocab_size = model.cfg.vocab_size as usize;
let context_length = if model.cfg.max_position_embeddings > 0 {
Some(model.cfg.max_position_embeddings as usize)
} else {
None
};
let quant_type = crate::serve::load_info::infer_quant_label(&gguf);
let load_duration = load_start.elapsed();
let loaded = Self {
model,
tokenizer,
chat_template,
model_id,
model_path: model_path.clone(),
eos_token_ids,
hidden_size,
vocab_size,
context_length,
quant_type,
load_duration,
provenance,
prompt_cache: HybridPromptCache::new(),
lcp_registry:
crate::serve::kv_persist::lcp_registry::LcpRegistry::with_byte_budget(
qwen35_lcp_registry_byte_budget(),
),
kv_metrics_sink: None,
disk_persistor: opts.kv_persist_dir.as_ref().and_then(|cache_dir| {
let budget_bytes: u64 = match std::env::var("HF2Q_KV_PERSIST_BUDGET_BYTES") {
Ok(raw) => match raw.trim().parse::<u64>() {
Ok(parsed) => parsed,
Err(err) => {
tracing::warn!(
raw = %raw,
error = %err,
"ADR-027 23d-γ: HF2Q_KV_PERSIST_BUDGET_BYTES parse failed; \
defaulting to 0 (unlimited)"
);
0
}
},
Err(_) => 0,
};
match crate::serve::kv_persist::families::qwen35_disk_persistor::Qwen35DiskPersistor::new_with_budget(cache_dir.clone(), budget_bytes) {
Ok(p) => {
tracing::info!(
cache_dir = %cache_dir.display(),
budget_bytes = budget_bytes,
"ADR-027 iter-6b.2 + 23d-γ: Qwen35DiskPersistor constructed; \
cold-process LCP resume enabled"
);
Some(std::sync::Arc::new(p))
}
Err(e) => {
tracing::warn!(
cache_dir = %cache_dir.display(),
error = %e,
"ADR-027 iter-6b.2: Qwen35DiskPersistor construction failed; \
falling back to in-process-only LCP"
);
None
}
}
}),
lcp_hydrated_for_cfg: std::collections::HashSet::new(),
tq_kv_active: crate::serve::api::tq_packed_descriptor::is_tq_active_mode(),
persistent_kv_cache: None,
};
if loaded.tq_kv_active && opts.kv_persist_dir.is_some() {
tracing::info!(
"ADR-027 sub-iter 23d-γ: HF2Q_TQ_KV=1 + HF2Q_KV_PERSIST both active — \
persist snapshots round-trip the TQ substrate (codec v4); fingerprint \
is substrate-namespaced so cross-mode hydration is a clean miss."
);
}
if loaded.tq_kv_active {
tracing::info!(
"ADR-027 sub-iter 23d-γ: HF2Q_TQ_KV=1; LCP resume restores TQ buffers \
per-slot (restore_partial 23d-γ) — coherent under the production \
TQ-only regime."
);
}
Ok(loaded)
}
pub fn store_lcp_with_disk_writeback(
&mut self,
kv_cache: &crate::inference::models::qwen35::kv_cache::HybridKvCache,
key: crate::serve::kv_persist::lcp_registry::LcpKey,
prompt_tokens: Vec<u32>,
snapshot: crate::inference::models::qwen35::kv_cache::HybridKvCacheSnapshot,
sliding_window: usize,
linear_capacity: usize,
) -> Result<(), crate::serve::kv_persist::lcp_registry::LcpStoreError> {
if let Some(persistor) = &self.disk_persistor {
match crate::serve::kv_persist::families::qwen35_hybrid_persistor::cfg_from_cache(
kv_cache,
crate::serve::kv_persist::families::qwen35_hybrid_persistor::FullAttnCodec::F32Dense,
) {
Ok(cfg) => {
let key_hex = lcp_key_to_filename_hex(&key);
let sidecar = crate::serve::kv_persist::families::qwen35_hybrid_persistor::LcpSidecarMetadata {
model_fingerprint: key.model_fingerprint.clone(),
tenant_id: key.tenant_id.clone(),
params_hash: key.params_hash,
prompt_tokens: prompt_tokens.clone(),
sliding_window: sliding_window as u64,
linear_capacity: linear_capacity as u64,
};
if let Err(e) = persistor.write(&cfg, &key_hex, &snapshot, &sidecar) {
tracing::warn!(
cache_dir = %persistor.cache_dir().display(),
key_hex = %key_hex,
error = %format!("{e:#}"),
"ADR-027 iter-6b.3: disk persistor write failed; \
in-memory store will still proceed"
);
}
}
Err(e) => {
tracing::warn!(
cache_dir = %persistor.cache_dir().display(),
error = %format!("{e:#}"),
"ADR-027 iter-6b.3: cfg_from_cache failed; disk \
write-through skipped (in-memory store still proceeds)"
);
}
}
}
self.lcp_registry.store(
key,
prompt_tokens,
vec![std::sync::Arc::new(snapshot)],
sliding_window,
linear_capacity,
)
}
pub fn hydrate_lcp_registry_from_disk(
&mut self,
kv_cache: &crate::inference::models::qwen35::kv_cache::HybridKvCache,
device: &mlx_native::MlxDevice,
) {
let persistor = match &self.disk_persistor {
Some(p) => std::sync::Arc::clone(p),
None => return, };
let cfg = match crate::serve::kv_persist::families::qwen35_hybrid_persistor::cfg_from_cache(
kv_cache,
crate::serve::kv_persist::families::qwen35_hybrid_persistor::FullAttnCodec::F32Dense,
) {
Ok(c) => c,
Err(e) => {
tracing::warn!(
error = %format!("{e:#}"),
"ADR-027 iter-6b.3: hydrate_lcp_registry: cfg_from_cache failed; skipping"
);
return;
}
};
let fingerprint_hex =
crate::serve::kv_persist::families::qwen35_disk_persistor::Qwen35DiskPersistor::fingerprint_hex_for(&cfg);
if self.lcp_hydrated_for_cfg.contains(&fingerprint_hex) {
return; }
self.lcp_hydrated_for_cfg.insert(fingerprint_hex.clone());
let triples = match persistor.hydrate_for_cfg(&cfg, device) {
Ok(v) => v,
Err(e) => {
tracing::warn!(
cache_dir = %persistor.cache_dir().display(),
fingerprint = %fingerprint_hex,
error = %format!("{e:#}"),
"ADR-027 iter-6b.3: hydrate_lcp_registry: hydrate_for_cfg failed; \
skipping (registry remains empty for this cfg)"
);
return;
}
};
let total = triples.len();
let mut inserted = 0usize;
for (key_hex, snap, sidecar) in triples {
let key = crate::serve::kv_persist::lcp_registry::LcpKey {
model_fingerprint: sidecar.model_fingerprint.clone(),
tenant_id: sidecar.tenant_id.clone(),
params_hash: sidecar.params_hash,
};
let sliding_window = sidecar.sliding_window as usize;
let linear_capacity = sidecar.linear_capacity as usize;
match self.lcp_registry.store(
key,
sidecar.prompt_tokens.clone(),
vec![std::sync::Arc::new(snap)],
sliding_window,
linear_capacity,
) {
Ok(()) => inserted += 1,
Err(e) => {
tracing::warn!(
key_hex = %key_hex,
error = ?e,
"ADR-027 iter-6b.3: lcp_registry.store rejected hydrated entry; \
skipping (file may exceed byte budget or be empty)"
);
}
}
}
tracing::info!(
cache_dir = %persistor.cache_dir().display(),
fingerprint = %fingerprint_hex,
files_on_disk = total,
entries_inserted = inserted,
registry_len = self.lcp_registry.len(),
"ADR-027 iter-6b.3: hydrate_lcp_registry_from_disk complete"
);
}
pub fn provision_multi_seq_kv_for_slot_aware(&mut self, max_slots: u32) -> anyhow::Result<()> {
use anyhow::Context;
if max_slots == 0 {
anyhow::bail!(
"ADR-040 C2d: provision_multi_seq_kv_for_slot_aware called with \
max_slots == 0; spawn_with_mode invariant is max_slots >= 1 \
(EngineMode::SlotAware variant enforces this at the API \
boundary — caller violated)"
);
}
let device = mlx_native::MlxDevice::new()
.context("ADR-040 C2d: MlxDevice::new for Qwen35 multi-seq KV provisioning")?;
let max_seq_len = self.model.cfg.max_position_embeddings / max_slots;
let cache = HybridKvCache::new_with_options(
&self.model.cfg,
&device,
max_seq_len,
max_slots,
self.tq_kv_active,
)
.with_context(|| {
format!(
"ADR-040 C2d: HybridKvCache::new_with_options(max_seq_len={}, \
n_seqs={}, tq_kv_active={}) for Qwen35 SlotAware provisioning",
max_seq_len, max_slots, self.tq_kv_active
)
})?;
self.persistent_kv_cache = Some(cache);
Ok(())
}
}
fn lcp_key_to_filename_hex(key: &crate::serve::kv_persist::lcp_registry::LcpKey) -> String {
use sha2::{Digest, Sha256};
let mut h = Sha256::new();
h.update(b"QH35-lcp-key-fname-v1");
h.update(&key.model_fingerprint.0);
h.update(key.tenant_id.as_bytes());
h.update(&key.params_hash.to_le_bytes());
let digest = h.finalize();
hex::encode(&digest[..16])
}
impl LoadInfoBuilder for Qwen35LoadedModel {
fn build_load_info(
&self,
gguf: &mlx_native::gguf::GgufFile,
load_wall_clock: Duration,
kv_cache_budget_bytes: Option<u64>,
kv_spill_active: bool,
) -> LoadInfo {
let cfg = &self.model.cfg;
LoadInfo {
model_id: self.model_id.clone(),
arch_str: load_info::arch_str_from_gguf(gguf),
arch_family: ArchFamily::Qwen35,
model_path: self.model_path.clone(),
on_disk_bytes: load_info::on_disk_bytes(&self.model_path),
backend_chip: mlx_native::MlxDevice::new()
.map(|d| d.name())
.unwrap_or_else(|_| "Apple GPU".to_string()),
backend: "mlx-native",
n_layers: cfg.num_hidden_layers,
hidden_size: self.hidden_size as u32,
vocab_size: self.vocab_size as u32,
n_attention_heads: cfg.num_attention_heads,
n_key_value_heads: cfg.num_key_value_heads,
head_dim: cfg.head_dim,
sliding_window: None,
full_attention_interval: Some(cfg.full_attention_interval),
max_context_length: self.context_length.map(|v| v as u32),
moe: cfg.moe.as_ref().map(|m| MoeShape {
n_experts: m.num_experts,
n_experts_per_tok: m.num_experts_per_tok,
}),
quant_label: self.quant_type.clone(),
quant_bpw: load_info::compute_bpw(gguf),
tokenizer_source: TokenizerSource::GgufEmbedded,
eos_token_ids: self.eos_token_ids.clone(),
bos_token_id: gguf.metadata_u32("tokenizer.ggml.bos_token_id"),
chat_template_source: if gguf.metadata_string("tokenizer.chat_template").is_some() {
ChatTemplateSource::GgufEmbedded
} else {
ChatTemplateSource::None
},
provenance: self.provenance.clone(),
vision_projector: None,
load_wall_clock,
resident_weight_bytes: None,
kv_cache_budget_bytes,
kv_spill_active,
tq_kv_active: self.tq_kv_active,
kv_bytes_per_token_override: None,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct HybridPromptCacheKey {
pub max_tokens: usize,
pub stop_strings: Vec<String>,
}
impl HybridPromptCacheKey {
pub fn from_params(params: &SamplingParams) -> Self {
Self {
max_tokens: params.max_tokens,
stop_strings: params.stop_strings.clone(),
}
}
}
#[derive(Debug)]
pub struct HybridPromptCache {
cached_prompt_tokens: Vec<u32>,
snapshot: Option<HybridKvCacheSnapshot>,
first_decoded_token: u32,
gen_params: Option<HybridPromptCacheKey>,
}
impl Default for HybridPromptCache {
fn default() -> Self {
Self::new()
}
}
impl HybridPromptCache {
pub fn new() -> Self {
Self {
cached_prompt_tokens: Vec::new(),
snapshot: None,
first_decoded_token: 0,
gen_params: None,
}
}
pub fn try_match(&self, new_prompt: &[u32], new_params: &SamplingParams) -> Option<usize> {
if !is_greedy_eligible(new_params) {
return None;
}
if self.snapshot.is_none() || self.gen_params.is_none() {
return None;
}
if self.cached_prompt_tokens.is_empty() {
return None;
}
if self.cached_prompt_tokens.as_slice() != new_prompt {
return None;
}
let request_key = HybridPromptCacheKey::from_params(new_params);
if self.gen_params.as_ref() != Some(&request_key) {
return None;
}
Some(new_prompt.len())
}
pub fn snapshot(&self) -> Option<&HybridKvCacheSnapshot> {
self.snapshot.as_ref()
}
pub fn first_decoded_token(&self) -> u32 {
self.first_decoded_token
}
pub fn update(
&mut self,
prompt: Vec<u32>,
snapshot: HybridKvCacheSnapshot,
first_decoded_token: u32,
params: &SamplingParams,
) {
if !is_greedy_eligible(params) {
return;
}
self.cached_prompt_tokens = prompt;
self.snapshot = Some(snapshot);
self.first_decoded_token = first_decoded_token;
self.gen_params = Some(HybridPromptCacheKey::from_params(params));
}
#[allow(dead_code)]
pub fn clear(&mut self) {
self.cached_prompt_tokens.clear();
self.snapshot = None;
self.first_decoded_token = 0;
self.gen_params = None;
}
pub fn has_entry(&self) -> bool {
self.snapshot.is_some()
}
}
fn is_greedy_eligible(params: &SamplingParams) -> bool {
!(params.temperature > 0.0
|| params.top_k > 0
|| params.top_p < 1.0
|| params.repetition_penalty != 1.0
|| params.seed.is_some())
}
static LCP_STORE_NOTIFY: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
fn lcp_notify_once(bit: u64, msg: std::string::String) {
let prev = LCP_STORE_NOTIFY.fetch_or(bit, std::sync::atomic::Ordering::Relaxed);
if prev & bit == 0 {
eprintln!("{msg}");
}
}
fn lcp_store_skip_notify(stride_aligned: bool, lcp_resume_enabled: bool, mid_store_disabled: bool) {
if !stride_aligned {
return;
}
if !lcp_resume_enabled {
lcp_notify_once(
1,
"[hf2q qwen35 lcp store] SKIPPED: HF2Q_KV_LCP_RESUME is off — \
no prefix checkpoints will be written (once-per-process notice)"
.to_string(),
);
} else if mid_store_disabled {
lcp_notify_once(
2,
"[hf2q qwen35 lcp store] SKIPPED: HF2Q_KV_LCP_DISABLE_MID_STORE=1 \
— mid-prefill checkpoints disabled (once-per-process notice)"
.to_string(),
);
}
}
fn lcp_snapshot_error_notify<E: std::fmt::Debug>(phase: &str, chunk_pos: usize, err: &E) {
lcp_notify_once(
4,
format!(
"[hf2q qwen35 lcp store] snapshot FAILED ({phase}, chunk_pos={chunk_pos}): \
{err:?} — checkpoints are NOT being written; further occurrences \
suppressed (once-per-process notice)"
),
);
}
fn lcp_store_error_notify<E: std::fmt::Debug>(phase: &str, chunk_pos: usize, err: &E) {
lcp_notify_once(
8,
format!(
"[hf2q qwen35 lcp store] store FAILED ({phase}, chunk_pos={chunk_pos}): \
{err:?} — check the registry byte budget; further occurrences \
suppressed (once-per-process notice)"
),
);
}
use crate::inference::models::qwen35::io_heads::greedy_argmax_last_token;
use crate::serve::sampler_pure::{self, SamplingParams as SamplerPureParams};
use mlx_native::MlxDevice;
use super::engine::GenerationResult;
use super::registry::{
ModelRegistration, ReasoningSplitter, SplitSlot, ToolCallEvent, ToolCallSplitter,
};
use super::sse::{DeltaKind, GenerationEvent, StreamStats};
fn prefill_positions_for(prompt_len: usize) -> Vec<i32> {
let mut flat = vec![0i32; 4 * prompt_len];
for axis in 0..4 {
for t in 0..prompt_len {
flat[axis * prompt_len + t] = t as i32;
}
}
flat
}
fn alloc_kv_cache_for_request(
qwen: &Qwen35LoadedModel,
device: &MlxDevice,
prompt_len: usize,
max_tokens: usize,
) -> Result<HybridKvCache> {
let max_seq = (prompt_len + max_tokens + 64)
.max(128)
.min(qwen.model.cfg.max_position_embeddings as usize);
HybridKvCache::new_with_options(
&qwen.model.cfg,
device,
max_seq as u32,
1,
qwen.tq_kv_active,
)
.context("HybridKvCache::new_with_options")
}
fn sample_logits_qwen35(logits: &mut [f32], params: &SamplingParams, generated: &[u32]) -> u32 {
let sp = SamplerPureParams {
temperature: params.temperature as f64,
top_p: params.top_p as f64,
top_k: params.top_k,
min_p: 0.0,
repetition_penalty: effective_repetition_penalty(params),
max_tokens: params.max_tokens,
};
sampler_pure::sample_token(logits, &sp, generated)
}
fn sample_logits_qwen35_with_logprob(
logits: &mut [f32],
params: &SamplingParams,
generated: &[u32],
) -> (u32, f32) {
let sp = SamplerPureParams {
temperature: params.temperature as f64,
top_p: params.top_p as f64,
top_k: params.top_k,
min_p: 0.0,
repetition_penalty: effective_repetition_penalty(params),
max_tokens: params.max_tokens,
};
sampler_pure::sample_token_with_logprob(logits, &sp, generated)
}
fn build_lcp_key_for_qwen35(
qwen: &Qwen35LoadedModel,
_params: &SamplingParams,
) -> crate::serve::kv_persist::lcp_registry::LcpKey {
use crate::serve::kv_persist::format::compute_model_fingerprint;
let (producer_version, source_sha256) = match &qwen.provenance {
crate::core::provenance::Provenance::Hf2q {
producer_version,
source_sha256,
..
} => (producer_version.as_str(), source_sha256.as_str()),
crate::core::provenance::Provenance::External => ("", ""),
};
let chat_template_hash = match &qwen.provenance {
crate::core::provenance::Provenance::Hf2q { .. } => {
super::kv_spill_descriptor::KvSpillProvenance::hash_chat_template(&qwen.chat_template)
}
crate::core::provenance::Provenance::External => String::new(),
};
let quant = qwen.quant_type.as_deref().unwrap_or("");
let fp = compute_model_fingerprint(
&qwen.model_id,
quant,
producer_version,
source_sha256,
&chat_template_hash,
);
crate::serve::kv_persist::lcp_registry::LcpKey {
model_fingerprint: fp,
tenant_id: String::new(),
params_hash: 0,
}
}
fn build_lcp_key_for_qwen35_chunk(
qwen: &Qwen35LoadedModel,
params: &SamplingParams,
chunk_position: usize,
) -> crate::serve::kv_persist::lcp_registry::LcpKey {
let mut key = build_lcp_key_for_qwen35(qwen, params);
if chunk_position > 0 {
key.tenant_id = format!("qwen35:lcp_chunk:{chunk_position}");
}
key
}
pub fn qwen35_lcp_registry_byte_budget() -> u64 {
crate::serve::kv_persist::lcp_registry::default_lcp_byte_budget()
}
fn qwen35_hit_stop_string(text: &str, stops: &[String]) -> bool {
if stops.is_empty() {
return false;
}
stops
.iter()
.any(|s| !s.is_empty() && text.ends_with(s.as_str()))
}
fn qwen35_strip_trailing_stop(text: &mut String, stops: &[String]) {
for s in stops {
if !s.is_empty() && text.ends_with(s) {
let new_len = text.len() - s.len();
text.truncate(new_len);
return;
}
}
}
pub fn generate_qwen35_once(
qwen: &mut Qwen35LoadedModel,
prompt_tokens: &[u32],
params: &SamplingParams,
registration: Option<&ModelRegistration>,
) -> Result<GenerationResult> {
anyhow::ensure!(
!prompt_tokens.is_empty(),
"generate_qwen35_once: empty prompt_tokens"
);
let prompt_len = prompt_tokens.len();
let max_tokens = params.max_tokens.max(1);
let is_greedy = is_greedy_eligible(params);
let want_logprobs = params.logprobs;
let mut logprobs_vec: Option<Vec<f32>> = if want_logprobs {
Some(Vec::with_capacity(max_tokens))
} else {
None
};
let device =
MlxDevice::new().map_err(|e| anyhow::anyhow!("MlxDevice::new (qwen35 generate): {e}"))?;
let mut kv_cache = alloc_kv_cache_for_request(qwen, &device, prompt_len, max_tokens)?;
qwen.hydrate_lcp_registry_from_disk(&kv_cache, &device);
let prompt_cache_hit = qwen.prompt_cache.try_match(prompt_tokens, params).is_some();
let mut lcp_resume_start: usize = 0;
if !prompt_cache_hit {
let stride_for_observe = crate::debug::INVESTIGATION_ENV.kv_lcp_deltanet_checkpoint_stride;
let base_key_for_observe = build_lcp_key_for_qwen35(qwen, params);
let detected = crate::serve::kv_persist::lcp_registry::probe_lcp_opportunity_chunk_aligned(
&mut qwen.lcp_registry,
prompt_tokens,
stride_for_observe,
false, |chunk_pos| {
let mut key = base_key_for_observe.clone();
if chunk_pos > 0 {
key.tenant_id = format!("qwen35:lcp_chunk:{chunk_pos}");
}
key
},
);
if let Some(sink) = qwen.kv_metrics_sink.as_ref() {
sink.record_lcp_probe(detected);
}
let _ = detected;
let lcp_resume_enabled = crate::serve::api::engine::effective_kv_lcp_resume(
crate::debug::INVESTIGATION_ENV.kv_lcp_resume,
true,
);
if lcp_resume_enabled {
let stride = crate::debug::INVESTIGATION_ENV.kv_lcp_deltanet_checkpoint_stride;
if stride == 0 {
eprintln!(
"[hf2q qwen35 lcp probe] stride=0; skipping descending \
scan (HF2Q_KV_LCP_DELTANET_CHECKPOINT_STRIDE must be > 0)"
);
}
let max_chunk_pos = if stride == 0 {
0
} else {
(prompt_tokens.len() / stride).saturating_mul(stride)
};
eprintln!(
"[hf2q qwen35 lcp probe] enabled, registry_len={}, prompt_len={}, \
stride={}, scanning chunk positions [{stride}..={max_chunk_pos}]",
qwen.lcp_registry.len(),
prompt_tokens.len(),
stride,
);
let mut hit = false;
let mut chunk_pos = max_chunk_pos;
while chunk_pos >= stride && !hit {
let chunk_key = build_lcp_key_for_qwen35_chunk(qwen, params, chunk_pos);
if let Some(prefix) = qwen.lcp_registry.lookup(&chunk_key, prompt_tokens) {
if prefix.k == prefix.cached_prompt_len && prefix.k < prompt_tokens.len() {
let snapshot: &HybridKvCacheSnapshot = &prefix.dense_kvs[0];
let restore_start = Instant::now();
kv_cache
.restore_partial(snapshot, prefix.k)
.context("qwen35 lcp_registry restore_partial")?;
let restore_ms = restore_start.elapsed().as_micros() as f64 / 1000.0;
lcp_resume_start = prefix.k;
eprintln!(
"[hf2q qwen35 lcp resume] STRIDE-ALIGNED HIT — \
restoring at k={} (cached_prompt_len={}, \
chunk_pos={}, restore_ms={:.3})",
prefix.k, prefix.cached_prompt_len, chunk_pos, restore_ms
);
hit = true;
} else {
eprintln!(
"[hf2q qwen35 lcp probe] PARTIAL HIT at \
chunk_pos={} — k={} < cached_prompt_len={}, \
continuing descent",
chunk_pos, prefix.k, prefix.cached_prompt_len
);
}
}
if chunk_pos == stride {
break;
}
chunk_pos -= stride;
}
if !hit {
eprintln!(
"[hf2q qwen35 lcp probe] no stride-aligned match \
(registry_len={})",
qwen.lcp_registry.len()
);
}
}
}
let prefill_start = Instant::now();
let mut next_token: u32;
if prompt_cache_hit {
let snap = qwen
.prompt_cache
.snapshot()
.expect("try_match returned Some implies snapshot Some");
kv_cache
.restore_from(snap)
.context("prompt_cache restore")?;
next_token = qwen.prompt_cache.first_decoded_token();
tracing::debug!(
"qwen35 prompt_cache: HIT — {} tokens; prefill skipped",
prompt_len
);
} else {
let stride = crate::debug::INVESTIGATION_ENV.kv_lcp_deltanet_checkpoint_stride;
let lcp_resume_enabled = crate::debug::INVESTIGATION_ENV.kv_lcp_resume;
let chunked_eligible = stride > 0
&& prompt_len > stride
&& crate::debug::INVESTIGATION_ENV.kv_lcp_chunked_prefill;
if lcp_resume_start > 0 && stride > 0 {
anyhow::ensure!(
lcp_resume_start % stride == 0,
"qwen35 prefill: lcp_resume_start ({}) must be stride-\
aligned ({}) — registry only stores at stride boundaries",
lcp_resume_start,
stride
);
}
let prefill_logits = if chunked_eligible {
anyhow::ensure!(
lcp_resume_start % stride == 0,
"qwen35 chunked prefill: lcp_resume_start ({}) must be \
stride-aligned ({}) — snapshots are stored only at \
stride boundaries, so this should be guaranteed by the \
probe site",
lcp_resume_start,
stride
);
let first_chunk_idx = lcp_resume_start / stride;
let n_chunks = (prompt_len + stride - 1) / stride;
eprintln!(
"[hf2q qwen35 chunked prefill] {} chunks (stride={}, \
prompt_len={}, first_chunk_idx={})",
n_chunks, stride, prompt_len, first_chunk_idx
);
let mut last_logits: Option<Vec<f32>> = None;
for chunk_idx in first_chunk_idx..n_chunks {
let k_start = chunk_idx * stride;
let k_end = ((chunk_idx + 1) * stride).min(prompt_len);
let chunk_seq_len = k_end - k_start;
let chunk_tokens = &prompt_tokens[k_start..k_end];
let mut chunk_positions = vec![0i32; 4 * chunk_seq_len];
for axis in 0..4 {
for t in 0..chunk_seq_len {
chunk_positions[axis * chunk_seq_len + t] = (k_start + t) as i32;
}
}
let logits = qwen
.model
.forward_gpu_last_logits(
chunk_tokens,
&chunk_positions,
&mut kv_cache,
SlotId(0),
)
.with_context(|| {
format!(
"qwen35 chunked prefill: chunk {}/{} (k_start={}, k_end={}, seq_len={})",
chunk_idx + 1, n_chunks, k_start, k_end, chunk_seq_len
)
})?;
if chunk_idx == n_chunks - 1 {
last_logits = Some(logits.clone());
}
let stride_aligned = k_end % stride == 0;
let mid_store_disabled =
std::env::var("HF2Q_KV_LCP_DISABLE_MID_STORE").as_deref() == Ok("1");
lcp_store_skip_notify(stride_aligned, lcp_resume_enabled, mid_store_disabled);
if lcp_resume_enabled && stride_aligned && !mid_store_disabled {
match kv_cache.snapshot(&device) {
Ok(snap) => {
let chunk_key = build_lcp_key_for_qwen35_chunk(qwen, params, k_end);
let linear_capacity = kv_cache
.linear_attn
.first()
.map(|s| s.recurrent.byte_len())
.unwrap_or(0);
if let Err(e) = qwen.store_lcp_with_disk_writeback(
&kv_cache,
chunk_key,
prompt_tokens[..k_end].to_vec(),
snap,
0,
linear_capacity,
) {
lcp_store_error_notify("mid-prefill", k_end, &e);
} else {
eprintln!(
"[hf2q qwen35 lcp store] mid-prefill snapshot \
at chunk_pos={k_end} (registry_len_after={})",
qwen.lcp_registry.len()
);
}
}
Err(e) => {
lcp_snapshot_error_notify("mid-prefill", k_end, &e);
}
}
}
}
last_logits.expect("at least one chunk in chunked prefill")
} else if lcp_resume_start > 0 {
let suffix_tokens = &prompt_tokens[lcp_resume_start..];
let suffix_len = suffix_tokens.len();
let mut suffix_positions = vec![0i32; 4 * suffix_len];
for axis in 0..4 {
for t in 0..suffix_len {
suffix_positions[axis * suffix_len + t] = (lcp_resume_start + t) as i32;
}
}
eprintln!(
"[hf2q qwen35 lcp resume] suffix prefill {} tokens \
(lcp_resume_start={}, prompt_len={})",
suffix_len, lcp_resume_start, prompt_len
);
qwen.model
.forward_gpu_last_logits(suffix_tokens, &suffix_positions, &mut kv_cache, SlotId(0))
.context("Qwen35Model::forward_gpu_last_logits (LCP resume suffix)")?
} else {
let positions = prefill_positions_for(prompt_len);
qwen.model
.forward_gpu_last_logits(prompt_tokens, &positions, &mut kv_cache, SlotId(0))
.context("Qwen35Model::forward_gpu_last_logits (prefill)")?
};
anyhow::ensure!(
prefill_logits.len() == qwen.vocab_size,
"qwen35 prefill logits len {} != vocab_size {}",
prefill_logits.len(),
qwen.vocab_size
);
if want_logprobs {
let mut logits = prefill_logits.clone();
let (tok, lp) = sample_logits_qwen35_with_logprob(&mut logits, params, &[]);
next_token = tok;
if let Some(v) = logprobs_vec.as_mut() {
v.push(lp);
}
} else if is_greedy {
next_token = greedy_argmax_last_token(&prefill_logits, qwen.vocab_size as u32);
} else {
let mut logits = prefill_logits.clone();
next_token = sample_logits_qwen35(&mut logits, params, &[]);
}
if is_greedy {
match kv_cache.snapshot(&device) {
Ok(snap) => {
qwen.prompt_cache
.update(prompt_tokens.to_vec(), snap, next_token, params)
}
Err(e) => {
eprintln!("[hf2q qwen35 lcp store] prompt_cache snapshot failed: {e}");
}
}
}
{
let chunked_already_stored = stride > 0
&& prompt_len > stride
&& crate::debug::INVESTIGATION_ENV.kv_lcp_chunked_prefill;
let prompt_stride_aligned = stride > 0 && prompt_len > 0 && prompt_len % stride == 0;
let should_store_eof =
lcp_resume_enabled && !chunked_already_stored && prompt_stride_aligned;
if should_store_eof {
eprintln!(
"[hf2q qwen35 lcp store] end-of-prefill snapshot for prompt_len={} \
(chunk_pos={prompt_len}, registry_len_before={})",
prompt_tokens.len(),
qwen.lcp_registry.len()
);
match kv_cache.snapshot(&device) {
Ok(snap) => {
let chunk_key = build_lcp_key_for_qwen35_chunk(qwen, params, prompt_len);
let linear_capacity = kv_cache
.linear_attn
.first()
.map(|s| s.recurrent.byte_len())
.unwrap_or(0);
if let Err(e) = qwen.store_lcp_with_disk_writeback(
&kv_cache,
chunk_key,
prompt_tokens.to_vec(),
snap,
0,
linear_capacity,
) {
lcp_store_error_notify("end-of-prefill", prompt_len, &e);
}
}
Err(e) => {
lcp_snapshot_error_notify("end-of-prefill", prompt_len, &e);
}
}
}
}
}
let prefill_duration = prefill_start.elapsed();
let decode_start = Instant::now();
let mut generated_tokens: Vec<u32> = Vec::with_capacity(max_tokens);
generated_tokens.push(next_token);
let first_fragment = qwen
.tokenizer
.decode(&[next_token], false)
.unwrap_or_default();
let mut decoded_text = first_fragment.clone();
let mut finish_reason: &'static str = "length";
if qwen.eos_token_ids.contains(&next_token) {
finish_reason = "stop";
} else if qwen35_hit_stop_string(&decoded_text, ¶ms.stop_strings) {
finish_reason = "stop";
qwen35_strip_trailing_stop(&mut decoded_text, ¶ms.stop_strings);
} else {
for step in 1..max_tokens {
let pos = (prompt_len + step - 1) as i32;
if pos as u32 >= kv_cache.max_seq_len {
tracing::warn!(
pos,
max_seq = kv_cache.max_seq_len,
"qwen35 decode: hit kv-cache bound; stopping with finish=length",
);
break;
}
let decode_positions = vec![pos; 4];
next_token = if want_logprobs {
let logits_full = qwen
.model
.forward_gpu_last_logits(
&[next_token],
&decode_positions,
&mut kv_cache,
SlotId(0),
)
.with_context(|| {
format!("forward_gpu_last_logits decode step {step} (logprobs)")
})?;
let mut logits = logits_full;
let (tok, lp) =
sample_logits_qwen35_with_logprob(&mut logits, params, &generated_tokens);
if let Some(v) = logprobs_vec.as_mut() {
v.push(lp);
}
tok
} else if is_greedy {
qwen.model
.forward_gpu_greedy(&[next_token], &decode_positions, &mut kv_cache, SlotId(0))
.with_context(|| format!("forward_gpu_greedy decode step {step}"))?
} else {
let logits_full = qwen
.model
.forward_gpu_last_logits(
&[next_token],
&decode_positions,
&mut kv_cache,
SlotId(0),
)
.with_context(|| format!("forward_gpu_last_logits decode step {step}"))?;
let mut logits = logits_full;
sample_logits_qwen35(&mut logits, params, &generated_tokens)
};
if qwen.eos_token_ids.contains(&next_token) {
finish_reason = "stop";
break;
}
generated_tokens.push(next_token);
let fragment = qwen
.tokenizer
.decode(&[next_token], false)
.unwrap_or_default();
decoded_text.push_str(&fragment);
if qwen35_hit_stop_string(&decoded_text, ¶ms.stop_strings) {
finish_reason = "stop";
qwen35_strip_trailing_stop(&mut decoded_text, ¶ms.stop_strings);
break;
}
}
}
let decode_duration = decode_start.elapsed();
let (content, reasoning_text) = match registration {
Some(reg) if reg.has_reasoning() => super::registry::split_full_output_forced(
reg,
&decoded_text,
params.reasoning_forced_open,
),
_ => (decoded_text, None),
};
let reasoning_token_count = match registration {
Some(reg) if reg.has_reasoning() => {
let mut sp =
super::registry::make_reasoning_splitter(reg, params.reasoning_forced_open);
let mut count = 0usize;
for &tok in &generated_tokens {
let frag = qwen.tokenizer.decode(&[tok], false).unwrap_or_default();
if let Some(splitter) = sp.as_mut() {
let _ = splitter.feed(&frag);
if splitter.in_reasoning() {
count += 1;
}
}
}
count
}
_ => 0,
};
Ok(GenerationResult {
text: content,
reasoning_text,
prompt_tokens: prompt_len,
completion_tokens: generated_tokens.len(),
reasoning_tokens: if reasoning_token_count > 0 {
Some(reasoning_token_count)
} else {
None
},
finish_reason,
prefill_duration,
decode_duration,
cached_tokens: if prompt_cache_hit { prompt_len } else { 0 },
logprobs: logprobs_vec,
})
}
pub fn generate_qwen35_once_slot_aware(
qwen: &mut Qwen35LoadedModel,
prompt_tokens: &[u32],
params: &SamplingParams,
registration: Option<&ModelRegistration>,
kv_cache: &mut HybridKvCache,
slot_id: SlotId,
) -> Result<GenerationResult> {
anyhow::ensure!(
!prompt_tokens.is_empty(),
"generate_qwen35_once_slot_aware: empty prompt_tokens"
);
anyhow::ensure!(
slot_id.0 < kv_cache.n_seqs,
"generate_qwen35_once_slot_aware: SlotOutOfRange slot={} max_slots={} \
(ADR-040 iter-C2d-cont-kernel iter-1)",
slot_id.0,
kv_cache.n_seqs,
);
let prompt_len = prompt_tokens.len();
let max_tokens = params.max_tokens.max(1);
let need_seq = prompt_len + max_tokens + 64;
if need_seq > kv_cache.max_seq_len as usize {
return Err(anyhow::anyhow!(
"generate_qwen35_once_slot_aware: per-request need_seq={} exceeds \
persistent cache max_seq_len={} (slot={} prompt_len={} max_tokens={}). \
ADR-040 iter-C2d-cont-kernel iter-1 sizes the persistent cache to \
cfg.max_position_embeddings; reduce max_tokens or use a shorter prompt.",
need_seq,
kv_cache.max_seq_len,
slot_id.0,
prompt_len,
max_tokens
));
}
let is_greedy = is_greedy_eligible(params);
let want_logprobs = params.logprobs;
let mut logprobs_vec: Option<Vec<f32>> = if want_logprobs {
Some(Vec::with_capacity(max_tokens))
} else {
None
};
kv_cache
.reset_for_slot(slot_id)
.context("ADR-040 iter-C2d-cont-kernel iter-1: reset_for_slot at entry")?;
let prompt_cache_hit = qwen.prompt_cache.try_match(prompt_tokens, params).is_some();
let prefill_start = Instant::now();
let next_token: u32;
if prompt_cache_hit {
let snap = qwen
.prompt_cache
.snapshot()
.expect("try_match returned Some implies snapshot Some");
kv_cache
.restore_partial(snap, prompt_len)
.context("ADR-040 iter-C2d-cont-kernel iter-1: prompt_cache restore_partial")?;
next_token = qwen.prompt_cache.first_decoded_token();
tracing::debug!(
"qwen35 slot-aware prompt_cache: HIT slot={} prompt_len={} prefill skipped",
slot_id.0,
prompt_len
);
} else {
let positions = prefill_positions_for(prompt_len);
let prefill_logits = qwen
.model
.forward_gpu_last_logits(prompt_tokens, &positions, kv_cache, slot_id)
.context("Qwen35Model::forward_gpu_last_logits (slot-aware prefill)")?;
anyhow::ensure!(
prefill_logits.len() == qwen.vocab_size,
"qwen35 slot-aware prefill logits len {} != vocab_size {}",
prefill_logits.len(),
qwen.vocab_size
);
if is_greedy && !want_logprobs {
next_token = greedy_argmax_last_token(&prefill_logits, qwen.vocab_size as u32);
} else {
let mut logits = prefill_logits;
if let Some(ref mut lps) = logprobs_vec {
let (tok, lp) = sample_logits_qwen35_with_logprob(&mut logits, params, &[]);
lps.push(lp);
next_token = tok;
} else {
next_token = sample_logits_qwen35(&mut logits, params, &[]);
}
}
}
let prefill_duration = prefill_start.elapsed();
let device = MlxDevice::new()
.map_err(|e| anyhow::anyhow!("MlxDevice::new (qwen35 slot-aware decode): {e}"))?;
let _ = &device; let decode_start = Instant::now();
let mut generated_tokens: Vec<u32> = Vec::with_capacity(max_tokens);
generated_tokens.push(next_token);
let mut decoded_text = qwen
.tokenizer
.decode(&[next_token], false)
.unwrap_or_default();
let stops = ¶ms.stop_strings;
let mut finish_reason: &'static str = "length";
if qwen.eos_token_ids.contains(&next_token) {
generated_tokens.pop();
decoded_text.clear();
finish_reason = "stop";
} else if qwen35_hit_stop_string(&decoded_text, stops) {
qwen35_strip_trailing_stop(&mut decoded_text, stops);
finish_reason = "stop";
}
let mut step = 1usize;
while step < max_tokens && finish_reason == "length" {
let pos = prompt_len + step - 1;
let pos_i32 = pos as i32;
let positions: Vec<i32> = vec![pos_i32; 4];
let last_input = &generated_tokens[generated_tokens.len() - 1..];
let tok = if is_greedy && !want_logprobs {
qwen.model
.forward_gpu_greedy(last_input, &positions, kv_cache, slot_id)
.with_context(|| {
format!(
"Qwen35Model::forward_gpu_greedy (slot-aware decode step {step}; \
ADR-040 §6.1.50 iter-G)"
)
})?
} else {
let logits = qwen
.model
.forward_gpu_last_logits(last_input, &positions, kv_cache, slot_id)
.with_context(|| {
format!("Qwen35Model::forward_gpu_last_logits (slot-aware decode step {step})")
})?;
anyhow::ensure!(
logits.len() == qwen.vocab_size,
"qwen35 slot-aware decode logits len {} != vocab_size {}",
logits.len(),
qwen.vocab_size
);
let mut logits = logits;
if let Some(ref mut lps) = logprobs_vec {
let (tok, lp) =
sample_logits_qwen35_with_logprob(&mut logits, params, &generated_tokens);
lps.push(lp);
tok
} else {
sample_logits_qwen35(&mut logits, params, &generated_tokens)
}
};
if qwen.eos_token_ids.contains(&tok) {
finish_reason = "stop";
break;
}
generated_tokens.push(tok);
let frag = qwen.tokenizer.decode(&[tok], false).unwrap_or_default();
decoded_text.push_str(&frag);
if qwen35_hit_stop_string(&decoded_text, stops) {
qwen35_strip_trailing_stop(&mut decoded_text, stops);
finish_reason = "stop";
break;
}
step += 1;
}
let decode_duration = decode_start.elapsed();
kv_cache
.reset_for_slot(slot_id)
.context("ADR-040 iter-C2d-cont-kernel iter-1: reset_for_slot at exit")?;
let (content_text, reasoning_text) = match registration {
Some(reg) if reg.has_reasoning() => super::registry::split_full_output_forced(
reg,
&decoded_text,
params.reasoning_forced_open,
),
_ => (decoded_text, None),
};
let reasoning_token_count = match registration {
Some(reg) if reg.has_reasoning() => {
let mut sp =
super::registry::make_reasoning_splitter(reg, params.reasoning_forced_open);
let mut count = 0usize;
for &tok in &generated_tokens {
let frag = qwen.tokenizer.decode(&[tok], false).unwrap_or_default();
if let Some(splitter) = sp.as_mut() {
let _ = splitter.feed(&frag);
if splitter.in_reasoning() {
count += 1;
}
}
}
count
}
_ => 0,
};
Ok(GenerationResult {
text: content_text,
reasoning_text,
prompt_tokens: prompt_len,
completion_tokens: generated_tokens.len(),
reasoning_tokens: if reasoning_token_count > 0 {
Some(reasoning_token_count)
} else {
None
},
finish_reason,
prefill_duration,
decode_duration,
cached_tokens: if prompt_cache_hit { prompt_len } else { 0 },
logprobs: logprobs_vec,
})
}
pub(crate) struct Qwen35TickOutcome {
pub fragment: String,
pub is_reasoning: bool,
pub finished: bool,
}
pub(crate) struct Qwen35DecodeState {
slot_id: SlotId,
prompt_len: usize,
max_tokens: usize,
is_greedy: bool,
want_logprobs: bool,
logprobs_vec: Option<Vec<f32>>,
prompt_cache_hit: bool,
params: SamplingParams,
next_token: u32,
generated_tokens: Vec<u32>,
decoded_text: String,
stop_strings: Vec<String>,
finish_reason: &'static str,
step: usize,
prefill_duration: Duration,
decode_start: Instant,
}
impl Qwen35DecodeState {
pub(crate) fn prefill_seed(
qwen: &mut Qwen35LoadedModel,
prompt_tokens: &[u32],
params: &SamplingParams,
kv_cache: &mut HybridKvCache,
slot_id: SlotId,
) -> Result<Self> {
anyhow::ensure!(
!prompt_tokens.is_empty(),
"Qwen35DecodeState::prefill_seed: empty prompt_tokens"
);
anyhow::ensure!(
slot_id.0 < kv_cache.n_seqs,
"Qwen35DecodeState::prefill_seed: SlotOutOfRange slot={} max_slots={} \
(ADR-040 Phase F M1)",
slot_id.0,
kv_cache.n_seqs,
);
let prompt_len = prompt_tokens.len();
let max_tokens = params.max_tokens.max(1);
let need_seq = prompt_len + max_tokens + 64;
if need_seq > kv_cache.max_seq_len as usize {
return Err(anyhow::anyhow!(
"Qwen35DecodeState::prefill_seed: per-request need_seq={} exceeds \
persistent cache max_seq_len={} (slot={} prompt_len={} max_tokens={}). \
ADR-040 Phase F M1; reduce max_tokens or use a shorter prompt.",
need_seq,
kv_cache.max_seq_len,
slot_id.0,
prompt_len,
max_tokens
));
}
let is_greedy = is_greedy_eligible(params);
let want_logprobs = params.logprobs;
let mut logprobs_vec: Option<Vec<f32>> = if want_logprobs {
Some(Vec::with_capacity(max_tokens))
} else {
None
};
kv_cache
.reset_for_slot(slot_id)
.context("ADR-040 Phase F M1: Qwen35 reset_for_slot at entry")?;
let prompt_cache_hit = false;
let prefill_start = Instant::now();
let next_token: u32;
if prompt_cache_hit {
let snap = qwen
.prompt_cache
.snapshot()
.expect("try_match returned Some implies snapshot Some");
kv_cache
.restore_partial(snap, prompt_len)
.context("ADR-040 Phase F M1: Qwen35 prompt_cache restore_partial")?;
next_token = qwen.prompt_cache.first_decoded_token();
} else {
let positions = prefill_positions_for(prompt_len);
let prefill_logits = qwen
.model
.forward_gpu_last_logits(prompt_tokens, &positions, kv_cache, slot_id)
.context("Qwen35Model::forward_gpu_last_logits (slot-aware prefill)")?;
anyhow::ensure!(
prefill_logits.len() == qwen.vocab_size,
"qwen35 slot-aware prefill logits len {} != vocab_size {}",
prefill_logits.len(),
qwen.vocab_size
);
if is_greedy && !want_logprobs {
next_token = greedy_argmax_last_token(&prefill_logits, qwen.vocab_size as u32);
} else {
let mut logits = prefill_logits;
if let Some(ref mut lps) = logprobs_vec {
let (tok, lp) = sample_logits_qwen35_with_logprob(&mut logits, params, &[]);
lps.push(lp);
next_token = tok;
} else {
next_token = sample_logits_qwen35(&mut logits, params, &[]);
}
}
}
let prefill_duration = prefill_start.elapsed();
let decode_start = Instant::now();
let mut generated_tokens: Vec<u32> = Vec::with_capacity(max_tokens);
generated_tokens.push(next_token);
let mut decoded_text = qwen
.tokenizer
.decode(&[next_token], false)
.unwrap_or_default();
let stops = params.stop_strings.clone();
let mut finish_reason: &'static str = "length";
if qwen.eos_token_ids.contains(&next_token) {
generated_tokens.pop();
decoded_text.clear();
finish_reason = "stop";
} else if qwen35_hit_stop_string(&decoded_text, &stops) {
qwen35_strip_trailing_stop(&mut decoded_text, &stops);
finish_reason = "stop";
}
Ok(Qwen35DecodeState {
slot_id,
prompt_len,
max_tokens,
is_greedy,
want_logprobs,
logprobs_vec,
prompt_cache_hit,
params: params.clone(),
next_token,
generated_tokens,
decoded_text,
stop_strings: stops,
finish_reason,
step: 1,
prefill_duration,
decode_start,
})
}
pub(crate) fn finished_at_seed(&self) -> bool {
self.finish_reason != "length"
}
pub(crate) fn seed_fragment(&self) -> String {
self.decoded_text.clone()
}
pub(crate) fn decode_tick(
&mut self,
qwen: &mut Qwen35LoadedModel,
kv_cache: &mut HybridKvCache,
) -> Result<Qwen35TickOutcome> {
if self.step >= self.max_tokens {
return Ok(Qwen35TickOutcome {
fragment: String::new(),
is_reasoning: false,
finished: true,
});
}
let pos = self.prompt_len + self.step - 1;
let pos_i32 = pos as i32;
let positions: Vec<i32> = vec![pos_i32; 4];
let last_input = &self.generated_tokens[self.generated_tokens.len() - 1..];
let tok = if self.is_greedy && !self.want_logprobs {
qwen.model
.forward_gpu_greedy(last_input, &positions, kv_cache, self.slot_id)
.with_context(|| {
format!(
"Qwen35Model::forward_gpu_greedy (slot-aware decode step {}; \
ADR-040 Phase F M1)",
self.step
)
})?
} else {
let logits = qwen
.model
.forward_gpu_last_logits(last_input, &positions, kv_cache, self.slot_id)
.with_context(|| {
format!(
"Qwen35Model::forward_gpu_last_logits (slot-aware decode step {})",
self.step
)
})?;
anyhow::ensure!(
logits.len() == qwen.vocab_size,
"qwen35 slot-aware decode logits len {} != vocab_size {}",
logits.len(),
qwen.vocab_size
);
let mut logits = logits;
if let Some(ref mut lps) = self.logprobs_vec {
let (tok, lp) = sample_logits_qwen35_with_logprob(
&mut logits,
&self.params,
&self.generated_tokens,
);
lps.push(lp);
tok
} else {
sample_logits_qwen35(&mut logits, &self.params, &self.generated_tokens)
}
};
if qwen.eos_token_ids.contains(&tok) {
self.finish_reason = "stop";
return Ok(Qwen35TickOutcome {
fragment: String::new(),
is_reasoning: false,
finished: true,
});
}
self.generated_tokens.push(tok);
let frag = qwen.tokenizer.decode(&[tok], false).unwrap_or_default();
self.decoded_text.push_str(&frag);
if qwen35_hit_stop_string(&self.decoded_text, &self.stop_strings) {
qwen35_strip_trailing_stop(&mut self.decoded_text, &self.stop_strings);
self.finish_reason = "stop";
return Ok(Qwen35TickOutcome {
fragment: frag,
is_reasoning: false,
finished: true,
});
}
self.step += 1;
let finished = self.step >= self.max_tokens;
Ok(Qwen35TickOutcome {
fragment: frag,
is_reasoning: false,
finished,
})
}
pub(crate) fn reset_at_exit(&self, kv_cache: &mut HybridKvCache) -> Result<()> {
kv_cache
.reset_for_slot(self.slot_id)
.context("ADR-040 Phase F M1: Qwen35 reset_for_slot at exit")
}
pub(crate) fn finish(
self,
qwen: &Qwen35LoadedModel,
registration: Option<&ModelRegistration>,
) -> GenerationResult {
let (content_text, reasoning_text) = match registration {
Some(reg) if reg.has_reasoning() => super::registry::split_full_output_forced(
reg,
&self.decoded_text,
self.params.reasoning_forced_open,
),
_ => (self.decoded_text.clone(), None),
};
let reasoning_token_count = match registration {
Some(reg) if reg.has_reasoning() => {
let mut sp = super::registry::make_reasoning_splitter(
reg,
self.params.reasoning_forced_open,
);
let mut count = 0usize;
for &tok in &self.generated_tokens {
let frag = qwen.tokenizer.decode(&[tok], false).unwrap_or_default();
if let Some(splitter) = sp.as_mut() {
let _ = splitter.feed(&frag);
if splitter.in_reasoning() {
count += 1;
}
}
}
count
}
_ => 0,
};
let decode_duration = self.decode_start.elapsed();
GenerationResult {
text: content_text,
reasoning_text,
prompt_tokens: self.prompt_len,
completion_tokens: self.generated_tokens.len(),
reasoning_tokens: if reasoning_token_count > 0 {
Some(reasoning_token_count)
} else {
None
},
finish_reason: self.finish_reason,
prefill_duration: self.prefill_duration,
decode_duration,
cached_tokens: if self.prompt_cache_hit {
self.prompt_len
} else {
0
},
logprobs: self.logprobs_vec,
}
}
}
#[allow(clippy::too_many_arguments)]
pub fn generate_stream_qwen35_once_extended_slot_aware(
qwen: &mut Qwen35LoadedModel,
prompt_tokens: &[u32],
soft_tokens: &[crate::serve::forward_prefill::SoftTokenInjection<'_>],
deepstack: Option<&crate::serve::forward_prefill::DeepstackInjection<'_>>,
positions_flat: Option<&[i32]>,
params: &SamplingParams,
events: &tokio::sync::mpsc::Sender<GenerationEvent>,
registration: Option<&ModelRegistration>,
cancellation_counter: Option<&std::sync::atomic::AtomicU64>,
kv_cache: &mut HybridKvCache,
slot_id: SlotId,
) {
macro_rules! send {
($ev:expr) => {
if events.blocking_send($ev).is_err() {
tracing::info!("SSE stream dropped by client; aborting qwen35 slot-aware decode");
if let Some(c) = cancellation_counter {
c.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
return;
}
};
}
if prompt_tokens.is_empty() {
send!(GenerationEvent::Error(
"generate_stream_qwen35_once_extended_slot_aware: empty prompt_tokens".into()
));
return;
}
if slot_id.0 >= kv_cache.n_seqs {
send!(GenerationEvent::Error(format!(
"capability_unsupported: ADR-040 iter-C2d-cont-kernel iter-2 — \
SlotOutOfRange slot={} max_slots={} (generate_stream_qwen35_\
once_extended_slot_aware)",
slot_id.0, kv_cache.n_seqs,
)));
return;
}
let has_extension = !soft_tokens.is_empty() || deepstack.is_some() || positions_flat.is_some();
if let Some(p) = positions_flat {
if p.len() != 4 * prompt_tokens.len() {
send!(GenerationEvent::Error(format!(
"qwen35 stream slot-aware (iter-4): positions_flat.len() = {} \
!= 4 * prompt_len = {}",
p.len(),
4 * prompt_tokens.len()
)));
return;
}
}
let prompt_len = prompt_tokens.len();
let max_tokens = params.max_tokens.max(1);
let need_seq = prompt_len + max_tokens + 64;
if need_seq > kv_cache.max_seq_len as usize {
send!(GenerationEvent::Error(format!(
"capability_unsupported: ADR-040 iter-C2d-cont-kernel iter-2 — \
per-request need_seq={} exceeds persistent cache \
max_seq_len={} (slot={} prompt_len={} max_tokens={}). \
Persistent cache is sized to cfg.max_position_embeddings; \
reduce max_tokens or use a shorter prompt.",
need_seq, kv_cache.max_seq_len, slot_id.0, prompt_len, max_tokens
)));
return;
}
let is_greedy = is_greedy_eligible(params);
let device = match MlxDevice::new() {
Ok(d) => d,
Err(e) => {
send!(GenerationEvent::Error(format!(
"qwen35 stream slot-aware: MlxDevice::new failed: {e}"
)));
return;
}
};
let _ = &device;
if let Err(e) = kv_cache.reset_for_slot(slot_id) {
send!(GenerationEvent::Error(format!(
"ADR-040 iter-C2d-cont-kernel iter-2: reset_for_slot at entry \
failed: {e:#}"
)));
return;
}
let pre_dispatches = mlx_native::dispatch_count();
let pre_syncs = mlx_native::sync_count();
let prompt_cache_hit =
!has_extension && qwen.prompt_cache.try_match(prompt_tokens, params).is_some();
let prefill_start = Instant::now();
let mut next_token: u32;
if prompt_cache_hit {
let snap = qwen
.prompt_cache
.snapshot()
.expect("try_match Some implies snapshot Some");
if let Err(e) = kv_cache.restore_partial(snap, prompt_len) {
send!(GenerationEvent::Error(format!(
"ADR-040 iter-C2d-cont-kernel iter-2: prompt_cache \
restore_partial failed: {e:#}"
)));
return;
}
next_token = qwen.prompt_cache.first_decoded_token();
tracing::debug!(
"qwen35 stream slot-aware prompt_cache: HIT slot={} prompt_len={} \
prefill skipped",
slot_id.0,
prompt_len
);
} else {
let positions_owned: Vec<i32>;
let positions_slice: &[i32] = match positions_flat {
Some(p) => p,
None => {
positions_owned = prefill_positions_for(prompt_len);
&positions_owned
}
};
let prefill_logits_res: Result<Vec<f32>> = if has_extension {
qwen.model
.forward_gpu_last_logits_with_soft_tokens_and_deepstack(
prompt_tokens,
positions_slice,
soft_tokens,
deepstack,
kv_cache,
slot_id,
)
} else {
qwen.model
.forward_gpu_last_logits(prompt_tokens, positions_slice, kv_cache, slot_id)
};
let prefill_logits = match prefill_logits_res {
Ok(l) => l,
Err(e) => {
send!(GenerationEvent::Error(format!(
"qwen35 stream slot-aware prefill failed: {e:#}"
)));
return;
}
};
if prefill_logits.len() != qwen.vocab_size {
send!(GenerationEvent::Error(format!(
"qwen35 stream slot-aware prefill logits len {} != \
vocab_size {}",
prefill_logits.len(),
qwen.vocab_size,
)));
return;
}
if is_greedy {
next_token = greedy_argmax_last_token(&prefill_logits, qwen.vocab_size as u32);
} else {
let mut logits = prefill_logits;
next_token = sample_logits_qwen35(&mut logits, params, &[]);
}
}
let prefill_duration = prefill_start.elapsed();
let t_post: i32 = match positions_flat {
Some(p) => {
let mut max_t = 0i32;
for i in 0..prompt_len {
let v = p[i]; if v > max_t {
max_t = v;
}
}
max_t.saturating_add(1)
}
None => prompt_len as i32,
};
let mut reasoning_splitter = registration
.and_then(|r| super::registry::make_reasoning_splitter(r, params.reasoning_forced_open));
let mut tool_splitter = registration.and_then(ToolCallSplitter::from_registration);
let mut tool_call_body: String = String::new();
let mut tool_call_index: usize = 0;
let mut saw_tool_call: bool = false;
fn route_content_qwen35_slot_aware(
tool_splitter: &mut Option<ToolCallSplitter>,
body: &mut String,
tc_index: &mut usize,
saw_tc: &mut bool,
registration: Option<&ModelRegistration>,
events: &tokio::sync::mpsc::Sender<GenerationEvent>,
text: &str,
) -> bool {
if text.is_empty() {
return true;
}
let Some(tcs) = tool_splitter.as_mut() else {
return events
.blocking_send(GenerationEvent::Delta {
kind: DeltaKind::Content,
text: text.to_string(),
})
.is_ok();
};
for ev in tcs.feed(text) {
match ev {
ToolCallEvent::Content(t) => {
if !t.is_empty()
&& events
.blocking_send(GenerationEvent::Delta {
kind: DeltaKind::Content,
text: t,
})
.is_err()
{
return false;
}
}
ToolCallEvent::ToolCallOpen => {
body.clear();
}
ToolCallEvent::ToolCallText(t) => {
body.push_str(&t);
}
ToolCallEvent::ToolCallClose => {
let parsed =
registration.and_then(|r| super::registry::parse_tool_call_body(r, body));
let body_dump = std::mem::take(body);
let sink = super::engine::EventSink::new(events);
if super::engine::emit_streaming_tool_call_close(
parsed,
body_dump,
params_tool_call_policy_for_qwen35_stream(),
tc_index,
saw_tc,
&sink,
)
.is_err()
{
return false;
}
}
}
}
true
}
fn emit_fragment_qwen35_slot_aware(
reasoning_splitter: &mut Option<ReasoningSplitter>,
tool_splitter: &mut Option<ToolCallSplitter>,
body: &mut String,
tc_index: &mut usize,
saw_tc: &mut bool,
registration: Option<&ModelRegistration>,
events: &tokio::sync::mpsc::Sender<GenerationEvent>,
fragment: &str,
) -> bool {
if fragment.is_empty() {
return true;
}
if let Some(rs) = reasoning_splitter.as_mut() {
for (slot, text) in rs.feed(fragment) {
match slot {
SplitSlot::Reasoning => {
if !text.is_empty()
&& events
.blocking_send(GenerationEvent::Delta {
kind: DeltaKind::Reasoning,
text,
})
.is_err()
{
return false;
}
}
SplitSlot::Content => {
if !route_content_qwen35_slot_aware(
tool_splitter,
body,
tc_index,
saw_tc,
registration,
events,
&text,
) {
return false;
}
}
}
}
true
} else {
route_content_qwen35_slot_aware(
tool_splitter,
body,
tc_index,
saw_tc,
registration,
events,
fragment,
)
}
}
let decode_start = Instant::now();
let mut completion_tokens = 0usize;
let mut accumulated_text = String::new();
let mut reasoning_token_count = 0usize;
let mut finish_reason: &'static str = "length";
let first_text = qwen
.tokenizer
.decode(&[next_token], false)
.unwrap_or_default();
let mut is_eos_first = qwen.eos_token_ids.contains(&next_token);
if !is_eos_first && !first_text.is_empty() {
accumulated_text.push_str(&first_text);
if !emit_fragment_qwen35_slot_aware(
&mut reasoning_splitter,
&mut tool_splitter,
&mut tool_call_body,
&mut tool_call_index,
&mut saw_tool_call,
registration,
events,
&first_text,
) {
if let Some(c) = cancellation_counter {
c.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
let _ = kv_cache.reset_for_slot(slot_id);
return;
}
}
completion_tokens += 1;
if reasoning_splitter
.as_ref()
.map(|s| s.in_reasoning())
.unwrap_or(false)
{
reasoning_token_count += 1;
}
if is_eos_first {
finish_reason = "stop";
} else if qwen35_hit_stop_string(&accumulated_text, ¶ms.stop_strings) {
finish_reason = "stop";
is_eos_first = true;
}
if !is_eos_first {
for step in 1..max_tokens {
let pos = t_post + (step as i32 - 1);
if pos as u32 >= kv_cache.max_seq_len {
break;
}
let decode_positions = vec![pos; 4];
let dec_result: Result<u32, anyhow::Error> = if is_greedy {
qwen.model
.forward_gpu_greedy(&[next_token], &decode_positions, kv_cache, slot_id)
.map_err(|e| {
anyhow::anyhow!(
"qwen35 stream slot-aware forward_gpu_greedy (ADR-040 \
§6.1.50 iter-G) step {step}: {e}"
)
})
} else {
match qwen.model.forward_gpu_last_logits(
&[next_token],
&decode_positions,
kv_cache,
slot_id,
) {
Ok(logits) => {
if logits.len() != qwen.vocab_size {
Err(anyhow::anyhow!(
"qwen35 stream slot-aware decode logits len {} \
!= vocab_size {}",
logits.len(),
qwen.vocab_size,
))
} else {
let mut tmp = logits;
Ok(sample_logits_qwen35(&mut tmp, params, &[next_token]))
}
}
Err(e) => Err(e),
}
};
next_token = match dec_result {
Ok(t) => t,
Err(e) => {
send!(GenerationEvent::Error(format!(
"qwen35 stream slot-aware decode step {step} failed: {e:#}"
)));
let _ = kv_cache.reset_for_slot(slot_id);
return;
}
};
if qwen.eos_token_ids.contains(&next_token) {
finish_reason = "stop";
break;
}
completion_tokens += 1;
let fragment = qwen
.tokenizer
.decode(&[next_token], false)
.unwrap_or_default();
accumulated_text.push_str(&fragment);
if !emit_fragment_qwen35_slot_aware(
&mut reasoning_splitter,
&mut tool_splitter,
&mut tool_call_body,
&mut tool_call_index,
&mut saw_tool_call,
registration,
events,
&fragment,
) {
if let Some(c) = cancellation_counter {
c.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
let _ = kv_cache.reset_for_slot(slot_id);
return;
}
if reasoning_splitter
.as_ref()
.map(|s| s.in_reasoning())
.unwrap_or(false)
{
reasoning_token_count += 1;
}
if qwen35_hit_stop_string(&accumulated_text, ¶ms.stop_strings) {
finish_reason = "stop";
break;
}
}
}
if let Some(rs) = reasoning_splitter.as_mut() {
if let Some((slot, tail)) = rs.finish() {
match slot {
SplitSlot::Reasoning => {
if !tail.is_empty()
&& events
.blocking_send(GenerationEvent::Delta {
kind: DeltaKind::Reasoning,
text: tail,
})
.is_err()
{
if let Some(c) = cancellation_counter {
c.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
let _ = kv_cache.reset_for_slot(slot_id);
return;
}
}
SplitSlot::Content => {
if !route_content_qwen35_slot_aware(
&mut tool_splitter,
&mut tool_call_body,
&mut tool_call_index,
&mut saw_tool_call,
registration,
events,
&tail,
) {
if let Some(c) = cancellation_counter {
c.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
let _ = kv_cache.reset_for_slot(slot_id);
return;
}
}
}
}
}
if let Some(tcs) = tool_splitter.as_mut() {
if let Some(ev) = tcs.finish() {
match ev {
ToolCallEvent::Content(t) => {
if !t.is_empty()
&& events
.blocking_send(GenerationEvent::Delta {
kind: DeltaKind::Content,
text: t,
})
.is_err()
{
if let Some(c) = cancellation_counter {
c.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
let _ = kv_cache.reset_for_slot(slot_id);
return;
}
}
ToolCallEvent::ToolCallText(t) => {
let prefix = registration.and_then(|r| r.tool_open).unwrap_or("");
let fallback = format!("{prefix}{t}");
if !fallback.is_empty()
&& events
.blocking_send(GenerationEvent::Delta {
kind: DeltaKind::Content,
text: fallback,
})
.is_err()
{
if let Some(c) = cancellation_counter {
c.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
let _ = kv_cache.reset_for_slot(slot_id);
return;
}
}
ToolCallEvent::ToolCallOpen | ToolCallEvent::ToolCallClose => {
}
}
}
}
if saw_tool_call {
finish_reason = "tool_calls";
}
let decode_duration = decode_start.elapsed();
if let Err(e) = kv_cache.reset_for_slot(slot_id) {
send!(GenerationEvent::Error(format!(
"ADR-040 iter-C2d-cont-kernel iter-2: reset_for_slot at exit \
failed: {e:#}"
)));
return;
}
let stats = StreamStats {
prefill_time_secs: Some(prefill_duration.as_secs_f64()),
decode_time_secs: Some(decode_duration.as_secs_f64()),
total_time_secs: Some((prefill_duration + decode_duration).as_secs_f64()),
time_to_first_token_ms: Some(prefill_duration.as_secs_f64() * 1000.0),
prefill_tokens_per_sec: Some(if prefill_duration.as_secs_f64() > 0.0 {
prompt_len as f64 / prefill_duration.as_secs_f64()
} else {
0.0
}),
decode_tokens_per_sec: Some(if decode_duration.as_secs_f64() > 0.0 {
completion_tokens as f64 / decode_duration.as_secs_f64()
} else {
0.0
}),
gpu_sync_count: Some(mlx_native::sync_count().saturating_sub(pre_syncs)),
gpu_dispatch_count: Some(mlx_native::dispatch_count().saturating_sub(pre_dispatches)),
cached_prompt_tokens: if prompt_cache_hit {
Some(prompt_len)
} else {
None
},
reasoning_tokens: if reasoning_token_count > 0 {
Some(reasoning_token_count)
} else {
None
},
};
send!(GenerationEvent::Done {
finish_reason,
prompt_tokens: prompt_len,
completion_tokens,
stats,
});
}
pub fn generate_qwen35_once_with_soft_tokens(
qwen: &mut Qwen35LoadedModel,
prompt_tokens: &[u32],
soft_tokens: &[crate::serve::forward_prefill::SoftTokenInjection<'_>],
params: &SamplingParams,
registration: Option<&ModelRegistration>,
) -> Result<GenerationResult> {
if soft_tokens.is_empty() {
return generate_qwen35_once(qwen, prompt_tokens, params, registration);
}
anyhow::ensure!(
!prompt_tokens.is_empty(),
"generate_qwen35_once_with_soft_tokens: empty prompt_tokens"
);
let prompt_len = prompt_tokens.len();
let max_tokens = params.max_tokens.max(1);
let is_greedy = is_greedy_eligible(params);
let want_logprobs = params.logprobs;
let mut logprobs_vec: Option<Vec<f32>> = if want_logprobs {
Some(Vec::with_capacity(max_tokens))
} else {
None
};
let device = MlxDevice::new()
.map_err(|e| anyhow::anyhow!("MlxDevice::new (qwen35 generate w/ soft tokens): {e}"))?;
let mut kv_cache = alloc_kv_cache_for_request(qwen, &device, prompt_len, max_tokens)?;
qwen.hydrate_lcp_registry_from_disk(&kv_cache, &device);
let prefill_start = Instant::now();
let positions = prefill_positions_for(prompt_len);
let prefill_logits = qwen
.model
.forward_gpu_last_logits_with_soft_tokens(
prompt_tokens,
&positions,
soft_tokens,
&mut kv_cache,
SlotId(0),
)
.context("Qwen35Model::forward_gpu_last_logits_with_soft_tokens (prefill)")?;
anyhow::ensure!(
prefill_logits.len() == qwen.vocab_size,
"qwen35 prefill (soft tokens) logits len {} != vocab_size {}",
prefill_logits.len(),
qwen.vocab_size
);
let mut next_token: u32 = if want_logprobs {
let mut logits = prefill_logits.clone();
let (tok, lp) = sample_logits_qwen35_with_logprob(&mut logits, params, &[]);
if let Some(v) = logprobs_vec.as_mut() {
v.push(lp);
}
tok
} else if is_greedy {
greedy_argmax_last_token(&prefill_logits, qwen.vocab_size as u32)
} else {
let mut logits = prefill_logits.clone();
sample_logits_qwen35(&mut logits, params, &[])
};
let prefill_duration = prefill_start.elapsed();
let decode_start = Instant::now();
let mut generated_tokens: Vec<u32> = Vec::with_capacity(max_tokens);
generated_tokens.push(next_token);
let first_fragment = qwen
.tokenizer
.decode(&[next_token], false)
.unwrap_or_default();
let mut decoded_text = first_fragment.clone();
let mut finish_reason: &'static str = "length";
if qwen.eos_token_ids.contains(&next_token) {
finish_reason = "stop";
} else if qwen35_hit_stop_string(&decoded_text, ¶ms.stop_strings) {
finish_reason = "stop";
qwen35_strip_trailing_stop(&mut decoded_text, ¶ms.stop_strings);
} else {
for step in 1..max_tokens {
let pos = (prompt_len + step - 1) as i32;
if pos as u32 >= kv_cache.max_seq_len {
tracing::warn!(
pos,
max_seq = kv_cache.max_seq_len,
"qwen35 decode (soft tokens): hit kv-cache bound; stopping with finish=length",
);
break;
}
let decode_positions = vec![pos; 4];
next_token = if want_logprobs {
let logits_full = qwen
.model
.forward_gpu_last_logits(
&[next_token],
&decode_positions,
&mut kv_cache,
SlotId(0),
)
.with_context(|| {
format!(
"forward_gpu_last_logits decode step {step} (soft tokens, logprobs)"
)
})?;
let mut logits = logits_full;
let (tok, lp) =
sample_logits_qwen35_with_logprob(&mut logits, params, &generated_tokens);
if let Some(v) = logprobs_vec.as_mut() {
v.push(lp);
}
tok
} else if is_greedy {
qwen.model
.forward_gpu_greedy(&[next_token], &decode_positions, &mut kv_cache, SlotId(0))
.with_context(|| {
format!("forward_gpu_greedy decode step {step} (soft tokens)")
})?
} else {
let logits_full = qwen
.model
.forward_gpu_last_logits(
&[next_token],
&decode_positions,
&mut kv_cache,
SlotId(0),
)
.with_context(|| {
format!("forward_gpu_last_logits decode step {step} (soft tokens)")
})?;
let mut logits = logits_full;
sample_logits_qwen35(&mut logits, params, &generated_tokens)
};
if qwen.eos_token_ids.contains(&next_token) {
finish_reason = "stop";
break;
}
generated_tokens.push(next_token);
let fragment = qwen
.tokenizer
.decode(&[next_token], false)
.unwrap_or_default();
decoded_text.push_str(&fragment);
if qwen35_hit_stop_string(&decoded_text, ¶ms.stop_strings) {
finish_reason = "stop";
qwen35_strip_trailing_stop(&mut decoded_text, ¶ms.stop_strings);
break;
}
}
}
let decode_duration = decode_start.elapsed();
let (content, reasoning_text) = match registration {
Some(reg) if reg.has_reasoning() => super::registry::split_full_output_forced(
reg,
&decoded_text,
params.reasoning_forced_open,
),
_ => (decoded_text, None),
};
let reasoning_token_count = match registration {
Some(reg) if reg.has_reasoning() => {
let mut sp =
super::registry::make_reasoning_splitter(reg, params.reasoning_forced_open);
let mut count = 0usize;
for &tok in &generated_tokens {
let frag = qwen.tokenizer.decode(&[tok], false).unwrap_or_default();
if let Some(splitter) = sp.as_mut() {
let _ = splitter.feed(&frag);
if splitter.in_reasoning() {
count += 1;
}
}
}
count
}
_ => 0,
};
Ok(GenerationResult {
text: content,
reasoning_text,
prompt_tokens: prompt_len,
completion_tokens: generated_tokens.len(),
reasoning_tokens: if reasoning_token_count > 0 {
Some(reasoning_token_count)
} else {
None
},
finish_reason,
prefill_duration,
decode_duration,
cached_tokens: 0,
logprobs: logprobs_vec,
})
}
pub fn generate_qwen35_once_with_soft_tokens_and_deepstack(
qwen: &mut Qwen35LoadedModel,
prompt_tokens: &[u32],
soft_tokens: &[crate::serve::forward_prefill::SoftTokenInjection<'_>],
deepstack: Option<&crate::serve::forward_prefill::DeepstackInjection<'_>>,
positions_flat: Option<&[i32]>,
params: &SamplingParams,
registration: Option<&ModelRegistration>,
) -> Result<GenerationResult> {
if soft_tokens.is_empty() && deepstack.is_none() && positions_flat.is_none() {
return generate_qwen35_once(qwen, prompt_tokens, params, registration);
}
anyhow::ensure!(
!prompt_tokens.is_empty(),
"generate_qwen35_once_with_soft_tokens_and_deepstack: empty prompt_tokens"
);
let prompt_len = prompt_tokens.len();
let max_tokens = params.max_tokens.max(1);
let is_greedy = is_greedy_eligible(params);
let want_logprobs = params.logprobs;
let mut logprobs_vec: Option<Vec<f32>> = if want_logprobs {
Some(Vec::with_capacity(max_tokens))
} else {
None
};
let device = MlxDevice::new()
.map_err(|e| anyhow::anyhow!("MlxDevice::new (qwen35 wedge-4d generate): {e}"))?;
let mut kv_cache = alloc_kv_cache_for_request(qwen, &device, prompt_len, max_tokens)?;
qwen.hydrate_lcp_registry_from_disk(&kv_cache, &device);
let prefill_start = Instant::now();
let positions_owned: Vec<i32>;
let positions: &[i32] = match positions_flat {
Some(p) => {
anyhow::ensure!(
p.len() == 4 * prompt_len,
"generate_qwen35_once_with_soft_tokens_and_deepstack: \
positions_flat.len() = {} != 4 * prompt_len = {}",
p.len(),
4 * prompt_len
);
p
}
None => {
positions_owned = prefill_positions_for(prompt_len);
&positions_owned
}
};
let prefill_logits = qwen
.model
.forward_gpu_last_logits_with_soft_tokens_and_deepstack(
prompt_tokens,
positions,
soft_tokens,
deepstack,
&mut kv_cache,
SlotId(0),
)
.context(
"Qwen35Model::forward_gpu_last_logits_with_soft_tokens_and_deepstack \
(prefill)",
)?;
anyhow::ensure!(
prefill_logits.len() == qwen.vocab_size,
"qwen35 prefill (wedge-4d) logits len {} != vocab_size {}",
prefill_logits.len(),
qwen.vocab_size
);
let mut next_token: u32 = if want_logprobs {
let mut logits = prefill_logits.clone();
let (tok, lp) = sample_logits_qwen35_with_logprob(&mut logits, params, &[]);
if let Some(v) = logprobs_vec.as_mut() {
v.push(lp);
}
tok
} else if is_greedy {
greedy_argmax_last_token(&prefill_logits, qwen.vocab_size as u32)
} else {
let mut logits = prefill_logits.clone();
sample_logits_qwen35(&mut logits, params, &[])
};
let prefill_duration = prefill_start.elapsed();
let t_post: i32 = match positions_flat {
Some(p) => {
let mut max_t = 0i32;
for i in 0..prompt_len {
let v = p[i]; if v > max_t {
max_t = v;
}
}
max_t.saturating_add(1)
}
None => prompt_len as i32,
};
let decode_start = Instant::now();
let mut generated_tokens: Vec<u32> = Vec::with_capacity(max_tokens);
generated_tokens.push(next_token);
let first_fragment = qwen
.tokenizer
.decode(&[next_token], false)
.unwrap_or_default();
let mut decoded_text = first_fragment.clone();
let mut finish_reason: &'static str = "length";
if qwen.eos_token_ids.contains(&next_token) {
finish_reason = "stop";
} else if qwen35_hit_stop_string(&decoded_text, ¶ms.stop_strings) {
finish_reason = "stop";
qwen35_strip_trailing_stop(&mut decoded_text, ¶ms.stop_strings);
} else {
for step in 1..max_tokens {
let pos = t_post + (step as i32 - 1);
if pos as u32 >= kv_cache.max_seq_len {
tracing::warn!(
pos,
max_seq = kv_cache.max_seq_len,
"qwen35 decode (wedge-4d): hit kv-cache bound; stopping with finish=length",
);
break;
}
let decode_positions = vec![pos; 4];
next_token = if want_logprobs {
let logits_full = qwen
.model
.forward_gpu_last_logits(
&[next_token],
&decode_positions,
&mut kv_cache,
SlotId(0),
)
.with_context(|| {
format!("forward_gpu_last_logits decode step {step} (wedge-4d, logprobs)")
})?;
let mut logits = logits_full;
let (tok, lp) =
sample_logits_qwen35_with_logprob(&mut logits, params, &generated_tokens);
if let Some(v) = logprobs_vec.as_mut() {
v.push(lp);
}
tok
} else if is_greedy {
qwen.model
.forward_gpu_greedy(&[next_token], &decode_positions, &mut kv_cache, SlotId(0))
.with_context(|| format!("forward_gpu_greedy decode step {step} (wedge-4d)"))?
} else {
let logits_full = qwen
.model
.forward_gpu_last_logits(
&[next_token],
&decode_positions,
&mut kv_cache,
SlotId(0),
)
.with_context(|| {
format!("forward_gpu_last_logits decode step {step} (wedge-4d)")
})?;
let mut logits = logits_full;
sample_logits_qwen35(&mut logits, params, &generated_tokens)
};
if qwen.eos_token_ids.contains(&next_token) {
finish_reason = "stop";
break;
}
generated_tokens.push(next_token);
let fragment = qwen
.tokenizer
.decode(&[next_token], false)
.unwrap_or_default();
decoded_text.push_str(&fragment);
if qwen35_hit_stop_string(&decoded_text, ¶ms.stop_strings) {
finish_reason = "stop";
qwen35_strip_trailing_stop(&mut decoded_text, ¶ms.stop_strings);
break;
}
}
}
let decode_duration = decode_start.elapsed();
let (content, reasoning_text) = match registration {
Some(reg) if reg.has_reasoning() => super::registry::split_full_output_forced(
reg,
&decoded_text,
params.reasoning_forced_open,
),
_ => (decoded_text, None),
};
let reasoning_token_count = match registration {
Some(reg) if reg.has_reasoning() => {
let mut sp =
super::registry::make_reasoning_splitter(reg, params.reasoning_forced_open);
let mut count = 0usize;
for &tok in &generated_tokens {
let frag = qwen.tokenizer.decode(&[tok], false).unwrap_or_default();
if let Some(splitter) = sp.as_mut() {
let _ = splitter.feed(&frag);
if splitter.in_reasoning() {
count += 1;
}
}
}
count
}
_ => 0,
};
Ok(GenerationResult {
text: content,
reasoning_text,
prompt_tokens: prompt_len,
completion_tokens: generated_tokens.len(),
reasoning_tokens: if reasoning_token_count > 0 {
Some(reasoning_token_count)
} else {
None
},
finish_reason,
prefill_duration,
decode_duration,
cached_tokens: 0,
logprobs: logprobs_vec,
})
}
pub fn generate_stream_qwen35_once(
qwen: &mut Qwen35LoadedModel,
prompt_tokens: &[u32],
params: &SamplingParams,
events: &tokio::sync::mpsc::Sender<GenerationEvent>,
registration: Option<&ModelRegistration>,
cancellation_counter: Option<&std::sync::atomic::AtomicU64>,
) {
generate_stream_qwen35_once_extended(
qwen,
prompt_tokens,
&[],
None,
None,
params,
events,
registration,
cancellation_counter,
)
}
#[allow(clippy::too_many_arguments)]
pub fn generate_stream_qwen35_once_extended(
qwen: &mut Qwen35LoadedModel,
prompt_tokens: &[u32],
soft_tokens: &[crate::serve::forward_prefill::SoftTokenInjection<'_>],
deepstack: Option<&crate::serve::forward_prefill::DeepstackInjection<'_>>,
positions_flat: Option<&[i32]>,
params: &SamplingParams,
events: &tokio::sync::mpsc::Sender<GenerationEvent>,
registration: Option<&ModelRegistration>,
cancellation_counter: Option<&std::sync::atomic::AtomicU64>,
) {
macro_rules! send {
($ev:expr) => {
if events.blocking_send($ev).is_err() {
tracing::info!("SSE stream dropped by client; aborting qwen35 decode");
if let Some(c) = cancellation_counter {
c.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
return;
}
};
}
if prompt_tokens.is_empty() {
send!(GenerationEvent::Error(
"generate_stream_qwen35_once: empty prompt_tokens".into()
));
return;
}
let prompt_len = prompt_tokens.len();
let max_tokens = params.max_tokens.max(1);
let is_greedy = is_greedy_eligible(params);
let has_extension = !soft_tokens.is_empty() || deepstack.is_some() || positions_flat.is_some();
if let Some(p) = positions_flat {
if p.len() != 4 * prompt_len {
send!(GenerationEvent::Error(format!(
"qwen35 stream (wedge-4e): positions_flat.len() = {} != 4 * prompt_len = {}",
p.len(),
4 * prompt_len
)));
return;
}
}
let device = match MlxDevice::new() {
Ok(d) => d,
Err(e) => {
send!(GenerationEvent::Error(format!(
"qwen35 stream: MlxDevice::new failed: {e}"
)));
return;
}
};
let mut kv_cache = match alloc_kv_cache_for_request(qwen, &device, prompt_len, max_tokens) {
Ok(k) => k,
Err(e) => {
send!(GenerationEvent::Error(format!(
"qwen35 stream: KV cache alloc failed: {e:#}"
)));
return;
}
};
qwen.hydrate_lcp_registry_from_disk(&kv_cache, &device);
let pre_dispatches = mlx_native::dispatch_count();
let pre_syncs = mlx_native::sync_count();
let prompt_cache_hit =
!has_extension && qwen.prompt_cache.try_match(prompt_tokens, params).is_some();
let mut lcp_resume_start: usize = 0;
if !prompt_cache_hit && !has_extension {
let stride_for_observe = crate::debug::INVESTIGATION_ENV.kv_lcp_deltanet_checkpoint_stride;
let base_key_for_observe = build_lcp_key_for_qwen35(qwen, params);
let detected = crate::serve::kv_persist::lcp_registry::probe_lcp_opportunity_chunk_aligned(
&mut qwen.lcp_registry,
prompt_tokens,
stride_for_observe,
false,
|chunk_pos| {
let mut key = base_key_for_observe.clone();
if chunk_pos > 0 {
key.tenant_id = format!("qwen35:lcp_chunk:{chunk_pos}");
}
key
},
);
if let Some(sink) = qwen.kv_metrics_sink.as_ref() {
sink.record_lcp_probe(detected);
}
let _ = detected;
let lcp_resume_enabled = crate::serve::api::engine::effective_kv_lcp_resume(
crate::debug::INVESTIGATION_ENV.kv_lcp_resume,
true,
);
if lcp_resume_enabled {
let stride = crate::debug::INVESTIGATION_ENV.kv_lcp_deltanet_checkpoint_stride;
let max_chunk_pos = if stride == 0 {
eprintln!(
"[hf2q qwen35 stream lcp probe] stride=0; skipping \
descending scan (HF2Q_KV_LCP_DELTANET_CHECKPOINT_STRIDE \
must be > 0)"
);
0
} else {
(prompt_tokens.len() / stride).saturating_mul(stride)
};
eprintln!(
"[hf2q qwen35 stream lcp probe] enabled, registry_len={}, \
prompt_len={}, stride={}, scanning [{stride}..={max_chunk_pos}]",
qwen.lcp_registry.len(),
prompt_tokens.len(),
stride,
);
let mut hit = false;
let mut chunk_pos = max_chunk_pos;
while stride > 0 && chunk_pos >= stride && !hit {
let chunk_key = build_lcp_key_for_qwen35_chunk(qwen, params, chunk_pos);
if let Some(prefix) = qwen.lcp_registry.lookup(&chunk_key, prompt_tokens) {
if prefix.k == prefix.cached_prompt_len && prefix.k < prompt_tokens.len() {
let snapshot: &HybridKvCacheSnapshot = &prefix.dense_kvs[0];
let restore_start = Instant::now();
if let Err(e) = kv_cache.restore_partial(snapshot, prefix.k) {
send!(GenerationEvent::Error(format!(
"qwen35 stream: lcp_registry restore_partial failed: {e:#}"
)));
return;
}
let restore_ms = restore_start.elapsed().as_micros() as f64 / 1000.0;
lcp_resume_start = prefix.k;
eprintln!(
"[hf2q qwen35 stream lcp resume] STRIDE-ALIGNED HIT \
— restoring at k={} (cached_prompt_len={}, chunk_pos={}, \
restore_ms={:.3})",
prefix.k, prefix.cached_prompt_len, chunk_pos, restore_ms
);
hit = true;
} else {
eprintln!(
"[hf2q qwen35 stream lcp probe] PARTIAL HIT at \
chunk_pos={} — k={} < cached_prompt_len={}",
chunk_pos, prefix.k, prefix.cached_prompt_len
);
}
}
if chunk_pos == stride {
break;
}
chunk_pos -= stride;
}
if !hit {
eprintln!(
"[hf2q qwen35 stream lcp probe] no stride-aligned match \
(registry_len={})",
qwen.lcp_registry.len()
);
}
}
}
let prefill_start = Instant::now();
let mut next_token: u32;
if prompt_cache_hit {
let snap = qwen
.prompt_cache
.snapshot()
.expect("try_match Some implies snapshot Some");
if let Err(e) = kv_cache.restore_from(snap) {
send!(GenerationEvent::Error(format!(
"qwen35 stream: prompt_cache restore failed: {e:#}"
)));
return;
}
next_token = qwen.prompt_cache.first_decoded_token();
tracing::debug!(
"qwen35 prompt_cache: STREAMING HIT — {} tokens; prefill skipped",
prompt_len
);
} else {
let positions_owned: Vec<i32>;
let positions_slice: &[i32] = match positions_flat {
Some(p) => p,
None => {
positions_owned = prefill_positions_for(prompt_len);
&positions_owned
}
};
let stride = crate::debug::INVESTIGATION_ENV.kv_lcp_deltanet_checkpoint_stride;
let lcp_resume_enabled = crate::debug::INVESTIGATION_ENV.kv_lcp_resume;
let chunked_eligible = !has_extension
&& stride > 0
&& prompt_len > stride
&& crate::debug::INVESTIGATION_ENV.kv_lcp_chunked_prefill;
if lcp_resume_start > 0 && stride > 0 && lcp_resume_start % stride != 0 {
send!(GenerationEvent::Error(format!(
"qwen35 stream: lcp_resume_start ({}) must be stride-aligned \
({}) — registry should only contain stride-aligned chunk-position \
keys",
lcp_resume_start, stride
)));
return;
}
let prefill_logits_res = if has_extension {
qwen.model
.forward_gpu_last_logits_with_soft_tokens_and_deepstack(
prompt_tokens,
positions_slice,
soft_tokens,
deepstack,
&mut kv_cache,
SlotId(0),
)
} else if chunked_eligible {
let first_chunk_idx = lcp_resume_start / stride;
let n_chunks = (prompt_len + stride - 1) / stride;
eprintln!(
"[hf2q qwen35 stream chunked prefill] {} chunks (stride={}, \
prompt_len={}, first_chunk_idx={})",
n_chunks, stride, prompt_len, first_chunk_idx
);
let mut last_logits_res: Result<Vec<f32>> = Err(anyhow::anyhow!("no chunks executed"));
for chunk_idx in first_chunk_idx..n_chunks {
let k_start = chunk_idx * stride;
let k_end = ((chunk_idx + 1) * stride).min(prompt_len);
let chunk_seq_len = k_end - k_start;
let chunk_tokens = &prompt_tokens[k_start..k_end];
let mut chunk_positions = vec![0i32; 4 * chunk_seq_len];
for axis in 0..4 {
for t in 0..chunk_seq_len {
chunk_positions[axis * chunk_seq_len + t] = (k_start + t) as i32;
}
}
let res = qwen.model.forward_gpu_last_logits(
chunk_tokens,
&chunk_positions,
&mut kv_cache,
SlotId(0),
);
let logits = match res {
Ok(l) => l,
Err(e) => {
send!(GenerationEvent::Error(format!(
"qwen35 stream chunked prefill chunk {}/{} failed: {e:#}",
chunk_idx + 1,
n_chunks
)));
return;
}
};
if chunk_idx == n_chunks - 1 {
last_logits_res = Ok(logits.clone());
}
let stride_aligned = k_end % stride == 0;
let mid_store_disabled =
std::env::var("HF2Q_KV_LCP_DISABLE_MID_STORE").as_deref() == Ok("1");
lcp_store_skip_notify(stride_aligned, lcp_resume_enabled, mid_store_disabled);
if lcp_resume_enabled && stride_aligned && !mid_store_disabled {
match kv_cache.snapshot(&device) {
Ok(snap) => {
let chunk_key = build_lcp_key_for_qwen35_chunk(qwen, params, k_end);
let linear_capacity = kv_cache
.linear_attn
.first()
.map(|s| s.recurrent.byte_len())
.unwrap_or(0);
if let Err(e) = qwen.store_lcp_with_disk_writeback(
&kv_cache,
chunk_key,
prompt_tokens[..k_end].to_vec(),
snap,
0,
linear_capacity,
) {
lcp_store_error_notify("stream mid-prefill", k_end, &e);
} else {
eprintln!(
"[hf2q qwen35 stream lcp store] mid-prefill \
snapshot at chunk_pos={k_end} \
(registry_len_after={})",
qwen.lcp_registry.len()
);
}
}
Err(e) => {
lcp_snapshot_error_notify("stream mid-prefill", k_end, &e);
}
}
}
}
last_logits_res
} else if lcp_resume_start > 0 {
let suffix_tokens = &prompt_tokens[lcp_resume_start..];
let suffix_len = suffix_tokens.len();
let mut suffix_positions = vec![0i32; 4 * suffix_len];
for axis in 0..4 {
for t in 0..suffix_len {
suffix_positions[axis * suffix_len + t] = (lcp_resume_start + t) as i32;
}
}
eprintln!(
"[hf2q qwen35 stream lcp resume] suffix prefill {} tokens \
(lcp_resume_start={}, prompt_len={})",
suffix_len, lcp_resume_start, prompt_len
);
qwen.model.forward_gpu_last_logits(
suffix_tokens,
&suffix_positions,
&mut kv_cache,
SlotId(0),
)
} else {
qwen.model.forward_gpu_last_logits(
prompt_tokens,
positions_slice,
&mut kv_cache,
SlotId(0),
)
};
let prefill_logits = match prefill_logits_res {
Ok(l) => l,
Err(e) => {
send!(GenerationEvent::Error(format!(
"qwen35 stream prefill failed: {e:#}"
)));
return;
}
};
if is_greedy {
next_token = greedy_argmax_last_token(&prefill_logits, qwen.vocab_size as u32);
} else {
let mut logits = prefill_logits.clone();
next_token = sample_logits_qwen35(&mut logits, params, &[]);
}
if is_greedy && !has_extension {
match kv_cache.snapshot(&device) {
Ok(snap) => {
qwen.prompt_cache
.update(prompt_tokens.to_vec(), snap, next_token, params)
}
Err(e) => {
eprintln!("[hf2q qwen35 stream lcp store] prompt_cache snapshot failed: {e}")
}
}
}
if !has_extension && lcp_resume_enabled {
let chunked_already_stored = stride > 0
&& prompt_len > stride
&& (crate::debug::INVESTIGATION_ENV.kv_lcp_chunked_prefill || lcp_resume_start > 0);
let prompt_stride_aligned = stride > 0 && prompt_len > 0 && prompt_len % stride == 0;
if !chunked_already_stored && prompt_stride_aligned {
match kv_cache.snapshot(&device) {
Ok(snap) => {
let chunk_key = build_lcp_key_for_qwen35_chunk(qwen, params, prompt_len);
let linear_capacity = kv_cache
.linear_attn
.first()
.map(|s| s.recurrent.byte_len())
.unwrap_or(0);
if let Err(e) = qwen.store_lcp_with_disk_writeback(
&kv_cache,
chunk_key,
prompt_tokens.to_vec(),
snap,
0,
linear_capacity,
) {
lcp_store_error_notify("stream end-of-prefill", prompt_len, &e);
}
}
Err(e) => {
lcp_snapshot_error_notify("stream end-of-prefill", prompt_len, &e);
}
}
}
}
}
let prefill_duration = prefill_start.elapsed();
let t_post: i32 = match positions_flat {
Some(p) => {
let mut max_t = 0i32;
for i in 0..prompt_len {
let v = p[i]; if v > max_t {
max_t = v;
}
}
max_t.saturating_add(1)
}
None => prompt_len as i32,
};
let mut reasoning_splitter = registration
.and_then(|r| super::registry::make_reasoning_splitter(r, params.reasoning_forced_open));
let mut tool_splitter = registration.and_then(ToolCallSplitter::from_registration);
let mut tool_call_body: String = String::new();
let mut tool_call_index: usize = 0;
let mut saw_tool_call: bool = false;
fn route_content_qwen35(
tool_splitter: &mut Option<ToolCallSplitter>,
body: &mut String,
tc_index: &mut usize,
saw_tc: &mut bool,
registration: Option<&ModelRegistration>,
events: &tokio::sync::mpsc::Sender<GenerationEvent>,
text: &str,
) -> bool {
if text.is_empty() {
return true;
}
let Some(tcs) = tool_splitter.as_mut() else {
return events
.blocking_send(GenerationEvent::Delta {
kind: DeltaKind::Content,
text: text.to_string(),
})
.is_ok();
};
for ev in tcs.feed(text) {
match ev {
ToolCallEvent::Content(t) => {
if !t.is_empty()
&& events
.blocking_send(GenerationEvent::Delta {
kind: DeltaKind::Content,
text: t,
})
.is_err()
{
return false;
}
}
ToolCallEvent::ToolCallOpen => {
body.clear();
}
ToolCallEvent::ToolCallText(t) => {
body.push_str(&t);
}
ToolCallEvent::ToolCallClose => {
let parsed =
registration.and_then(|r| super::registry::parse_tool_call_body(r, body));
let body_dump = std::mem::take(body);
let sink = super::engine::EventSink::new(events);
if super::engine::emit_streaming_tool_call_close(
parsed,
body_dump,
params_tool_call_policy_for_qwen35_stream(),
tc_index,
saw_tc,
&sink,
)
.is_err()
{
return false;
}
}
}
}
true
}
fn emit_fragment_qwen35(
reasoning_splitter: &mut Option<ReasoningSplitter>,
tool_splitter: &mut Option<ToolCallSplitter>,
body: &mut String,
tc_index: &mut usize,
saw_tc: &mut bool,
registration: Option<&ModelRegistration>,
events: &tokio::sync::mpsc::Sender<GenerationEvent>,
fragment: &str,
) -> bool {
if fragment.is_empty() {
return true;
}
if let Some(rs) = reasoning_splitter.as_mut() {
for (slot, text) in rs.feed(fragment) {
match slot {
SplitSlot::Reasoning => {
if !text.is_empty()
&& events
.blocking_send(GenerationEvent::Delta {
kind: DeltaKind::Reasoning,
text,
})
.is_err()
{
return false;
}
}
SplitSlot::Content => {
if !route_content_qwen35(
tool_splitter,
body,
tc_index,
saw_tc,
registration,
events,
&text,
) {
return false;
}
}
}
}
true
} else {
route_content_qwen35(
tool_splitter,
body,
tc_index,
saw_tc,
registration,
events,
fragment,
)
}
}
let decode_start = Instant::now();
let mut completion_tokens = 0usize;
let mut accumulated_text = String::new();
let mut reasoning_token_count = 0usize;
let mut finish_reason: &'static str = "length";
let first_text = qwen
.tokenizer
.decode(&[next_token], false)
.unwrap_or_default();
let mut is_eos_first = qwen.eos_token_ids.contains(&next_token);
if !is_eos_first && !first_text.is_empty() {
accumulated_text.push_str(&first_text);
if !emit_fragment_qwen35(
&mut reasoning_splitter,
&mut tool_splitter,
&mut tool_call_body,
&mut tool_call_index,
&mut saw_tool_call,
registration,
events,
&first_text,
) {
if let Some(c) = cancellation_counter {
c.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
return;
}
}
completion_tokens += 1;
if reasoning_splitter
.as_ref()
.map(|s| s.in_reasoning())
.unwrap_or(false)
{
reasoning_token_count += 1;
}
if is_eos_first {
finish_reason = "stop";
} else if qwen35_hit_stop_string(&accumulated_text, ¶ms.stop_strings) {
finish_reason = "stop";
is_eos_first = true;
}
if !is_eos_first {
for step in 1..max_tokens {
let pos = t_post + (step as i32 - 1);
if pos as u32 >= kv_cache.max_seq_len {
break;
}
let decode_positions = vec![pos; 4];
let dec_result = if is_greedy {
qwen.model
.forward_gpu_greedy(&[next_token], &decode_positions, &mut kv_cache, SlotId(0))
} else {
match qwen.model.forward_gpu_last_logits(
&[next_token],
&decode_positions,
&mut kv_cache,
SlotId(0),
) {
Ok(logits) => {
let mut tmp = logits;
Ok(sample_logits_qwen35(&mut tmp, params, &[next_token]))
}
Err(e) => Err(e),
}
};
next_token = match dec_result {
Ok(t) => t,
Err(e) => {
send!(GenerationEvent::Error(format!(
"qwen35 stream decode step {step} failed: {e:#}"
)));
return;
}
};
if qwen.eos_token_ids.contains(&next_token) {
finish_reason = "stop";
break;
}
completion_tokens += 1;
let fragment = qwen
.tokenizer
.decode(&[next_token], false)
.unwrap_or_default();
accumulated_text.push_str(&fragment);
if !emit_fragment_qwen35(
&mut reasoning_splitter,
&mut tool_splitter,
&mut tool_call_body,
&mut tool_call_index,
&mut saw_tool_call,
registration,
events,
&fragment,
) {
if let Some(c) = cancellation_counter {
c.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
return;
}
if reasoning_splitter
.as_ref()
.map(|s| s.in_reasoning())
.unwrap_or(false)
{
reasoning_token_count += 1;
}
if qwen35_hit_stop_string(&accumulated_text, ¶ms.stop_strings) {
finish_reason = "stop";
break;
}
}
}
if let Some(rs) = reasoning_splitter.as_mut() {
if let Some((slot, tail)) = rs.finish() {
match slot {
SplitSlot::Reasoning => {
if !tail.is_empty()
&& events
.blocking_send(GenerationEvent::Delta {
kind: DeltaKind::Reasoning,
text: tail,
})
.is_err()
{
if let Some(c) = cancellation_counter {
c.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
return;
}
}
SplitSlot::Content => {
if !route_content_qwen35(
&mut tool_splitter,
&mut tool_call_body,
&mut tool_call_index,
&mut saw_tool_call,
registration,
events,
&tail,
) {
if let Some(c) = cancellation_counter {
c.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
return;
}
}
}
}
}
if let Some(tcs) = tool_splitter.as_mut() {
if let Some(ev) = tcs.finish() {
match ev {
ToolCallEvent::Content(t) => {
if !t.is_empty()
&& events
.blocking_send(GenerationEvent::Delta {
kind: DeltaKind::Content,
text: t,
})
.is_err()
{
if let Some(c) = cancellation_counter {
c.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
return;
}
}
ToolCallEvent::ToolCallText(t) => {
let prefix = registration.and_then(|r| r.tool_open).unwrap_or("");
let fallback = format!("{prefix}{t}");
if !fallback.is_empty()
&& events
.blocking_send(GenerationEvent::Delta {
kind: DeltaKind::Content,
text: fallback,
})
.is_err()
{
if let Some(c) = cancellation_counter {
c.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
return;
}
}
ToolCallEvent::ToolCallOpen | ToolCallEvent::ToolCallClose => {
}
}
}
}
if saw_tool_call {
finish_reason = "tool_calls";
}
let decode_duration = decode_start.elapsed();
let stats = StreamStats {
prefill_time_secs: Some(prefill_duration.as_secs_f64()),
decode_time_secs: Some(decode_duration.as_secs_f64()),
total_time_secs: Some((prefill_duration + decode_duration).as_secs_f64()),
time_to_first_token_ms: Some(prefill_duration.as_secs_f64() * 1000.0),
prefill_tokens_per_sec: Some(if prefill_duration.as_secs_f64() > 0.0 {
prompt_len as f64 / prefill_duration.as_secs_f64()
} else {
0.0
}),
decode_tokens_per_sec: Some(if decode_duration.as_secs_f64() > 0.0 {
completion_tokens as f64 / decode_duration.as_secs_f64()
} else {
0.0
}),
gpu_sync_count: Some(mlx_native::sync_count().saturating_sub(pre_syncs)),
gpu_dispatch_count: Some(mlx_native::dispatch_count().saturating_sub(pre_dispatches)),
cached_prompt_tokens: if prompt_cache_hit {
Some(prompt_len)
} else {
None
},
reasoning_tokens: if reasoning_token_count > 0 {
Some(reasoning_token_count)
} else {
None
},
};
send!(GenerationEvent::Done {
finish_reason,
prompt_tokens: prompt_len,
completion_tokens,
stats,
});
}
pub fn embed_qwen35(qwen: &mut Qwen35LoadedModel, prompt_tokens: &[u32]) -> Result<Vec<f32>> {
anyhow::ensure!(
!prompt_tokens.is_empty(),
"embed_qwen35: empty prompt_tokens"
);
let device =
MlxDevice::new().map_err(|e| anyhow::anyhow!("MlxDevice::new (qwen35 embed): {e}"))?;
let mut kv_cache = alloc_kv_cache_for_request(qwen, &device, prompt_tokens.len(), 0)?;
let positions = prefill_positions_for(prompt_tokens.len());
qwen.model
.forward_embed_last(prompt_tokens, &positions, &mut kv_cache, SlotId(0))
.context("Qwen35Model::forward_embed_last")
}
pub fn embed_qwen35_slot_aware(
qwen: &mut Qwen35LoadedModel,
prompt_tokens: &[u32],
kv_cache: &mut HybridKvCache,
slot_id: SlotId,
) -> Result<Vec<f32>> {
anyhow::ensure!(
!prompt_tokens.is_empty(),
"embed_qwen35_slot_aware: empty prompt_tokens"
);
anyhow::ensure!(
slot_id.0 < kv_cache.n_seqs,
"embed_qwen35_slot_aware: SlotOutOfRange slot={} max_slots={} \
(ADR-040 iter-C2d-cont-kernel iter-3)",
slot_id.0,
kv_cache.n_seqs,
);
let prompt_len = prompt_tokens.len();
let need_seq = prompt_len + 64;
if need_seq > kv_cache.max_seq_len as usize {
return Err(anyhow::anyhow!(
"embed_qwen35_slot_aware: per-request need_seq={} exceeds \
persistent cache max_seq_len={} (slot={} prompt_len={}). \
ADR-040 iter-C2d-cont-kernel iter-3 sizes the persistent \
cache to cfg.max_position_embeddings; use a shorter prompt.",
need_seq,
kv_cache.max_seq_len,
slot_id.0,
prompt_len
));
}
kv_cache
.reset_for_slot(slot_id)
.context("ADR-040 iter-C2d-cont-kernel iter-3: reset_for_slot at entry")?;
let positions = prefill_positions_for(prompt_len);
let embed_result = qwen
.model
.forward_embed_last(prompt_tokens, &positions, kv_cache, slot_id)
.context(
"Qwen35Model::forward_embed_last (ADR-040 iter-C2d-cont-kernel iter-3 slot-aware)",
);
let reset_exit_res = kv_cache.reset_for_slot(slot_id);
match (embed_result, reset_exit_res) {
(Ok(vec), Ok(())) => Ok(vec),
(Err(e), _) => Err(e),
(Ok(_), Err(e)) => {
Err(e).context("ADR-040 iter-C2d-cont-kernel iter-3: reset_for_slot at exit")
}
}
}
pub fn generate_qwen35_once_with_soft_tokens_slot_aware(
qwen: &mut Qwen35LoadedModel,
prompt_tokens: &[u32],
soft_tokens: &[crate::serve::forward_prefill::SoftTokenInjection<'_>],
params: &SamplingParams,
registration: Option<&ModelRegistration>,
kv_cache: &mut HybridKvCache,
slot_id: SlotId,
) -> Result<GenerationResult> {
if soft_tokens.is_empty() {
return generate_qwen35_once_slot_aware(
qwen,
prompt_tokens,
params,
registration,
kv_cache,
slot_id,
);
}
anyhow::ensure!(
!prompt_tokens.is_empty(),
"generate_qwen35_once_with_soft_tokens_slot_aware: empty prompt_tokens"
);
anyhow::ensure!(
slot_id.0 < kv_cache.n_seqs,
"generate_qwen35_once_with_soft_tokens_slot_aware: SlotOutOfRange slot={} \
max_slots={} (ADR-040 iter-C2d-cont-kernel iter-4)",
slot_id.0,
kv_cache.n_seqs,
);
let prompt_len = prompt_tokens.len();
let max_tokens = params.max_tokens.max(1);
let need_seq = prompt_len + max_tokens + 64;
if need_seq > kv_cache.max_seq_len as usize {
return Err(anyhow::anyhow!(
"generate_qwen35_once_with_soft_tokens_slot_aware: per-request \
need_seq={} exceeds persistent cache max_seq_len={} (slot={} \
prompt_len={} max_tokens={}). ADR-040 iter-C2d-cont-kernel iter-4 \
sizes the persistent cache to cfg.max_position_embeddings; reduce \
max_tokens or use a shorter prompt.",
need_seq,
kv_cache.max_seq_len,
slot_id.0,
prompt_len,
max_tokens
));
}
let is_greedy = is_greedy_eligible(params);
let want_logprobs = params.logprobs;
let mut logprobs_vec: Option<Vec<f32>> = if want_logprobs {
Some(Vec::with_capacity(max_tokens))
} else {
None
};
kv_cache
.reset_for_slot(slot_id)
.context("ADR-040 iter-C2d-cont-kernel iter-4: reset_for_slot at entry")?;
let prefill_start = Instant::now();
let positions = prefill_positions_for(prompt_len);
let prefill_logits = qwen
.model
.forward_gpu_last_logits_with_soft_tokens(
prompt_tokens,
&positions,
soft_tokens,
kv_cache,
slot_id,
)
.context(
"Qwen35Model::forward_gpu_last_logits_with_soft_tokens \
(slot-aware prefill, ADR-040 iter-C2d-cont-kernel iter-4)",
)?;
anyhow::ensure!(
prefill_logits.len() == qwen.vocab_size,
"qwen35 slot-aware soft-tokens prefill logits len {} != vocab_size {}",
prefill_logits.len(),
qwen.vocab_size
);
let mut next_token: u32 = if want_logprobs {
let mut logits = prefill_logits.clone();
let (tok, lp) = sample_logits_qwen35_with_logprob(&mut logits, params, &[]);
if let Some(v) = logprobs_vec.as_mut() {
v.push(lp);
}
tok
} else if is_greedy {
greedy_argmax_last_token(&prefill_logits, qwen.vocab_size as u32)
} else {
let mut logits = prefill_logits.clone();
sample_logits_qwen35(&mut logits, params, &[])
};
let prefill_duration = prefill_start.elapsed();
let decode_start = Instant::now();
let mut generated_tokens: Vec<u32> = Vec::with_capacity(max_tokens);
generated_tokens.push(next_token);
let first_fragment = qwen
.tokenizer
.decode(&[next_token], false)
.unwrap_or_default();
let mut decoded_text = first_fragment.clone();
let mut finish_reason: &'static str = "length";
if qwen.eos_token_ids.contains(&next_token) {
finish_reason = "stop";
} else if qwen35_hit_stop_string(&decoded_text, ¶ms.stop_strings) {
finish_reason = "stop";
qwen35_strip_trailing_stop(&mut decoded_text, ¶ms.stop_strings);
} else {
for step in 1..max_tokens {
let pos = (prompt_len + step - 1) as i32;
if pos as u32 >= kv_cache.max_seq_len {
tracing::warn!(
pos,
max_seq = kv_cache.max_seq_len,
"qwen35 slot-aware decode (soft tokens): hit kv-cache bound; \
stopping with finish=length",
);
break;
}
let decode_positions = vec![pos; 4];
next_token = if want_logprobs {
let logits_full = qwen
.model
.forward_gpu_last_logits(&[next_token], &decode_positions, kv_cache, slot_id)
.with_context(|| {
format!(
"forward_gpu_last_logits slot-aware decode step {step} \
(soft tokens, logprobs)"
)
})?;
let mut logits = logits_full;
let (tok, lp) =
sample_logits_qwen35_with_logprob(&mut logits, params, &generated_tokens);
if let Some(v) = logprobs_vec.as_mut() {
v.push(lp);
}
tok
} else if is_greedy {
qwen.model
.forward_gpu_greedy(&[next_token], &decode_positions, kv_cache, slot_id)
.with_context(|| {
format!(
"forward_gpu_greedy slot-aware decode step {step} \
(soft tokens; ADR-040 §6.1.50 iter-G)"
)
})?
} else {
let logits_full = qwen
.model
.forward_gpu_last_logits(&[next_token], &decode_positions, kv_cache, slot_id)
.with_context(|| {
format!(
"forward_gpu_last_logits slot-aware decode step {step} \
(soft tokens)"
)
})?;
let mut logits = logits_full;
sample_logits_qwen35(&mut logits, params, &generated_tokens)
};
if qwen.eos_token_ids.contains(&next_token) {
finish_reason = "stop";
break;
}
generated_tokens.push(next_token);
let fragment = qwen
.tokenizer
.decode(&[next_token], false)
.unwrap_or_default();
decoded_text.push_str(&fragment);
if qwen35_hit_stop_string(&decoded_text, ¶ms.stop_strings) {
finish_reason = "stop";
qwen35_strip_trailing_stop(&mut decoded_text, ¶ms.stop_strings);
break;
}
}
}
let decode_duration = decode_start.elapsed();
kv_cache
.reset_for_slot(slot_id)
.context("ADR-040 iter-C2d-cont-kernel iter-4: reset_for_slot at exit")?;
let (content, reasoning_text) = match registration {
Some(reg) if reg.has_reasoning() => super::registry::split_full_output_forced(
reg,
&decoded_text,
params.reasoning_forced_open,
),
_ => (decoded_text, None),
};
let reasoning_token_count = match registration {
Some(reg) if reg.has_reasoning() => {
let mut sp =
super::registry::make_reasoning_splitter(reg, params.reasoning_forced_open);
let mut count = 0usize;
for &tok in &generated_tokens {
let frag = qwen.tokenizer.decode(&[tok], false).unwrap_or_default();
if let Some(splitter) = sp.as_mut() {
let _ = splitter.feed(&frag);
if splitter.in_reasoning() {
count += 1;
}
}
}
count
}
_ => 0,
};
Ok(GenerationResult {
text: content,
reasoning_text,
prompt_tokens: prompt_len,
completion_tokens: generated_tokens.len(),
reasoning_tokens: if reasoning_token_count > 0 {
Some(reasoning_token_count)
} else {
None
},
finish_reason,
prefill_duration,
decode_duration,
cached_tokens: 0,
logprobs: logprobs_vec,
})
}
#[allow(clippy::too_many_arguments)]
pub fn generate_qwen35_once_with_soft_tokens_and_deepstack_slot_aware(
qwen: &mut Qwen35LoadedModel,
prompt_tokens: &[u32],
soft_tokens: &[crate::serve::forward_prefill::SoftTokenInjection<'_>],
deepstack: Option<&crate::serve::forward_prefill::DeepstackInjection<'_>>,
positions_flat: Option<&[i32]>,
params: &SamplingParams,
registration: Option<&ModelRegistration>,
kv_cache: &mut HybridKvCache,
slot_id: SlotId,
) -> Result<GenerationResult> {
if soft_tokens.is_empty() && deepstack.is_none() && positions_flat.is_none() {
return generate_qwen35_once_slot_aware(
qwen,
prompt_tokens,
params,
registration,
kv_cache,
slot_id,
);
}
anyhow::ensure!(
!prompt_tokens.is_empty(),
"generate_qwen35_once_with_soft_tokens_and_deepstack_slot_aware: \
empty prompt_tokens"
);
anyhow::ensure!(
slot_id.0 < kv_cache.n_seqs,
"generate_qwen35_once_with_soft_tokens_and_deepstack_slot_aware: \
SlotOutOfRange slot={} max_slots={} (ADR-040 iter-C2d-cont-kernel iter-4)",
slot_id.0,
kv_cache.n_seqs,
);
let prompt_len = prompt_tokens.len();
let max_tokens = params.max_tokens.max(1);
let need_seq = prompt_len + max_tokens + 64;
if need_seq > kv_cache.max_seq_len as usize {
return Err(anyhow::anyhow!(
"generate_qwen35_once_with_soft_tokens_and_deepstack_slot_aware: \
per-request need_seq={} exceeds persistent cache max_seq_len={} \
(slot={} prompt_len={} max_tokens={}). ADR-040 iter-C2d-cont-\
kernel iter-4 sizes the persistent cache to \
cfg.max_position_embeddings; reduce max_tokens or use a shorter \
prompt.",
need_seq,
kv_cache.max_seq_len,
slot_id.0,
prompt_len,
max_tokens
));
}
let is_greedy = is_greedy_eligible(params);
let want_logprobs = params.logprobs;
let mut logprobs_vec: Option<Vec<f32>> = if want_logprobs {
Some(Vec::with_capacity(max_tokens))
} else {
None
};
kv_cache
.reset_for_slot(slot_id)
.context("ADR-040 iter-C2d-cont-kernel iter-4: reset_for_slot at entry (deepstack)")?;
let prefill_start = Instant::now();
let positions_owned: Vec<i32>;
let positions: &[i32] = match positions_flat {
Some(p) => {
anyhow::ensure!(
p.len() == 4 * prompt_len,
"generate_qwen35_once_with_soft_tokens_and_deepstack_slot_aware: \
positions_flat.len() = {} != 4 * prompt_len = {}",
p.len(),
4 * prompt_len
);
p
}
None => {
positions_owned = prefill_positions_for(prompt_len);
&positions_owned
}
};
let prefill_logits = qwen
.model
.forward_gpu_last_logits_with_soft_tokens_and_deepstack(
prompt_tokens,
positions,
soft_tokens,
deepstack,
kv_cache,
slot_id,
)
.context(
"Qwen35Model::forward_gpu_last_logits_with_soft_tokens_and_deepstack \
(slot-aware prefill, ADR-040 iter-C2d-cont-kernel iter-4)",
)?;
anyhow::ensure!(
prefill_logits.len() == qwen.vocab_size,
"qwen35 slot-aware deepstack prefill logits len {} != vocab_size {}",
prefill_logits.len(),
qwen.vocab_size
);
let mut next_token: u32 = if want_logprobs {
let mut logits = prefill_logits.clone();
let (tok, lp) = sample_logits_qwen35_with_logprob(&mut logits, params, &[]);
if let Some(v) = logprobs_vec.as_mut() {
v.push(lp);
}
tok
} else if is_greedy {
greedy_argmax_last_token(&prefill_logits, qwen.vocab_size as u32)
} else {
let mut logits = prefill_logits.clone();
sample_logits_qwen35(&mut logits, params, &[])
};
let prefill_duration = prefill_start.elapsed();
let t_post: i32 = match positions_flat {
Some(p) => {
let mut max_t = 0i32;
for i in 0..prompt_len {
let v = p[i]; if v > max_t {
max_t = v;
}
}
max_t.saturating_add(1)
}
None => prompt_len as i32,
};
let decode_start = Instant::now();
let mut generated_tokens: Vec<u32> = Vec::with_capacity(max_tokens);
generated_tokens.push(next_token);
let first_fragment = qwen
.tokenizer
.decode(&[next_token], false)
.unwrap_or_default();
let mut decoded_text = first_fragment.clone();
let mut finish_reason: &'static str = "length";
if qwen.eos_token_ids.contains(&next_token) {
finish_reason = "stop";
} else if qwen35_hit_stop_string(&decoded_text, ¶ms.stop_strings) {
finish_reason = "stop";
qwen35_strip_trailing_stop(&mut decoded_text, ¶ms.stop_strings);
} else {
for step in 1..max_tokens {
let pos = t_post + (step as i32 - 1);
if pos as u32 >= kv_cache.max_seq_len {
tracing::warn!(
pos,
max_seq = kv_cache.max_seq_len,
"qwen35 slot-aware decode (deepstack): hit kv-cache bound; \
stopping with finish=length",
);
break;
}
let decode_positions = vec![pos; 4];
next_token = if want_logprobs {
let logits_full = qwen
.model
.forward_gpu_last_logits(&[next_token], &decode_positions, kv_cache, slot_id)
.with_context(|| {
format!(
"forward_gpu_last_logits slot-aware decode step {step} \
(deepstack, logprobs)"
)
})?;
let mut logits = logits_full;
let (tok, lp) =
sample_logits_qwen35_with_logprob(&mut logits, params, &generated_tokens);
if let Some(v) = logprobs_vec.as_mut() {
v.push(lp);
}
tok
} else if is_greedy {
qwen.model
.forward_gpu_greedy(&[next_token], &decode_positions, kv_cache, slot_id)
.with_context(|| {
format!(
"forward_gpu_greedy slot-aware decode step {step} \
(deepstack; ADR-040 §6.1.50 iter-G)"
)
})?
} else {
let logits_full = qwen
.model
.forward_gpu_last_logits(&[next_token], &decode_positions, kv_cache, slot_id)
.with_context(|| {
format!(
"forward_gpu_last_logits slot-aware decode step {step} \
(deepstack)"
)
})?;
let mut logits = logits_full;
sample_logits_qwen35(&mut logits, params, &generated_tokens)
};
if qwen.eos_token_ids.contains(&next_token) {
finish_reason = "stop";
break;
}
generated_tokens.push(next_token);
let fragment = qwen
.tokenizer
.decode(&[next_token], false)
.unwrap_or_default();
decoded_text.push_str(&fragment);
if qwen35_hit_stop_string(&decoded_text, ¶ms.stop_strings) {
finish_reason = "stop";
qwen35_strip_trailing_stop(&mut decoded_text, ¶ms.stop_strings);
break;
}
}
}
let decode_duration = decode_start.elapsed();
kv_cache
.reset_for_slot(slot_id)
.context("ADR-040 iter-C2d-cont-kernel iter-4: reset_for_slot at exit (deepstack)")?;
let (content, reasoning_text) = match registration {
Some(reg) if reg.has_reasoning() => super::registry::split_full_output_forced(
reg,
&decoded_text,
params.reasoning_forced_open,
),
_ => (decoded_text, None),
};
let reasoning_token_count = match registration {
Some(reg) if reg.has_reasoning() => {
let mut sp =
super::registry::make_reasoning_splitter(reg, params.reasoning_forced_open);
let mut count = 0usize;
for &tok in &generated_tokens {
let frag = qwen.tokenizer.decode(&[tok], false).unwrap_or_default();
if let Some(splitter) = sp.as_mut() {
let _ = splitter.feed(&frag);
if splitter.in_reasoning() {
count += 1;
}
}
}
count
}
_ => 0,
};
Ok(GenerationResult {
text: content,
reasoning_text,
prompt_tokens: prompt_len,
completion_tokens: generated_tokens.len(),
reasoning_tokens: if reasoning_token_count > 0 {
Some(reasoning_token_count)
} else {
None
},
finish_reason,
prefill_duration,
decode_duration,
cached_tokens: 0,
logprobs: logprobs_vec,
})
}
fn params_tool_call_policy_for_qwen35_stream() -> super::engine::ToolCallPolicy {
super::engine::ToolCallPolicy::Auto
}
#[cfg(test)]
mod tests {
use super::*;
use crate::inference::models::qwen35::kv_cache::HybridKvCache;
use crate::inference::models::qwen35::{
default_layer_types, Qwen35Config, Qwen35MoeConfig, Qwen35Variant,
};
use mlx_native::MlxDevice;
fn moe_cfg_40layer_for_cache_test() -> Qwen35Config {
Qwen35Config {
variant: Qwen35Variant::Moe,
hidden_size: 64,
num_hidden_layers: 4,
num_attention_heads: 4,
num_key_value_heads: 2,
head_dim: 16,
linear_num_key_heads: 4,
linear_num_value_heads: 8,
linear_key_head_dim: 16,
linear_value_head_dim: 16,
linear_conv_kernel_dim: 4,
full_attention_interval: 4,
layer_types: default_layer_types(4, 4),
partial_rotary_factor: 0.25,
rope_theta: 1e7,
rotary_dim: 4,
mrope_section: [1, 1, 0, 0],
mrope_interleaved: true,
rms_norm_eps: 1e-6,
max_position_embeddings: 1024,
vocab_size: 256,
attn_output_gate: true,
mtp_num_hidden_layers: 0,
mtp_use_dedicated_embeddings: true,
intermediate_size: None,
moe: Some(Qwen35MoeConfig {
moe_intermediate_size: 16,
num_experts: 4,
num_experts_per_tok: 2,
shared_expert_intermediate_size: 16,
}),
}
}
fn greedy_params() -> SamplingParams {
SamplingParams {
max_tokens: 16,
..SamplingParams::default()
}
}
fn sampling_params_with_temperature() -> SamplingParams {
SamplingParams {
temperature: 0.7,
max_tokens: 16,
..SamplingParams::default()
}
}
#[test]
fn hybrid_prompt_cache_new_is_empty() {
let cache = HybridPromptCache::new();
assert!(!cache.has_entry());
assert!(cache.snapshot().is_none());
assert!(cache.try_match(&[1, 2, 3], &greedy_params()).is_none());
}
#[test]
fn hybrid_prompt_cache_invalidates_on_prompt_divergence() {
let cfg = moe_cfg_40layer_for_cache_test();
let device = MlxDevice::new().expect("device");
let kv = HybridKvCache::new(&cfg, &device, 16, 1).expect("kv");
let snap = kv.snapshot(&device).expect("snap");
let prompt = vec![10u32, 20, 30, 40];
let mut cache = HybridPromptCache::new();
cache.update(prompt.clone(), snap, 99u32, &greedy_params());
assert!(cache.has_entry());
assert_eq!(
cache.try_match(&prompt, &greedy_params()),
Some(prompt.len()),
"exact-match prompt should hit"
);
let mut diverged = prompt.clone();
diverged[2] = 999;
assert!(
cache.try_match(&diverged, &greedy_params()).is_none(),
"divergent prompt should miss"
);
let mut shorter = prompt.clone();
shorter.pop();
assert!(
cache.try_match(&shorter, &greedy_params()).is_none(),
"shorter prompt should miss"
);
let mut longer = prompt.clone();
longer.push(50);
assert!(
cache.try_match(&longer, &greedy_params()).is_none(),
"longer prompt should miss"
);
}
#[test]
fn hybrid_prompt_cache_invalidates_on_genparams_mismatch() {
let cfg = moe_cfg_40layer_for_cache_test();
let device = MlxDevice::new().expect("device");
let kv = HybridKvCache::new(&cfg, &device, 16, 1).expect("kv");
let snap = kv.snapshot(&device).expect("snap");
let prompt = vec![1u32, 2, 3];
let mut cache = HybridPromptCache::new();
let stored_params = SamplingParams {
max_tokens: 32,
stop_strings: vec!["</done>".into()],
..SamplingParams::default()
};
cache.update(prompt.clone(), snap, 7u32, &stored_params);
assert_eq!(
cache.try_match(&prompt, &stored_params),
Some(prompt.len()),
"matching key should hit"
);
let diff_max = SamplingParams {
max_tokens: 64,
stop_strings: vec!["</done>".into()],
..SamplingParams::default()
};
assert!(
cache.try_match(&prompt, &diff_max).is_none(),
"max_tokens mismatch must miss"
);
let diff_stop = SamplingParams {
max_tokens: 32,
stop_strings: vec!["</STOP>".into()],
..SamplingParams::default()
};
assert!(
cache.try_match(&prompt, &diff_stop).is_none(),
"stop_strings mismatch must miss"
);
}
#[test]
fn hybrid_prompt_cache_sampling_mode_bypasses_lookup_and_store() {
let cfg = moe_cfg_40layer_for_cache_test();
let device = MlxDevice::new().expect("device");
let kv = HybridKvCache::new(&cfg, &device, 16, 1).expect("kv");
let snap = kv.snapshot(&device).expect("snap");
let prompt = vec![1u32, 2, 3];
let mut cache = HybridPromptCache::new();
cache.update(prompt.clone(), snap, 1u32, &greedy_params());
assert!(cache.has_entry());
assert!(
cache
.try_match(&prompt, &sampling_params_with_temperature())
.is_none(),
"sampling-mode lookup must miss"
);
let kv2 = HybridKvCache::new(&cfg, &device, 16, 1).expect("kv2");
let snap2 = kv2.snapshot(&device).expect("snap2");
let mut cache2 = HybridPromptCache::new();
cache2.update(
prompt.clone(),
snap2,
5u32,
&sampling_params_with_temperature(),
);
assert!(!cache2.has_entry(), "sampling-mode update must be a no-op");
}
#[test]
fn splitter_helper_extracts_reasoning_from_qwen35_thinkblocks() {
let reg = super::super::registry::QWEN35;
let raw = "Sure! <think>Let me solve this step by step.</think>The answer is 42.";
let (content, reasoning) = super::super::registry::split_full_output(®, raw);
assert_eq!(
content, "Sure! The answer is 42.",
"content must exclude the <think>...</think> span"
);
assert_eq!(
reasoning.as_deref(),
Some("Let me solve this step by step."),
"reasoning must contain the inner span"
);
}
#[test]
fn splitter_helper_extracts_tool_calls_from_qwen35_toolblocks() {
let reg = super::super::registry::QWEN35;
let mut sp = super::super::registry::ToolCallSplitter::from_registration(®)
.expect("QWEN35 has tool markers");
let raw =
"Let me search.<tool_call><function=search><parameter=q>weather</parameter></function></tool_call> Done.";
let mut events = Vec::new();
events.extend(sp.feed(raw));
if let Some(tail) = sp.finish() {
events.push(tail);
}
let mut saw_open = false;
let mut saw_text = false;
let mut saw_close = false;
let mut content_runs: Vec<String> = Vec::new();
for ev in events {
use super::super::registry::ToolCallEvent::*;
match ev {
Content(t) => content_runs.push(t),
ToolCallOpen => saw_open = true,
ToolCallText(_) => saw_text = true,
ToolCallClose => saw_close = true,
}
}
assert!(saw_open, "must observe ToolCallOpen for QWEN35 marker");
assert!(saw_text, "must observe ToolCallText body");
assert!(saw_close, "must observe ToolCallClose");
let joined: String = content_runs.join("");
assert!(
joined.contains("Let me search."),
"preamble content must round-trip"
);
assert!(
joined.contains(" Done."),
"post-close content must round-trip"
);
}
#[test]
fn qwen35_loaded_model_has_initialized_prompt_cache() {
let cache = HybridPromptCache::default();
assert!(!cache.has_entry());
assert!(cache.snapshot().is_none());
assert!(cache.try_match(&[1, 2], &greedy_params()).is_none());
}
#[test]
fn qwen35_loaded_model_load_errors_when_path_missing() {
let opts = LoadOptions {
model_path: std::path::PathBuf::from("/tmp/iter-215-does-not-exist.gguf"),
tokenizer_path: None,
config_path: None,
dwq_overlay_path: None,
kv_persist_dir: None,
};
let res = Qwen35LoadedModel::load(&opts);
assert!(res.is_err());
let msg = format!("{:#}", res.err().unwrap());
assert!(
msg.contains("Model not found"),
"expected 'Model not found' in error; got: {msg}"
);
}
use super::super::registry::{
SplitSlot as _SplitSlot, ToolCallEvent as _ToolCallEvent,
ToolCallSplitter as _ToolCallSplitter, QWEN35,
};
#[test]
fn wedge4e_reasoning_splitter_is_mode_invariant() {
let mut sp = super::super::registry::make_reasoning_splitter(&QWEN35, false)
.expect("Qwen35 has reasoning markers");
let mut all_pairs: Vec<(_SplitSlot, String)> = Vec::new();
for frag in ["<thi", "nk>let me reason", " more</thin", "k>final answer"] {
for pair in sp.feed(frag) {
all_pairs.push(pair);
}
}
if let Some(tail) = sp.finish() {
all_pairs.push(tail);
}
let mut reasoning = String::new();
let mut content = String::new();
for (slot, text) in &all_pairs {
match slot {
_SplitSlot::Reasoning => reasoning.push_str(text),
_SplitSlot::Content => content.push_str(text),
}
}
assert_eq!(
reasoning, "let me reason more",
"Wedge-4e: reasoning text must be cleanly extracted"
);
assert_eq!(
content, "final answer",
"Wedge-4e: content text must NOT contain reasoning brackets"
);
}
#[test]
fn wedge4e_tool_call_splitter_is_mode_invariant() {
let mut tcs =
_ToolCallSplitter::from_registration(&QWEN35).expect("Qwen35 has tool markers");
let mut events: Vec<_ToolCallEvent> = Vec::new();
for frag in [
"let me search.<tool_",
"call>{\"name\":\"search\",\"arguments\":{\"q\":\"x\"}}</tool_",
"call> done.",
] {
for ev in tcs.feed(frag) {
events.push(ev);
}
}
if let Some(tail) = tcs.finish() {
events.push(tail);
}
let mut saw_open = false;
let mut saw_close = false;
let mut body = String::new();
let mut content_runs: Vec<String> = Vec::new();
for ev in events {
match ev {
_ToolCallEvent::Content(t) => content_runs.push(t),
_ToolCallEvent::ToolCallOpen => saw_open = true,
_ToolCallEvent::ToolCallText(t) => body.push_str(&t),
_ToolCallEvent::ToolCallClose => saw_close = true,
}
}
assert!(saw_open, "Wedge-4e: must observe ToolCallOpen");
assert!(saw_close, "Wedge-4e: must observe ToolCallClose");
assert!(
body.contains("\"name\":\"search\""),
"Wedge-4e: tool-call body must round-trip; got {body:?}"
);
let joined: String = content_runs.join("");
assert!(
joined.contains("let me search."),
"Wedge-4e: pre-tool-call content must round-trip"
);
assert!(
joined.contains(" done."),
"Wedge-4e: post-tool-call content must round-trip"
);
}
#[test]
fn wedge4e_legacy_stream_entry_is_thin_wrapper() {
let src = include_str!("engine_qwen35.rs");
assert!(
src.contains("generate_stream_qwen35_once_extended(\n qwen,\n prompt_tokens,\n &[],\n None,\n None,"),
"Wedge-4e: generate_stream_qwen35_once must delegate to \
generate_stream_qwen35_once_extended with empty extensions \
— the byte-identical text-only regression contract \
requires this exact shape"
);
}
#[test]
fn wedge4e_extended_stream_validates_positions_len() {
let prompt_len = 5usize;
let bad_positions = vec![0i32; 17]; let expected_err = format!(
"qwen35 stream (wedge-4e): positions_flat.len() = {} != 4 * prompt_len = {}",
bad_positions.len(),
4 * prompt_len
);
let src = include_str!("engine_qwen35.rs");
assert!(
src.contains("qwen35 stream (wedge-4e): positions_flat.len() = "),
"Wedge-4e: positions_flat length validator must surface \
the actionable diagnostic byte string"
);
assert!(
expected_err.contains("17 != 4 * prompt_len = 20"),
"expected_err format check"
);
}
#[test]
fn wedge4e_t_post_advance_rule_matches_non_streaming_sibling() {
let src = include_str!("engine_qwen35.rs");
assert!(
src.contains("max_t.saturating_add(1)"),
"Wedge-4e: t_post must use saturating_add(1) over axis-0 max"
);
assert!(
src.contains("None => prompt_len as i32,"),
"Wedge-4e: t_post must default to prompt_len when no 3D \
positions supplied (text-only byte-identity)"
);
assert!(
src.contains("let pos = t_post + (step as i32 - 1);"),
"Wedge-4e: streaming decode position formula must use t_post \
advance — not the legacy (prompt_len + step - 1) form, \
which would silently misalign on multi-image prefill"
);
}
#[test]
fn wedge4e_extended_stream_bypasses_prompt_cache_on_extension() {
let src = include_str!("engine_qwen35.rs");
assert!(
src.contains("let prompt_cache_hit = !has_extension"),
"Wedge-4e: streaming prompt-cache MUST be bypassed when \
any extension is present (cache key is prompt_tokens \
only — same placeholder ids + different image ⇒ false \
hit)"
);
assert!(
src.contains("if is_greedy && !has_extension {"),
"Wedge-4e: streaming prompt-cache write MUST be skipped on \
extension paths to avoid poisoning subsequent text-only \
requests with a soft-token-tainted KV snapshot"
);
}
#[test]
fn lcp_prefix_stores_are_not_greedy_gated() {
let src = include_str!("engine_qwen35.rs");
let forbidden = concat!("lcp_resume_enabled", " && ", "is_greedy");
assert!(
!src.contains(forbidden),
"LCP prefix-KV stores must NOT be greedy-gated — KV state is \
sampling-independent; greedy gating belongs only on \
HybridPromptCache's decoded-token replay"
);
let gate = concat!(
"if lcp_resume_enabled && stride_aligned",
" && !mid_store_disabled {"
);
assert_eq!(
src.matches(gate).count(),
2,
"non-stream + stream mid-prefill store gates must both be \
`lcp_resume_enabled && stride_aligned && !mid_store_disabled`"
);
assert!(
src.contains("if is_greedy && !has_extension {"),
"HybridPromptCache write must REMAIN greedy-gated (first \
decoded token replay is sampling-dependent)"
);
}
#[test]
fn wedge4e_phase_2c_soft_token_guard_is_removed() {
let src = include_str!("engine.rs");
assert!(
!src.contains("Qwen35 streaming path does not yet support"),
"Wedge-4e: Phase-2c soft_token guard at engine.rs's \
LoadedModel::Qwen35 streaming arm MUST be removed — the \
extended streaming entry now threads soft_tokens + \
deepstack + positions through to the LM forward"
);
assert!(
!src.contains("For Qwen3-VL image-bearing chat, set \\\"stream\\\": false."),
"Wedge-4e: actionable diagnostic about set stream=false \
MUST be removed (the streaming path is now the \
production path)"
);
assert!(
src.contains("generate_stream_qwen35_once_extended"),
"Wedge-4e: streaming arm must dispatch through the \
extended entry"
);
}
#[test]
fn wedge4e_handler_streaming_501_reject_is_removed() {
let src = include_str!("handlers.rs");
assert!(
!src.contains("streaming chat with Qwen3-VL DeepStack injection is not yet"),
"Wedge-4e: handler-side streaming 501 reject MUST be \
removed — streaming Qwen3-VL chat now flows through \
generate_stream_with_deepstack"
);
assert!(
src.contains("generate_stream_with_deepstack"),
"Wedge-4e: handler must call generate_stream_with_deepstack \
so soft_tokens + deepstack + positions reach the worker"
);
}
fn synth_loaded_model_for_alloc_test(
cfg: Qwen35Config,
tq_kv_active: bool,
) -> Qwen35LoadedModel {
Qwen35LoadedModel {
model: super::Qwen35Model::empty_from_cfg(cfg),
tokenizer: tokenizers::Tokenizer::new(tokenizers::models::bpe::BPE::default()),
chat_template: "{{ messages }}".to_string(),
model_id: "iter-12-test".to_string(),
model_path: std::path::PathBuf::from("/tmp/iter-12-test.gguf"),
eos_token_ids: vec![151_645],
hidden_size: 64,
vocab_size: 256,
context_length: Some(1024),
quant_type: Some("Q4_K".to_string()),
load_duration: std::time::Duration::from_millis(1),
provenance: crate::core::provenance::Provenance::External,
prompt_cache: HybridPromptCache::new(),
lcp_registry: crate::serve::kv_persist::lcp_registry::LcpRegistry::new(1),
kv_metrics_sink: None,
disk_persistor: None,
lcp_hydrated_for_cfg: std::collections::HashSet::new(),
tq_kv_active,
persistent_kv_cache: None,
}
}
#[test]
fn alloc_kv_cache_for_request_tq_off_keeps_full_attn_tq_none() {
let device = match MlxDevice::new() {
Ok(d) => d,
Err(e) => {
eprintln!("skipping: no Metal device: {e}");
return;
}
};
let cfg = moe_cfg_40layer_for_cache_test();
let qwen = synth_loaded_model_for_alloc_test(cfg, false);
let cache =
alloc_kv_cache_for_request(&qwen, &device, 32, 16).expect("alloc_kv_cache_for_request");
assert!(!cache.full_attn.is_empty(), "fixture has full-attn layers");
for (i, slot) in cache.full_attn.iter().enumerate() {
assert!(
slot.tq.is_none(),
"tq_kv_active=false: full_attn[{i}].tq must be None \
(legacy F32 path preserved)"
);
}
}
#[test]
fn alloc_kv_cache_for_request_tq_on_populates_tq_per_full_attn_slot() {
let device = match MlxDevice::new() {
Ok(d) => d,
Err(e) => {
eprintln!("skipping: no Metal device: {e}");
return;
}
};
let mut cfg = moe_cfg_40layer_for_cache_test();
cfg.head_dim = 256;
cfg.num_attention_heads = 8;
cfg.num_key_value_heads = 2;
let qwen = synth_loaded_model_for_alloc_test(cfg, true);
let cache =
alloc_kv_cache_for_request(&qwen, &device, 32, 16).expect("alloc_kv_cache_for_request");
assert!(!cache.full_attn.is_empty());
for (i, slot) in cache.full_attn.iter().enumerate() {
assert!(
slot.tq.is_some(),
"tq_kv_active=true: full_attn[{i}].tq must be populated"
);
let tq = slot.tq.as_ref().unwrap();
assert_eq!(tq.norms_per_pos, 1, "head_dim=256 → norms_per_pos=1");
}
}
}