hugo/DS4Server
1
0 Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

Optimize Qwen3.8 inference on Apple silicon

Browse Source
This commit is contained in:
Georg Bauer committed 2026-09-04 23:27:09 +02:00
1 parent bf82df77cb
commit bd6353804b
14 files changed
+10149 -1145
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/config.rs
+41 -6
View File
@@ -24,6 +24,7 @@ impl ModelPreferences {
Self {
reasoning_mode: model.default_reasoning_mode(),
speculative: SpeculativePreferences {
glm_mtp: model == ModelChoice::Qwen38FlashNext,
dspark_enabled: model == ModelChoice::DeepSeekV4Flash0731,
..SpeculativePreferences::default()
},
@@ -308,7 +309,13 @@ impl Config {
.map_err(|error| format!("Could not read {}: {error}", path.display()))?;
drop_legacy_model_settings(&mut value);
let migrated = migrate_deprecated_flash(&mut value);
let dspark_explicit = deepseek_0731_dspark_is_explicit(&value);
let dspark_explicit = model_speculative_field_is_explicit(
&value,
ModelChoice::DeepSeekV4Flash0731,
"dspark_enabled",
);
let qwen_mtp_explicit =
model_speculative_field_is_explicit(&value, ModelChoice::Qwen38FlashNext, "glm_mtp");
let explicit_reasoning =
MODEL_CHOICES.map(|model| (model, model_reasoning_mode_is_explicit(&value, model)));
let mut config: Self = serde_norway::from_value(value)
@@ -331,6 +338,14 @@ impl Config {
.speculative
.dspark_enabled = true;
}
if !qwen_mtp_explicit {
config
.model_profiles
.get_mut(&ModelChoice::Qwen38FlashNext)
.expect("Qwen model profile is present")
.speculative
.glm_mtp = true;
}
config.validate()?;
if migrated {
config.save(path)?;
@@ -446,19 +461,17 @@ fn model_reasoning_mode_is_explicit(value: &Value, model: ModelChoice) -> bool {
.is_some_and(|profile| profile.contains_key(Value::String("reasoning_mode".into())))
}
fn deepseek_0731_dspark_is_explicit(value: &Value) -> bool {
fn model_speculative_field_is_explicit(value: &Value, model: ModelChoice, field: &str) -> bool {
let Value::Mapping(root) = value else {
return false;
};
root.get(Value::String("model_profiles".into()))
.and_then(Value::as_mapping)
.and_then(|profiles| {
profiles.get(Value::String(ModelChoice::DeepSeekV4Flash0731.id().into()))
})
.and_then(|profiles| profiles.get(Value::String(model.id().into())))
.and_then(Value::as_mapping)
.and_then(|profile| profile.get(Value::String("speculative".into())))
.and_then(Value::as_mapping)
.is_some_and(|section| section.contains_key(Value::String("dspark_enabled".into())))
.is_some_and(|section| section.contains_key(Value::String(field.into())))
}
fn drop_legacy_model_settings(value: &mut Value) {
@@ -603,6 +616,28 @@ mod tests {
ReasoningMode::Medium
);
assert_eq!(loaded.active_generation().context_tokens, 131_072);
assert!(
loaded
.runtime_for(ModelChoice::Qwen38FlashNext)
.speculative
.glm_mtp
);
let mut opted_out = loaded;
opted_out
.model_profiles
.get_mut(&ModelChoice::Qwen38FlashNext)
.unwrap()
.speculative
.glm_mtp = false;
opted_out.save(&path).unwrap();
assert!(
!Config::load(&path)
.unwrap()
.runtime_for(ModelChoice::Qwen38FlashNext)
.speculative
.glm_mtp
);
fs::remove_dir_all(directory).unwrap();
}
src/engine.rs
+342 -36
View File
@@ -792,6 +792,9 @@ impl Generator {
let result = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress);
self.publish_execution_stats();
let (mut output, prompt_complete) = result?;
if prompt_complete {
self.executor.finalize_pending_sampled_token()?;
}
let mut completed = messages.to_vec();
completed.push(output.message.clone());
output.previous_checkpoint_bytes =
@@ -870,6 +873,9 @@ impl Generator {
let result = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress);
self.publish_execution_stats();
let (mut output, prompt_complete) = result?;
if prompt_complete {
self.executor.finalize_pending_sampled_token()?;
}
let mut completed = messages.to_vec();
completed.push(output.message.clone());
output.previous_checkpoint_bytes = previous_checkpoint
@@ -1462,7 +1468,12 @@ impl Generator {
self.executor.set_vision_overlays(overlays)?;
self.metrics.kv_prefix_reused(reused);
progress(self.executor.position(), self.executor.context(), None);
let mut rng = Rng::new(settings.seed.unwrap_or(0x4453_3453_4552_5645));
let seed = settings.seed.unwrap_or(0x4453_3453_4552_5645);
let mut rng = if matches!(self.executor.model(), ModelRef::Qwen(_)) {
Rng::new_qwen(seed)
} else {
Rng::new(seed)
};
let mut reasoning = settings.reasoning_mode != ReasoningMode::Direct;
let mut generated = ChatTurn {
user: false,
@@ -1525,14 +1536,32 @@ impl Generator {
true,
));
}
let token = sample(
self.executor.logits(),
settings.temperature,
settings.top_p,
settings.min_p,
settings.top_k,
&mut rng,
);
let token = self
.executor
.take_pending_sampled_token()
.unwrap_or_else(|| {
if matches!(self.executor.model(), ModelRef::Qwen(_)) {
sample_probabilities_f64(
&qwen_sampling_probabilities(
self.executor.logits(),
settings.temperature,
settings.top_p,
settings.min_p,
settings.top_k,
),
&mut rng,
)
} else {
sample(
self.executor.logits(),
settings.temperature,
settings.top_p,
settings.min_p,
settings.top_k,
&mut rng,
)
}
});
if self
.executor
.model()
@@ -1903,6 +1932,123 @@ fn sample(
sample_probabilities(&probabilities, rng, None)
}
#[cfg(any(target_os = "macos", test))]
fn qwen_sampling_probabilities(
logits: &[f32],
temperature: f32,
top_p: f32,
min_p: f32,
top_k: i32,
) -> Vec<(usize, f64)> {
let greedy = || {
vec![(
logits
.iter()
.enumerate()
.max_by(|a, b| a.1.total_cmp(b.1))
.map_or(0, |(index, _)| index),
1.0,
)]
};
if temperature <= 0.0 {
return greedy();
}
let maximum = logits
.iter()
.copied()
.filter(|value| value.is_finite())
.fold(f32::NEG_INFINITY, f32::max);
if !maximum.is_finite() {
return greedy();
}
let top_p = if top_p <= 0.0 || top_p > 1.0 {
1.0
} else {
f64::from(top_p)
};
let rank = |a: &(usize, f64), b: &(usize, f64)| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0));
let inverse_temperature = temperature.recip();
let requested = usize::try_from(top_k.max(0)).unwrap_or(usize::MAX);
let mut probabilities = Vec::with_capacity(if requested > 0 {
requested.min(logits.len())
} else {
logits.len()
});
let mut total = 0.0_f32;
for (index, logit) in logits.iter().copied().enumerate() {
if !logit.is_finite() {
continue;
}
let probability = ((logit - maximum) * inverse_temperature).exp();
total += probability;
let candidate = (index, f64::from(probability));
if requested == 0 {
probabilities.push(candidate);
continue;
}
let slot = probabilities.partition_point(|current| {
current.1 > candidate.1 || (current.1 == candidate.1 && current.0 < candidate.0)
});
if slot < requested {
probabilities.insert(slot, candidate);
if probabilities.len() > requested {
probabilities.pop();
}
}
}
if probabilities.is_empty() || !total.is_finite() || total <= 0.0 {
return greedy();
}
for (_, probability) in &mut probabilities {
*probability = f64::from((*probability as f32) / total);
}
let count = probabilities.len();
if requested == 0 {
probabilities.sort_unstable_by(rank);
}
let minimum = f64::from(min_p.max(0.0));
let mut kept = 0.0;
let mut end = 0;
for (_, probability) in probabilities.iter().take(count) {
if end > 0 && *probability < minimum * probabilities[0].1 {
break;
}
if kept >= top_p {
break;
}
kept += *probability;
end += 1;
}
probabilities.truncate(end);
if probabilities.is_empty() || !kept.is_finite() || kept <= 0.0 {
return greedy();
}
for (_, probability) in &mut probabilities {
*probability /= kept;
}
probabilities.sort_unstable_by_key(|(token, _)| *token);
probabilities
}
#[cfg(any(target_os = "macos", test))]
fn sample_probabilities_f64(probabilities: &[(usize, f64)], rng: &mut Rng) -> i32 {
let mut choice = rng.unit_f64() * probabilities.iter().map(|(_, p)| p).sum::<f64>();
for (token, probability) in probabilities {
if !probability.is_finite() || *probability <= 0.0 {
continue;
}
choice -= probability;
if choice <= 0.0 {
return *token as i32;
}
}
probabilities
.iter()
.rev()
.find(|(_, probability)| probability.is_finite() && *probability > 0.0)
.map_or(0, |(token, _)| *token as i32)
}
#[cfg(any(target_os = "macos", test))]
fn sampling_probabilities(
logits: &[f32],
@@ -1938,20 +2084,29 @@ fn sampling_probabilities(
top_p
};
let min_p = min_p.max(0.0);
let mut probabilities: Vec<(usize, f32)> = logits
let rank = |a: &(usize, f32), b: &(usize, f32)| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0));
let mut candidates = logits
.iter()
.copied()
.enumerate()
.filter(|(_, logit)| logit.is_finite())
.map(|(index, logit)| (index, ((*logit - maximum) / temperature).exp()))
.collect();
.collect::<Vec<_>>();
if top_k > 0 {
let count = candidates.len().min((top_k as usize).min(1024));
if count < candidates.len() {
candidates.select_nth_unstable_by(count, rank);
candidates.truncate(count);
}
}
let mut probabilities = candidates
.into_iter()
.map(|(index, logit)| (index, ((logit - maximum) / temperature).exp()))
.collect::<Vec<_>>();
if probabilities.is_empty() {
return greedy();
}
if top_p < 1.0 || top_k > 0 || min_p > 0.0 {
probabilities.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
if top_k > 0 {
probabilities.truncate(probabilities.len().min((top_k as usize).min(1024)));
}
probabilities.sort_unstable_by(rank);
}
let total: f32 = probabilities
.iter()
@@ -2037,7 +2192,7 @@ fn exact_delta_sample(
)
}
#[cfg(any(target_os = "macos", test))]
#[cfg(test)]
fn sample_from_logits(
logits: &[f32],
temperature: f32,
@@ -2046,14 +2201,13 @@ fn sample_from_logits(
top_k: i32,
rng: &mut Rng,
) -> i32 {
sample_probabilities(
&sampling_probabilities(logits, temperature, top_p, min_p, top_k),
sample_probabilities_f64(
&qwen_sampling_probabilities(logits, temperature, top_p, min_p, top_k),
rng,
None,
)
}
#[cfg(any(target_os = "macos", test))]
#[cfg(test)]
#[allow(clippy::too_many_arguments)]
fn exact_speculative_sample(
target_logits: &[f32],
@@ -2065,19 +2219,53 @@ fn exact_speculative_sample(
top_k: i32,
rng: &mut Rng,
) -> (i32, bool) {
let mut target = sampling_probabilities(target_logits, temperature, top_p, min_p, top_k);
let mut proposal = sampling_probabilities(draft_logits, temperature, top_p, min_p, top_k);
target.sort_unstable_by_key(|(token, _)| *token);
proposal.sort_unstable_by_key(|(token, _)| *token);
let probability = |values: &[(usize, f32)], token: usize| {
let proposal = qwen_sampling_probabilities(draft_logits, temperature, top_p, min_p, top_k);
exact_speculative_sample_with_probabilities(
target_logits,
&proposal,
draft,
temperature,
top_p,
min_p,
top_k,
rng,
)
}
#[cfg(test)]
#[allow(clippy::too_many_arguments)]
fn exact_speculative_sample_with_probabilities(
target_logits: &[f32],
proposal: &[(usize, f64)],
draft: i32,
temperature: f32,
top_p: f32,
min_p: f32,
top_k: i32,
rng: &mut Rng,
) -> (i32, bool) {
let target = qwen_sampling_probabilities(target_logits, temperature, top_p, min_p, top_k);
exact_speculative_sample_from_probabilities(&target, proposal, draft, rng)
}
#[cfg(any(target_os = "macos", test))]
fn exact_speculative_sample_from_probabilities(
target: &[(usize, f64)],
proposal: &[(usize, f64)],
draft: i32,
rng: &mut Rng,
) -> (i32, bool) {
let probability = |values: &[(usize, f64)], token: usize| {
values
.binary_search_by_key(&token, |(candidate, _)| *candidate)
.ok()
.map_or(0.0, |index| values[index].1)
};
let target_probability = probability(&target, draft as usize);
let draft_probability = probability(&proposal, draft as usize);
if draft_probability > 0.0 && rng.unit() <= (target_probability / draft_probability).min(1.0) {
let target_probability = probability(target, draft as usize);
let draft_probability = probability(proposal, draft as usize);
if draft_probability > 0.0
&& rng.unit_f64() <= (target_probability / draft_probability).min(1.0)
{
return (draft, true);
}
@@ -2086,8 +2274,8 @@ fn exact_speculative_sample(
.max_by(|left, right| left.1.total_cmp(&right.1))
.map_or(draft, |(token, _)| *token as i32);
let mut residual = Vec::with_capacity(target.len());
for (token, target_probability) in target {
let remaining = target_probability - probability(&proposal, token);
for &(token, target_probability) in target {
let remaining = target_probability - probability(proposal, token);
if remaining > 0.0 {
residual.push((token, remaining));
}
@@ -2095,27 +2283,113 @@ fn exact_speculative_sample(
if residual.is_empty() {
return (target_best, false);
}
(sample_probabilities(&residual, rng, None), false)
(sample_probabilities_f64(&residual, rng), false)
}
#[cfg(any(target_os = "macos", test))]
struct Rng(u64);
enum Rng {
Ds4(u64),
Qwen { state: u128, increment: u128 },
}
#[cfg(any(target_os = "macos", test))]
impl Rng {
fn new(seed: u64) -> Self {
Self(seed.max(1))
Self::Ds4(seed.max(1))
}
fn new_qwen(seed: u64) -> Self {
const INIT_A: u32 = 0x43b0_d7e5;
const MULT_A: u32 = 0x931e_8875;
const INIT_B: u32 = 0x8b51_f9dd;
const MULT_B: u32 = 0x58f3_8ded;
const MIX_MULT_L: u32 = 0xca01_f9dd;
const MIX_MULT_R: u32 = 0x4973_f715;
fn hash_mix(value: u32, hash: &mut u32) -> u32 {
let mut value = value ^ *hash;
*hash = hash.wrapping_mul(MULT_A);
value = value.wrapping_mul(*hash);
value ^ (value >> 16)
}
fn mix(left: u32, right: u32) -> u32 {
let value = MIX_MULT_L
.wrapping_mul(left)
.wrapping_sub(MIX_MULT_R.wrapping_mul(right));
value ^ (value >> 16)
}
let entropy = [seed as u32, (seed >> 32) as u32];
let mut pool = [0_u32; 4];
let mut hash = INIT_A;
for (index, slot) in pool.iter_mut().enumerate() {
*slot = hash_mix(entropy.get(index).copied().unwrap_or(0), &mut hash);
}
for source in 0..pool.len() {
for destination in 0..pool.len() {
if source != destination {
pool[destination] = mix(pool[destination], hash_mix(pool[source], &mut hash));
}
}
}
let mut words = [0_u32; 8];
hash = INIT_B;
for (index, word) in words.iter_mut().enumerate() {
let mut value = pool[index % pool.len()] ^ hash;
hash = hash.wrapping_mul(MULT_B);
value = value.wrapping_mul(hash);
*word = value ^ (value >> 16);
}
let word64 = |index| u64::from(words[index]) | (u64::from(words[index + 1]) << 32);
let initial_state = u128::from(word64(2)) | (u128::from(word64(0)) << 64);
let stream = u128::from(word64(6)) | (u128::from(word64(4)) << 64);
let increment = stream.wrapping_mul(2) | 1;
let mut rng = Self::Qwen {
state: 0,
increment,
};
rng.advance_qwen();
if let Self::Qwen { state, .. } = &mut rng {
*state = state.wrapping_add(initial_state);
}
rng.advance_qwen();
rng
}
fn advance_qwen(&mut self) {
const MULTIPLIER: u128 = 47_026_247_687_942_121_848_144_207_491_837_523_525;
if let Self::Qwen { state, increment } = self {
*state = state.wrapping_mul(MULTIPLIER).wrapping_add(*increment);
}
}
fn unit_f64(&mut self) -> f64 {
if matches!(self, Self::Qwen { .. }) {
self.advance_qwen();
let Self::Qwen { state, .. } = self else {
unreachable!()
};
let mixed = ((*state >> 64) as u64) ^ (*state as u64);
let raw = mixed.rotate_right((*state >> 122) as u32);
return (raw >> 11) as f64 / 9_007_199_254_740_992.0;
}
f64::from(self.unit())
}
fn unit(&mut self) -> f32 {
let mut value = self.0;
let Self::Ds4(state) = self else {
return self.unit_f64() as f32;
};
let mut value = *state;
if value == 0 {
value = 0x9e37_79b9_7f4a_7c15;
}
value ^= value >> 12;
value ^= value << 25;
value ^= value >> 27;
self.0 = value;
*state = value;
let value = value.wrapping_mul(0x2545_f491_4f6c_dd1d);
((value >> 40) & 0xff_ffff) as f32 / 16_777_216.0
}
@@ -2125,6 +2399,21 @@ impl Rng {
mod sampling_tests {
use super::*;
#[test]
fn qwen_rng_matches_numpy_pcg64() {
let mut rng = Rng::new_qwen(12_345);
let expected = [
0.227_336_022_467_169_66,
0.316_758_339_709_752_87,
0.797_365_457_332_734_1,
0.676_254_670_750_974_6,
0.391_109_550_601_909,
];
for expected in expected {
assert_eq!(rng.unit_f64(), expected);
}
}
#[test]
fn checkpoint_headers_identify_the_exact_model() {
let path =
@@ -2229,6 +2518,23 @@ mod sampling_tests {
);
}
#[test]
fn qwen_sampling_uses_mtplx_global_nucleus_mass_before_top_k() {
let probabilities = qwen_sampling_probabilities(&[0.0, 0.0, 0.0], 1.0, 0.4, 0.0, 2);
assert_eq!(
probabilities
.iter()
.map(|(token, _)| *token)
.collect::<Vec<_>>(),
[0, 1]
);
assert!(
probabilities
.iter()
.all(|(_, probability)| (*probability - 0.5).abs() < f64::EPSILON)
);
}
#[test]
fn split_utf8_token_bytes_are_joined_before_decoding() {
let mut generated = ChatTurn {
src/engine/metal.rs
+19 -2
View File
@@ -13,10 +13,13 @@ use profile::ExpertProfile;
use qwen::QwenExecutor;
use super::gguf::{BF16, F16, F32, Gguf, IQ2_XXS, MXFP4, Q2_K, Q4_K, Q8_0, Tensor as GgufTensor};
#[cfg(test)]
use super::sample_from_logits;
use super::validation::{DsparkConfig, SupportKind, dspark_config};
use super::{
LoadedModel, Model, ModelFamily, ModelRef, Rng, exact_delta_sample, exact_speculative_sample,
sample_from_logits,
LoadedModel, Model, ModelFamily, ModelRef, Rng, exact_delta_sample,
exact_speculative_sample_from_probabilities, qwen_sampling_probabilities,
sample_probabilities_f64,
};
use crate::model::ModelChoice;
use crate::settings::{
@@ -4517,6 +4520,20 @@ impl Executor {
}
}
pub(super) fn take_pending_sampled_token(&mut self) -> Option<i32> {
match self {
Self::Qwen(executor) => executor.take_pending_sampled_token(),
Self::DeepSeek(_) | Self::Glm(_) => None,
}
}
pub(super) fn finalize_pending_sampled_token(&mut self) -> Result<(), String> {
match self {
Self::Qwen(executor) => executor.finalize_pending_sampled_token(),
Self::DeepSeek(_) | Self::Glm(_) => Ok(()),
}
}
pub(super) fn eval(&mut self, token: i32) -> Result<(), String> {
match self {
Self::DeepSeek(executor) => executor.eval(token),
src/engine/metal/gpu.rs
+13
View File
@@ -8,6 +8,7 @@ pub(super) struct GpuTensor {
#[derive(Clone, Copy)]
#[repr(C)]
pub(super) struct QwenWeightView {
pub(super) tensor: *const GpuTensor,
pub(super) map: *const c_void,
pub(super) size: u64,
pub(super) offset: u64,
@@ -90,6 +91,7 @@ pub(super) struct StreamExpertCacheStats {
unsafe extern "C" {
pub(super) fn ds4_gpu_init() -> i32;
pub(super) fn ds4_gpu_cleanup();
pub(super) fn ds4_gpu_metal4_tensor_api_enabled() -> i32;
pub(super) fn ds4_gpu_set_model_map_range(
model_map: *const c_void,
model_size: u64,
@@ -103,6 +105,7 @@ unsafe extern "C" {
a: *const GpuTensor,
b: *const GpuTensor,
c: *const GpuTensor,
d: *const GpuTensor,
weights: *const QwenWeightView,
weight_count: u32,
args: *const QwenKernelArgs,
@@ -132,6 +135,7 @@ unsafe extern "C" {
) -> i32;
pub(super) fn ds4_gpu_set_quality(quality: bool);
pub(super) fn ds4_gpu_set_glm_model(enabled: bool);
pub(super) fn ds4_gpu_set_qwen_model(enabled: bool);
pub(super) fn ds4_gpu_set_ssd_streaming(enabled: bool);
pub(super) fn ds4_gpu_set_model_fd(fd: i32) -> i32;
pub(super) fn ds4_gpu_set_streaming_expert_cache_budget(experts: u32);
@@ -199,6 +203,7 @@ unsafe extern "C" {
pub(super) fn ds4_gpu_parallel_ffn_finish() -> i32;
pub(super) fn ds4_gpu_parallel_ffn_abort();
pub(super) fn ds4_gpu_tensor_alloc(bytes: u64) -> *mut GpuTensor;
pub(super) fn ds4_gpu_tensor_alloc_untracked(bytes: u64) -> *mut GpuTensor;
pub(super) fn ds4_gpu_tensor_view(
base: *const GpuTensor,
offset: u64,
@@ -1847,6 +1852,7 @@ impl Context {
check(unsafe { ds4_gpu_init() }, "Metal initialization")?;
unsafe {
ds4_gpu_set_glm_model(model.shape.family == ModelFamily::Glm);
ds4_gpu_set_qwen_model(false);
ds4_gpu_set_ssd_streaming(ssd_streaming);
// Decode enables this only around DS4's eligible resident pre-M5
// MXFP4 token path; all other work starts from the portable path.
@@ -1944,6 +1950,7 @@ impl Context {
check(unsafe { ds4_gpu_init() }, "Metal initialization")?;
unsafe {
ds4_gpu_set_glm_model(false);
ds4_gpu_set_qwen_model(true);
ds4_gpu_set_ssd_streaming(false);
ds4_gpu_set_decode_pipeline_fast_lookup(0);
}
@@ -2071,6 +2078,12 @@ impl Buffer {
.ok_or_else(|| format!("Metal could not allocate {bytes} bytes"))
}
pub(super) fn untracked(bytes: u64) -> Result<Self, String> {
NonNull::new(unsafe { ds4_gpu_tensor_alloc_untracked(bytes) })
.map(Self)
.ok_or_else(|| format!("Metal could not allocate {bytes} untracked bytes"))
}
pub(super) fn view(&self, offset: u64, bytes: u64) -> Result<Self, String> {
// SAFETY: `ds4_gpu_tensor_view` in native/metal/ds4_metal.m bounds-checks
// the view, ARC-retains base_obj.buffer, and marks the view non-owning, so
src/engine/metal/qwen.rs
+6015 -829
View File
File diff suppressed because it is too large. Load diff
src/engine/qwen.rs
+314 -13
View File
@@ -3,15 +3,19 @@ use super::{ChatTurn, ModelSummary};
use crate::model::ModelChoice;
use crate::settings::ReasoningMode;
use memmap2::{Mmap, MmapOptions};
use rustc_hash::FxHashMap as HashMap;
use serde::de::{MapAccess, Visitor};
use serde::{Deserialize, Deserializer};
use serde_json::Value;
use sha2::{Digest, Sha256};
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::collections::{BTreeMap, BTreeSet, HashSet, VecDeque};
use std::fmt;
use std::fs::{self, File};
use std::io::Read;
use std::os::unix::fs::FileExt;
use std::path::{Path, PathBuf};
use std::sync::mpsc;
use std::thread;
const MAX_SAFETENSORS_HEADER: u64 = 16 * 1024 * 1024;
const MANIFEST: &[u8] = include_bytes!("../../assets/models/qwen38-flash-next-bare-speed.json");
@@ -32,6 +36,9 @@ const GDN_CONV_BYTES: u64 = 2_211_840;
const PLE_CONV_BYTES: u64 = 184_320;
const MTP_CAPTURE_HIDDEN_BYTES: u64 = 10_240 * 4;
const MTP_CAPTURE_LOGITS_BYTES: u64 = 248_320 * 4;
const PLE_ROW_BYTES: usize = 100;
const PLE_READERS: usize = 16;
const PLE_HOT_BYTES: usize = 1024 * 1024 * 1024;
#[derive(Deserialize)]
struct Manifest {
@@ -104,7 +111,9 @@ pub(super) struct ArtifactBindings {
pub(super) struct QwenMap {
path: PathBuf,
map: Mmap,
file: File,
bytes: u64,
map: Option<Mmap>,
}
#[derive(Clone)]
@@ -123,9 +132,227 @@ pub(super) struct QwenModel {
memory: MemoryPlan,
maps: Vec<QwenMap>,
tensors: HashMap<String, QwenTensor>,
ple_reader: PleReader,
identity: [u8; 32],
}
struct PleTask {
row: u64,
response: mpsc::Sender<(u64, Result<[u8; PLE_ROW_BYTES], String>)>,
}
struct PleReader {
workers: Vec<mpsc::Sender<Option<PleTask>>>,
threads: Vec<thread::JoinHandle<()>>,
hot: HashMap<u64, ([u8; PLE_ROW_BYTES], u64)>,
order: VecDeque<(u64, u64)>,
generation: u64,
capacity: usize,
hits: u64,
misses: u64,
evictions: u64,
read_bytes: u64,
read_ns: u64,
}
#[derive(Clone, Copy, Default)]
pub(super) struct PleStats {
pub(super) entries: u64,
pub(super) cache_bytes: u64,
pub(super) hits: u64,
pub(super) misses: u64,
pub(super) evictions: u64,
pub(super) read_bytes: u64,
pub(super) read_ms: u64,
}
impl PleReader {
fn new(maps: &[QwenMap], tensors: &HashMap<String, QwenTensor>) -> Result<Self, String> {
let packed = tensors
.get("ngram.weight")
.ok_or_else(|| "Qwen tensor is missing: ngram.weight".to_owned())?;
let scales = tensors
.get("ngram.scales")
.ok_or_else(|| "Qwen tensor is missing: ngram.scales".to_owned())?;
let biases = tensors
.get("ngram.biases")
.ok_or_else(|| "Qwen tensor is missing: ngram.biases".to_owned())?;
if scales.map != packed.map || biases.map != packed.map {
return Err("Qwen PLE tensors do not share their sidecar file".into());
}
let ranges = [
(packed.range.clone(), 80_u64),
(scales.range.clone(), 10_u64),
(biases.range.clone(), 10_u64),
];
let mut workers = Vec::with_capacity(PLE_READERS);
let mut threads = Vec::with_capacity(PLE_READERS);
for index in 0..PLE_READERS {
let file = maps[packed.map]
.file
.try_clone()
.map_err(|error| format!("cloning the Qwen PLE sidecar: {error}"))?;
let ranges = ranges.clone();
let (sender, receiver) = mpsc::channel::<Option<PleTask>>();
let handle = thread::Builder::new()
.name(format!("qwen-ple-{index}"))
.spawn(move || {
while let Ok(Some(task)) = receiver.recv() {
let result = read_ple_row(&file, &ranges, task.row);
let _ = task.response.send((task.row, result));
}
})
.map_err(|error| format!("starting a Qwen PLE reader: {error}"))?;
workers.push(sender);
threads.push(handle);
}
Ok(Self {
workers,
threads,
hot: HashMap::default(),
order: VecDeque::new(),
generation: 0,
capacity: PLE_HOT_BYTES / PLE_ROW_BYTES,
hits: 0,
misses: 0,
evictions: 0,
read_bytes: 0,
read_ns: 0,
})
}
fn rows(&mut self, rows: &[u64]) -> Result<Vec<[u8; PLE_ROW_BYTES]>, String> {
let mut unique = Vec::with_capacity(rows.len());
let mut seen = HashSet::with_capacity(rows.len());
for &row in rows {
if seen.insert(row) {
unique.push(row);
}
}
let missing = unique
.iter()
.copied()
.filter(|row| !self.hot.contains_key(row))
.collect::<Vec<_>>();
self.hits = self
.hits
.saturating_add((unique.len() - missing.len()) as u64);
self.misses = self.misses.saturating_add(missing.len() as u64);
if !missing.is_empty() {
let started = std::time::Instant::now();
let (response, received) = mpsc::channel();
for (index, row) in missing.iter().copied().enumerate() {
self.workers[index % self.workers.len()]
.send(Some(PleTask {
row,
response: response.clone(),
}))
.map_err(|_| "Qwen PLE reader stopped unexpectedly".to_owned())?;
}
drop(response);
for _ in 0..missing.len() {
let (row, result) = received
.recv()
.map_err(|_| "Qwen PLE reader returned too few rows".to_owned())?;
self.insert(row, result?);
}
self.read_bytes = self
.read_bytes
.saturating_add((missing.len() * PLE_ROW_BYTES) as u64);
self.read_ns = self
.read_ns
.saturating_add(u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX));
}
rows.iter()
.map(|row| {
let value = self
.hot
.get(row)
.map(|(value, _)| *value)
.ok_or_else(|| format!("Qwen PLE row {row} was not cached"))?;
self.touch(*row);
Ok(value)
})
.collect()
}
fn insert(&mut self, row: u64, value: [u8; PLE_ROW_BYTES]) {
self.generation = self.generation.wrapping_add(1);
self.hot.insert(row, (value, self.generation));
self.order.push_back((row, self.generation));
while self.hot.len() > self.capacity {
let Some((old, generation)) = self.order.pop_front() else {
break;
};
if self
.hot
.get(&old)
.is_some_and(|entry| entry.1 == generation)
{
self.hot.remove(&old);
self.evictions = self.evictions.saturating_add(1);
}
}
}
fn touch(&mut self, row: u64) {
self.generation = self.generation.wrapping_add(1);
if let Some(entry) = self.hot.get_mut(&row) {
entry.1 = self.generation;
self.order.push_back((row, self.generation));
}
}
fn stats(&self) -> PleStats {
PleStats {
entries: self.hot.len() as u64,
cache_bytes: (self.hot.len() * PLE_ROW_BYTES) as u64,
hits: self.hits,
misses: self.misses,
evictions: self.evictions,
read_bytes: self.read_bytes,
read_ms: self.read_ns / 1_000_000,
}
}
}
impl Drop for PleReader {
fn drop(&mut self) {
for worker in &self.workers {
let _ = worker.send(None);
}
for thread in self.threads.drain(..) {
let _ = thread.join();
}
}
}
fn read_ple_row(
file: &File,
ranges: &[(std::ops::Range<u64>, u64); 3],
row: u64,
) -> Result<[u8; PLE_ROW_BYTES], String> {
let mut output = [0_u8; PLE_ROW_BYTES];
let mut cursor = 0;
for (range, width) in ranges {
let offset = range
.start
.checked_add(
row.checked_mul(*width)
.ok_or_else(|| "Qwen PLE row offset overflows".to_owned())?,
)
.ok_or_else(|| "Qwen PLE row offset overflows".to_owned())?;
if offset + width > range.end {
return Err(format!("Qwen PLE row {row} is truncated"));
}
let end = cursor + *width as usize;
file.read_exact_at(&mut output[cursor..end], offset)
.map_err(|error| format!("reading Qwen PLE row {row}: {error}"))?;
cursor = end;
}
Ok(output)
}
impl QwenModel {
#[cfg(test)]
pub(super) fn open(root: &Path, context: u32) -> Result<Self, String> {
@@ -150,14 +377,28 @@ impl QwenModel {
paths.sort();
paths.dedup();
let mut maps = Vec::with_capacity(paths.len());
let mut map_indices = HashMap::with_capacity(paths.len());
let mut map_indices = HashMap::with_capacity_and_hasher(paths.len(), Default::default());
for path in paths {
let file = File::open(&path).map_err(|error| format!("{}: {error}", path.display()))?;
// SAFETY: verified managed artifacts remain read-only while the model owns each mapping.
let map = unsafe { MmapOptions::new().map(&file) }
.map_err(|error| format!("cannot map {}: {error}", path.display()))?;
if path
.file_name()
.is_some_and(|name| name == "ngram-table.safetensors")
{
// SAFETY: the read-only mapping remains valid for this advisory call.
unsafe {
libc::madvise(map.as_ptr().cast_mut().cast(), map.len(), libc::MADV_RANDOM);
}
}
map_indices.insert(path.clone(), maps.len());
maps.push(QwenMap { path, map });
maps.push(QwenMap {
path,
file,
bytes: map.len() as u64,
map: Some(map),
});
}
let tensors = bindings
.into_iter()
@@ -174,6 +415,7 @@ impl QwenModel {
(binding.name, tensor)
})
.collect::<HashMap<_, _>>();
let ple_reader = PleReader::new(&maps, &tensors)?;
let mut hash = Sha256::new();
hash.update(b"DS4Server Qwen3.8 checkpoint identity v1");
hash.update(MANIFEST);
@@ -183,6 +425,7 @@ impl QwenModel {
memory: loaded.memory,
maps,
tensors,
ple_reader,
identity,
})
}
@@ -193,12 +436,25 @@ impl QwenModel {
.ok_or_else(|| format!("Qwen tensor is missing: {name}"))
}
pub(super) fn tensors(&self) -> impl Iterator<Item = &QwenTensor> {
self.tensors.values()
}
pub(super) fn map(&self, index: usize) -> (&[u8], &Path) {
(&self.maps[index].map, &self.maps[index].path)
(
self.maps[index]
.map
.as_deref()
.expect("materialized Qwen weights must not use their released source mapping"),
&self.maps[index].path,
)
}
pub(super) fn tensor_bytes<'a>(&'a self, tensor: &QwenTensor) -> Result<&'a [u8], String> {
let map = &self.maps[tensor.map].map;
let map = self.maps[tensor.map]
.map
.as_deref()
.ok_or_else(|| format!("{} source mapping was already released", tensor.name))?;
let start = usize::try_from(tensor.range.start)
.map_err(|_| format!("{} starts beyond this platform", tensor.name))?;
let end = usize::try_from(tensor.range.end)
@@ -207,6 +463,48 @@ impl QwenModel {
.ok_or_else(|| format!("{} is outside its mapped artifact", tensor.name))
}
pub(super) fn read_tensor_at(
&self,
tensor: &QwenTensor,
offset: u64,
output: &mut [u8],
) -> Result<(), String> {
let bytes = tensor.range.end - tensor.range.start;
let end = offset
.checked_add(output.len() as u64)
.ok_or_else(|| format!("{} read offset overflows", tensor.name))?;
if end > bytes {
return Err(format!("{} read exceeds its tensor range", tensor.name));
}
self.maps[tensor.map]
.file
.read_exact_at(output, tensor.range.start + offset)
.map_err(|error| format!("reading {} at byte {offset}: {error}", tensor.name))
}
pub(super) fn release_materialized_mappings(&mut self) {
for (index, map) in self.maps.iter_mut().enumerate() {
if !self
.tensors
.values()
.any(|tensor| tensor.map == index && tensor.name.starts_with("ngram."))
{
map.map = None;
}
}
}
pub(super) fn read_ple_rows(
&mut self,
rows: &[u64],
) -> Result<Vec<[u8; PLE_ROW_BYTES]>, String> {
self.ple_reader.rows(rows)
}
pub(super) fn ple_stats(&self) -> PleStats {
self.ple_reader.stats()
}
pub(super) fn checkpoint_identity(&self) -> [u8; 32] {
self.identity
}
@@ -214,7 +512,7 @@ impl QwenModel {
pub(super) fn summary(&self) -> ModelSummary {
ModelSummary {
model: ModelChoice::Qwen38FlashNext,
mapped_bytes: self.maps.iter().map(|item| item.map.len() as u64).sum(),
mapped_bytes: self.maps.iter().map(|item| item.bytes).sum(),
tensor_count: self.tensors.len(),
vocabulary_size: self.tokenizer.vocab_size(),
support_loaded: false,
@@ -284,12 +582,15 @@ impl QwenModel {
let mut core = 0_u64;
let mut ple = 0_u64;
for item in &self.maps {
let mut pages = vec![0_i8; item.map.len().div_ceil(page)];
let Some(map) = &item.map else {
continue;
};
let mut pages = vec![0_i8; map.len().div_ceil(page)];
// SAFETY: each read-only mmap and residency vector remain valid for this call.
if unsafe {
libc::mincore(
item.map.as_ptr().cast_mut().cast(),
item.map.len(),
map.as_ptr().cast_mut().cast(),
map.len(),
pages.as_mut_ptr(),
)
} != 0
@@ -301,7 +602,7 @@ impl QwenModel {
));
}
let bytes = (pages.iter().filter(|value| **value & 1 != 0).count() * page)
.min(item.map.len()) as u64;
.min(map.len()) as u64;
if item
.path
.file_name()
@@ -846,8 +1147,8 @@ mod tests {
assert_eq!(
direct,
[
248_045, 846, 198, 12_675, 248_046, 198, 248_045, 74_455, 198, 13_314, 741, 29,
271, 510, 26_003, 29, 271,
248_045, 846, 198, 12_675, 248_046, 198, 248_045, 74_455, 198, 248_068, 271,
248_069, 271,
]
);
assert_eq!(
src/engine/tokenizer.rs
+1 -4
View File
@@ -62,7 +62,6 @@ struct QwenTokenizerFile {
struct QwenAddedToken {
id: usize,
content: String,
special: bool,
}
#[derive(Deserialize)]
@@ -225,9 +224,7 @@ impl Tokenizer {
));
}
tokens[token.id] = token.content.as_bytes().to_vec();
if token.special {
rendered_specials.push((token.content.into_bytes(), token.id as i32));
}
rendered_specials.push((token.content.into_bytes(), token.id as i32));
}
if tokens[..=maximum].iter().any(Vec::is_empty) {
return Err("Qwen tokenizer token ids are not contiguous".into());