811 lines
28 KiB
Rust
811 lines
28 KiB
Rust
//! Qwen4ExpTextModel/TextModel graph orchestration from the pinned MTPLX.
|
|
//! GDN runs use the reference simplify/fuse passes. Synthetic run graphs are
|
|
//! checked; real-weight and serving integration remain outside this test port.
|
|
use super::array::{Array, Dtype};
|
|
use super::binary::{Binary, binary};
|
|
use super::decoder::{Attention, AttentionCache, Cache, DecoderLayer, Options};
|
|
use super::mlp::Linear;
|
|
use super::stream::{Stream, Streams};
|
|
use super::{cache::QsaCache, compiled::Replay, gdn, hyper, indexing, linear, ops, ple};
|
|
use std::collections::HashMap;
|
|
use std::ops::Range;
|
|
|
|
#[derive(Clone, Copy, Default)]
|
|
pub(super) struct CompilePolicy {
|
|
explicit_off: bool,
|
|
environment: bool,
|
|
}
|
|
|
|
impl CompilePolicy {
|
|
pub(super) fn from_values(gdn: Option<&str>, family: Option<&str>) -> Self {
|
|
let truthy = |s: &str| {
|
|
matches!(
|
|
s.trim().to_ascii_lowercase().as_str(),
|
|
"1" | "true" | "yes" | "on"
|
|
)
|
|
};
|
|
let explicit_off = [gdn, family].into_iter().flatten().any(|s| !truthy(s));
|
|
Self {
|
|
explicit_off,
|
|
environment: !explicit_off && [gdn, family].into_iter().flatten().any(truthy),
|
|
}
|
|
}
|
|
|
|
#[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
|
|
pub(super) fn from_env() -> Self {
|
|
Self::from_values(
|
|
std::env::var("MTPLX_COMPILED_GDN").ok().as_deref(),
|
|
std::env::var("MTPLX_QWEN4EXP_COMPILE").ok().as_deref(),
|
|
)
|
|
}
|
|
|
|
fn enabled(self) -> bool {
|
|
!self.explicit_off && self.environment
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, PartialEq, Eq)]
|
|
enum Run {
|
|
Eager(usize),
|
|
Gdn(Range<usize>),
|
|
}
|
|
|
|
fn runs(layers: impl IntoIterator<Item = (bool, bool)>) -> Vec<Run> {
|
|
let mut result = Vec::new();
|
|
let mut start = None;
|
|
let mut end = 0;
|
|
for (i, (linear, ple)) in layers.into_iter().enumerate() {
|
|
end = i + 1;
|
|
if linear && !ple {
|
|
start.get_or_insert(i);
|
|
} else {
|
|
if let Some(start) = start.take() {
|
|
result.push(Run::Gdn(start..i));
|
|
}
|
|
result.push(Run::Eager(i));
|
|
}
|
|
}
|
|
if let Some(start) = start {
|
|
result.push(Run::Gdn(start..end));
|
|
}
|
|
result
|
|
}
|
|
|
|
#[derive(Hash, PartialEq, Eq)]
|
|
struct RunKey {
|
|
range: Range<usize>,
|
|
capture: bool,
|
|
stream: Stream,
|
|
signature: Vec<(Vec<i32>, Dtype)>,
|
|
}
|
|
|
|
pub(super) struct TextModel<'a> {
|
|
/// Original Module.parameters tree order, populated by checkpoint binding.
|
|
pub(super) parameter_order: Vec<String>,
|
|
pub(super) embedding: Linear<'a>,
|
|
pub(super) layers: Vec<DecoderLayer<'a>>,
|
|
pub(super) mixer: hyper::GatedResidual<'a>,
|
|
pub(super) head: Option<Linear<'a>>, // None means tied embedding weights.
|
|
pub(super) mtp: Option<super::mtp::Mtp<'a>>,
|
|
pub(super) vision: Option<super::vision::Input>,
|
|
pub(super) last_widened: Option<Array>,
|
|
policy: CompilePolicy,
|
|
decode_runs: Option<Vec<Run>>,
|
|
replays: HashMap<RunKey, Replay>,
|
|
commit_refusals: usize,
|
|
}
|
|
|
|
/// ArraysCache.state before verification; QSA entries need no snapshot here.
|
|
pub(super) type ArraySnapshot = Vec<Option<Array>>;
|
|
|
|
#[derive(Clone, Copy)]
|
|
pub(super) struct OutputOptions {
|
|
pub(super) emit_logits: bool,
|
|
pub(super) logits_keep: i32,
|
|
pub(super) return_hidden: bool,
|
|
}
|
|
|
|
pub(super) struct Output {
|
|
pub(super) logits: Option<Array>,
|
|
pub(super) hidden: Option<Array>,
|
|
}
|
|
|
|
pub(super) fn embedding(layer: Linear<'_>, ids: &Array, stream: Stream) -> Result<Array, String> {
|
|
match layer {
|
|
Linear::Dense(weight) => indexing::take_rows(weight, ids, stream),
|
|
Linear::Quantized(layer) => linear::dequantize(
|
|
&indexing::take_rows(&layer.weight, ids, stream)?,
|
|
&indexing::take_rows(&layer.scales, ids, stream)?,
|
|
&indexing::take_rows(&layer.biases, ids, stream)?,
|
|
layer.group,
|
|
layer.bits,
|
|
stream,
|
|
),
|
|
}
|
|
}
|
|
|
|
fn mask(cache: Option<&Cache>, rows: i32, stream: Stream) -> Result<Option<Array>, String> {
|
|
let Some(Cache {
|
|
attention: AttentionCache::Linear(cache),
|
|
..
|
|
}) = cache
|
|
else {
|
|
return Ok(None);
|
|
};
|
|
let (metadata, op) = if let Some(padding) = &cache.left_padding {
|
|
(padding, Binary::GreaterEqual)
|
|
} else if let Some(lengths) = &cache.lengths {
|
|
(lengths, Binary::Less)
|
|
} else {
|
|
return Ok(None);
|
|
};
|
|
binary(
|
|
&indexing::arange_i32(rows, stream)?,
|
|
&ops::expand_dims(metadata, &[-1], stream)?,
|
|
op,
|
|
stream,
|
|
)
|
|
.map(Some)
|
|
}
|
|
|
|
impl<'a> TextModel<'a> {
|
|
#[cfg(test)]
|
|
pub(super) fn replay_count(&self) -> usize {
|
|
self.replays.len()
|
|
}
|
|
pub(super) fn new(
|
|
embedding: Linear<'a>,
|
|
layers: Vec<DecoderLayer<'a>>,
|
|
mixer: hyper::GatedResidual<'a>,
|
|
head: Option<Linear<'a>>,
|
|
policy: CompilePolicy,
|
|
) -> Result<Self, String> {
|
|
if layers.is_empty() || mixer.inject.is_some() {
|
|
return Err("Qwen requires decoder layers and a non-combining final mixer".into());
|
|
}
|
|
Ok(Self {
|
|
embedding,
|
|
layers,
|
|
mixer,
|
|
head,
|
|
parameter_order: Vec::new(),
|
|
mtp: None,
|
|
vision: None,
|
|
last_widened: None,
|
|
policy,
|
|
decode_runs: None,
|
|
replays: HashMap::new(),
|
|
commit_refusals: 0,
|
|
})
|
|
}
|
|
|
|
pub(super) fn make_cache(&self, ratio: i32) -> Vec<Option<Cache>> {
|
|
self.layers
|
|
.iter()
|
|
.map(|layer| {
|
|
Some(Cache {
|
|
attention: match layer.attention {
|
|
Attention::Linear(_) => AttentionCache::Linear(gdn::Cache::default()),
|
|
Attention::Qsa(_) => {
|
|
AttentionCache::Qsa(QsaCache::new(if ratio == 0 { 4 } else { ratio }))
|
|
}
|
|
},
|
|
ple: ple::Cache::default(),
|
|
capture_ple: None,
|
|
})
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
pub(super) fn make_mtp_cache(&self, ratio: i32) -> [Option<Cache>; 1] {
|
|
super::mtp::make_cache(ratio)
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub(super) fn mtp_forward(
|
|
&mut self,
|
|
hidden: &Array,
|
|
ids: &Array,
|
|
cache: Option<&mut [Option<Cache>]>,
|
|
return_hidden: bool,
|
|
options: &Options,
|
|
vision: Option<(Option<&Array>, i32)>,
|
|
streams: &Streams,
|
|
stream: Stream,
|
|
) -> Result<Output, String> {
|
|
let image = self.vision.clone();
|
|
let vision = vision.or_else(|| image.as_ref().map(|v| (Some(&v.table), v.delta)));
|
|
let emb = embedding(self.embedding, ids, stream)?;
|
|
let mtp = self
|
|
.mtp
|
|
.as_mut()
|
|
.ok_or("Qwen MTP weights are not attached")?;
|
|
let h = mtp.fuse_and_run(hidden, &emb, cache, options, vision, streams, stream)?;
|
|
let mixed = mtp
|
|
.mixer
|
|
.forward(&h, options.hyper_fused, options.hyper_v3, streams, stream)?
|
|
.mixed;
|
|
let logits = self.head.unwrap_or(self.embedding).apply(&mixed, stream)?;
|
|
Ok(Output {
|
|
logits: Some(logits),
|
|
hidden: return_hidden.then_some(h),
|
|
})
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub(super) fn mtp_update_cache(
|
|
&mut self,
|
|
hidden: &Array,
|
|
ids: &Array,
|
|
cache: Option<&mut [Option<Cache>]>,
|
|
input_embeddings: Option<&Array>,
|
|
options: &Options,
|
|
vision: Option<(Option<&Array>, i32)>,
|
|
streams: &Streams,
|
|
stream: Stream,
|
|
) -> Result<Array, String> {
|
|
let image = self.vision.clone();
|
|
let vision = vision.or_else(|| image.as_ref().map(|v| (Some(&v.table), v.delta)));
|
|
let spliced = match (&image, input_embeddings) {
|
|
(Some(v), None) => v.embeddings(
|
|
self.embedding,
|
|
ids,
|
|
super::vision::cache_position(cache.as_deref()) + 1,
|
|
stream,
|
|
)?,
|
|
_ => None,
|
|
};
|
|
let emb = match input_embeddings.or(spliced.as_ref()) {
|
|
Some(emb) => emb.clone(),
|
|
None => embedding(self.embedding, ids, stream)?,
|
|
};
|
|
self.mtp
|
|
.as_mut()
|
|
.ok_or("Qwen MTP weights are not attached")?
|
|
.fuse_and_run(hidden, &emb, cache, options, vision, streams, stream)
|
|
}
|
|
|
|
/// Model.post_weight_load: attach each PLE table after trunk weight loading.
|
|
pub(super) fn post_weight_load(
|
|
&mut self,
|
|
root: &std::path::Path,
|
|
streams: &Streams,
|
|
stream: Stream,
|
|
) -> Result<(), String> {
|
|
let path = root.join("ngram-table.safetensors");
|
|
if !path.exists() {
|
|
return Ok(());
|
|
}
|
|
let resident = super::ngram::resident_policy();
|
|
let mut sidecar = None;
|
|
for layer in &mut self.layers {
|
|
let Some(ple) = &mut layer.ple else { continue };
|
|
let table = &mut ple.embedding.table;
|
|
table.attach_sidecar(&path)?;
|
|
if resident {
|
|
// Original attach_resident reports failure and retains the SSD fallback.
|
|
table.attach_resident(&path, streams, stream);
|
|
}
|
|
sidecar = table.sidecar.as_ref();
|
|
}
|
|
if !resident && let Some(sidecar) = sidecar {
|
|
println!(
|
|
"[qwen4_exp] ngram table: streamed from SSD (file-backed pages, hot cache {}M, prefetch {}) — resident only at >=160G RAM or MTPLX_NGRAM_RESIDENT=1",
|
|
sidecar.hot_mebibytes(),
|
|
if sidecar.prefetch_enabled() {
|
|
"on"
|
|
} else {
|
|
"off"
|
|
}
|
|
);
|
|
let _ = std::io::Write::flush(&mut std::io::stdout());
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
|
|
pub(super) fn clear_verify_capture(cache: &mut [Option<Cache>]) {
|
|
for cache in cache.iter_mut().flatten() {
|
|
cache.capture_ple = None;
|
|
if let AttentionCache::Linear(c) = &mut cache.attention {
|
|
c.capture = None;
|
|
}
|
|
}
|
|
}
|
|
|
|
fn refuse_commit(&mut self, layer: usize, reason: &str) -> bool {
|
|
if self.commit_refusals < 3 {
|
|
use std::io::Write;
|
|
println!(
|
|
"[qwen4_exp] capture-commit refused (layer {layer}: {reason}) — falling back to re-forward"
|
|
);
|
|
let _ = std::io::stdout().flush();
|
|
}
|
|
self.commit_refusals += 1;
|
|
false
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub(super) fn commit_verified_window(
|
|
&mut self,
|
|
cache: &mut [Option<Cache>],
|
|
snapshots: Option<&[Option<ArraySnapshot>]>,
|
|
keep: i32,
|
|
verified: i32,
|
|
options: &Options,
|
|
vision: Option<(Option<&Array>, i32)>,
|
|
streams: &Streams,
|
|
stream: Stream,
|
|
) -> Result<bool, String> {
|
|
if keep < 1 || keep > verified || cache.len() != self.layers.len() {
|
|
return Ok(false);
|
|
}
|
|
// First validate EVERY layer, without cache mutations or GPU work.
|
|
// A refusal must leave the full verify window available for fallback.
|
|
for (i, (layer, entry)) in self.layers.iter().zip(cache.iter()).enumerate() {
|
|
let Some(entry) = entry else {
|
|
return Ok(self.refuse_commit(i, "entry_missing"));
|
|
};
|
|
let AttentionCache::Linear(c) = &entry.attention else {
|
|
continue;
|
|
};
|
|
let snapshot = snapshots.and_then(|s| s.get(i)).and_then(Option::as_ref);
|
|
let validation = if layer.ple.is_some() {
|
|
match (snapshot, &entry.capture_ple) {
|
|
(None, _) => Err("snapshot_missing".to_owned()),
|
|
(_, None) => Err("ple_rows_missing".to_owned()),
|
|
(Some(_), Some((h, _))) if h.layout().shape()[1] != verified => Err(format!(
|
|
"ple_rows_width_{}_vs_{verified}",
|
|
h.layout().shape()[1]
|
|
)),
|
|
(Some(pre), _) if pre.len() < 4 => Err("ple_snapshot_short".to_owned()),
|
|
_ => Ok(()),
|
|
}
|
|
} else {
|
|
c.validate_verified_window(snapshot.map(Vec::as_slice), verified)
|
|
};
|
|
if let Err(reason) = validation {
|
|
return Ok(self.refuse_commit(i, &reason));
|
|
}
|
|
}
|
|
for (i, (layer, entry)) in self.layers.iter_mut().zip(cache).enumerate() {
|
|
let entry = entry.as_mut().unwrap();
|
|
if let AttentionCache::Qsa(c) = &mut entry.attention {
|
|
if keep != verified {
|
|
c.trim(verified - keep);
|
|
}
|
|
continue;
|
|
}
|
|
let snapshot = snapshots.unwrap()[i].as_ref().unwrap();
|
|
if layer.ple.is_some() {
|
|
let (hidden, ids) = entry.capture_ple.as_ref().unwrap();
|
|
let prefix = |a: &Array| {
|
|
let mut end = a.layout().shape().to_vec();
|
|
end[1] = keep;
|
|
ops::slice(a, &vec![0; end.len()], &end, &vec![1; end.len()], stream)
|
|
};
|
|
let hidden = prefix(hidden)?;
|
|
let ids = prefix(ids)?;
|
|
let AttentionCache::Linear(c) = &mut entry.attention else {
|
|
unreachable!()
|
|
};
|
|
c.conv = snapshot[0].clone();
|
|
c.delta = snapshot[1].clone();
|
|
entry.ple.conv = snapshot[2].clone();
|
|
entry.ple.history = snapshot[3].clone();
|
|
layer.ple.as_mut().unwrap().embedding.stage(
|
|
&ids,
|
|
entry.ple.history.as_ref(),
|
|
&options.stage,
|
|
streams,
|
|
stream,
|
|
);
|
|
layer.call(
|
|
&hidden,
|
|
Some(&ids),
|
|
None,
|
|
Some(entry),
|
|
options,
|
|
vision,
|
|
streams,
|
|
stream,
|
|
)?;
|
|
entry.capture_ple = None;
|
|
} else {
|
|
let AttentionCache::Linear(c) = &mut entry.attention else {
|
|
unreachable!()
|
|
};
|
|
let Attention::Linear(weights) = &layer.attention else {
|
|
return Err("GDN capture-commit layer/cache mismatch".into());
|
|
};
|
|
c.replay_verified_prefix(snapshot, keep, weights.a_log, weights.dt_bias, stream)?;
|
|
}
|
|
}
|
|
Ok(true)
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub(super) fn call(
|
|
&mut self,
|
|
ids: &Array,
|
|
cache: Option<&mut [Option<Cache>]>,
|
|
input_embeddings: Option<&Array>,
|
|
output: OutputOptions,
|
|
options: &Options,
|
|
vision: Option<(Option<&Array>, i32)>,
|
|
streams: &Streams,
|
|
stream: Stream,
|
|
) -> Result<Output, String> {
|
|
let image = self.vision.clone();
|
|
let vision = vision.or_else(|| image.as_ref().map(|v| (Some(&v.table), v.delta)));
|
|
let spliced = match (&image, input_embeddings) {
|
|
(Some(v), None) => v.embeddings(
|
|
self.embedding,
|
|
ids,
|
|
super::vision::cache_position(cache.as_deref()),
|
|
stream,
|
|
)?,
|
|
_ => None,
|
|
};
|
|
let h = match input_embeddings.or(spliced.as_ref()) {
|
|
Some(h) => h.clone(),
|
|
None => embedding(self.embedding, ids, stream)?,
|
|
};
|
|
let shape = h.layout().shape().to_vec();
|
|
if shape.len() != 3
|
|
|| shape[0] < 1
|
|
|| shape[1] < 1
|
|
|| shape[2] != 2560
|
|
|| ids.layout().shape() != &shape[..2]
|
|
{
|
|
return Err(
|
|
"installed Qwen requires [B,S,2560] embeddings and matching token IDs".into(),
|
|
);
|
|
}
|
|
let mut empty = Vec::new();
|
|
let cache = match cache {
|
|
Some(cache) => cache,
|
|
None => {
|
|
empty.resize_with(self.layers.len(), || None);
|
|
&mut empty
|
|
}
|
|
};
|
|
if cache.len() != self.layers.len() {
|
|
return Err("Qwen cache/layer count mismatch".into());
|
|
}
|
|
let ssm = self
|
|
.layers
|
|
.iter()
|
|
.position(|l| matches!(l.attention, Attention::Linear(_)))
|
|
.unwrap_or(0);
|
|
let ssm_mask = mask(cache[ssm].as_ref(), shape[1], stream)?;
|
|
if let Some(i) = self.layers.iter().position(|l| l.ple.is_some()) {
|
|
self.layers[i].ple.as_mut().unwrap().embedding.stage(
|
|
ids,
|
|
cache[i].as_ref().and_then(|c| c.ple.history.as_ref()),
|
|
&options.stage,
|
|
streams,
|
|
stream,
|
|
);
|
|
}
|
|
// tile(h, (1,1,4)): original reshape -> broadcast -> reshape.
|
|
let h = ops::reshape(&h, &[shape[0], shape[1], 1, 2560], stream)?;
|
|
let h = ops::broadcast_to(&h, &[shape[0], shape[1], 4, 2560], stream)?;
|
|
let mut h = ops::reshape(&h, &[shape[0], shape[1], 10240], stream)?;
|
|
if (1..=4).contains(&shape[1])
|
|
&& ssm_mask.is_none()
|
|
&& self.policy.enabled()
|
|
&& cache[ssm].is_some()
|
|
{
|
|
h = self.decode(h, ids, cache, options, vision, streams, stream)?;
|
|
} else {
|
|
for (layer, cache) in self.layers.iter_mut().zip(cache) {
|
|
h = eager(
|
|
layer,
|
|
&h,
|
|
ids,
|
|
ssm_mask.as_ref(),
|
|
cache.as_mut(),
|
|
options,
|
|
vision,
|
|
streams,
|
|
stream,
|
|
)?;
|
|
}
|
|
}
|
|
self.last_widened = Some(h.clone());
|
|
let mixed = self
|
|
.mixer
|
|
.forward(&h, options.hyper_fused, options.hyper_v3, streams, stream)?
|
|
.mixed;
|
|
let logits = if output.emit_logits {
|
|
let mixed = if output.logits_keep != 0 {
|
|
let keep = output.logits_keep.max(1).min(shape[1]);
|
|
ops::slice(
|
|
&mixed,
|
|
&[0, shape[1] - keep, 0],
|
|
&[shape[0], shape[1], 2560],
|
|
&[1; 3],
|
|
stream,
|
|
)?
|
|
} else {
|
|
mixed
|
|
};
|
|
Some(self.head.unwrap_or(self.embedding).apply(&mixed, stream)?)
|
|
} else {
|
|
None
|
|
};
|
|
Ok(Output {
|
|
logits,
|
|
hidden: output.return_hidden.then_some(h),
|
|
})
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)]
|
|
fn decode(
|
|
&mut self,
|
|
mut h: Array,
|
|
ids: &Array,
|
|
cache: &mut [Option<Cache>],
|
|
options: &Options,
|
|
vision: Option<(Option<&Array>, i32)>,
|
|
streams: &Streams,
|
|
stream: Stream,
|
|
) -> Result<Array, String> {
|
|
let runs = self.decode_runs.get_or_insert_with(|| {
|
|
runs(
|
|
self.layers
|
|
.iter()
|
|
.map(|l| (matches!(l.attention, Attention::Linear(_)), l.ple.is_some())),
|
|
)
|
|
});
|
|
for run in runs {
|
|
let range = match run {
|
|
Run::Eager(i) => {
|
|
h = eager(
|
|
&mut self.layers[*i],
|
|
&h,
|
|
ids,
|
|
None,
|
|
cache[*i].as_mut(),
|
|
options,
|
|
vision,
|
|
streams,
|
|
stream,
|
|
)?;
|
|
continue;
|
|
}
|
|
Run::Gdn(range) => range.clone(),
|
|
};
|
|
let mut flat = vec![h.clone()];
|
|
for c in &cache[range.clone()] {
|
|
let Some(Cache {
|
|
attention: AttentionCache::Linear(c),
|
|
..
|
|
}) = c
|
|
else {
|
|
break;
|
|
};
|
|
let (Some(conv), Some(delta)) = (&c.conv, &c.delta) else {
|
|
break;
|
|
};
|
|
flat.extend([conv.clone(), delta.clone()]);
|
|
}
|
|
if flat.len() != 1 + 2 * range.len() {
|
|
for i in range {
|
|
h = eager(
|
|
&mut self.layers[i],
|
|
&h,
|
|
ids,
|
|
None,
|
|
cache[i].as_mut(),
|
|
options,
|
|
vision,
|
|
streams,
|
|
stream,
|
|
)?;
|
|
}
|
|
continue;
|
|
}
|
|
let key = RunKey {
|
|
range: range.clone(),
|
|
capture: options.gdn.capture,
|
|
stream,
|
|
signature: flat
|
|
.iter()
|
|
.map(|a| (a.layout().shape().to_vec(), a.layout().dtype()))
|
|
.collect(),
|
|
};
|
|
let replay = match self.replays.entry(key) {
|
|
std::collections::hash_map::Entry::Occupied(entry) => entry.into_mut(),
|
|
std::collections::hash_map::Entry::Vacant(entry) => {
|
|
let replay = Replay::trace_compiled(&flat, |args| {
|
|
let mut h = args[0].clone();
|
|
let mut states = Vec::new();
|
|
let mut captured = Vec::new();
|
|
for (n, i) in range.clone().enumerate() {
|
|
let mut c = Cache {
|
|
attention: AttentionCache::Linear(gdn::Cache {
|
|
conv: Some(args[1 + 2 * n].clone()),
|
|
delta: Some(args[2 + 2 * n].clone()),
|
|
..Default::default()
|
|
}),
|
|
ple: ple::Cache::default(),
|
|
capture_ple: None,
|
|
};
|
|
h = self.layers[i].call(
|
|
&h,
|
|
None,
|
|
None,
|
|
Some(&mut c),
|
|
options,
|
|
vision,
|
|
streams,
|
|
stream,
|
|
)?;
|
|
let AttentionCache::Linear(c) = c.attention else {
|
|
unreachable!()
|
|
};
|
|
states.extend([
|
|
c.conv.ok_or("compiled GDN conv missing")?,
|
|
c.delta.ok_or("compiled GDN state missing")?,
|
|
]);
|
|
if options.gdn.capture {
|
|
captured.extend(c.capture.ok_or("compiled GDN capture missing")?);
|
|
}
|
|
}
|
|
Ok(std::iter::once(h).chain(states).chain(captured).collect())
|
|
})?;
|
|
entry.insert(replay)
|
|
}
|
|
};
|
|
let out = replay.run(&flat)?;
|
|
h = out[0].clone();
|
|
for (n, i) in range.clone().enumerate() {
|
|
let AttentionCache::Linear(c) = &mut cache[i].as_mut().unwrap().attention else {
|
|
unreachable!()
|
|
};
|
|
c.conv = Some(out[1 + 2 * n].clone());
|
|
c.delta = Some(out[2 + 2 * n].clone());
|
|
if options.gdn.capture {
|
|
let start = 1 + 2 * range.len() + 6 * n;
|
|
c.capture = Some(std::array::from_fn(|j| out[start + j].clone()));
|
|
}
|
|
}
|
|
}
|
|
Ok(h)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
impl TextModel<'_> {
|
|
pub(super) fn run_contract(&self, rows: i32, capture: bool) -> serde_json::Value {
|
|
use super::ops::Operation as O;
|
|
let replay = &self
|
|
.replays
|
|
.iter()
|
|
.find(|(key, _)| key.capture == capture && key.signature[0].0 == [1, rows, 10240])
|
|
.expect("compiled fixture run missing")
|
|
.1;
|
|
let mut counts = std::collections::BTreeMap::<String, usize>::new();
|
|
let contract = replay.contract(
|
|
&["hidden", "conv", "delta"],
|
|
|op| {
|
|
let name = match op {
|
|
O::Compiled(expression) => expression.name(),
|
|
_ => match op {
|
|
O::Reshape => "Reshape",
|
|
O::Flatten { .. } => "Flatten",
|
|
O::Unflatten { .. } => "Unflatten",
|
|
O::Transpose(_) => "Transpose",
|
|
O::ExpandDims(_) => "ExpandDims",
|
|
O::Squeeze(_) => "Squeeze",
|
|
O::Slice { .. } => "Slice",
|
|
O::Contiguous(_) => "Contiguous",
|
|
O::RmsNorm(_) => "RMSNorm",
|
|
O::Softmax(_) => "Softmax",
|
|
O::QuantizedLinearBf16 { .. } => "QuantizedMatmul",
|
|
O::GatherQmmBf16 { .. } => "GatherQMM",
|
|
O::DenseMatmul => "Matmul",
|
|
O::Sum(_) => "Sum",
|
|
O::Sort {
|
|
indices: true,
|
|
kth: Some(_),
|
|
..
|
|
} => "ArgPartition",
|
|
O::Split { .. } => "Split",
|
|
O::HyperRead(_)
|
|
| O::GdnStepBf16
|
|
| O::GdnConvNorm(_)
|
|
| O::GatedDeltaBf16(_) => "CustomKernel",
|
|
O::Concatenate(_) => "Concatenate",
|
|
O::DepthwiseConv1d => "Convolution",
|
|
O::GatherAxis(_) => "GatherAxis",
|
|
O::ArangeI32(_) | O::ArangeU32 => "Arange",
|
|
_ => super::fused::operator(op).expect("uncaptured GDN primitive"),
|
|
}
|
|
.to_owned(),
|
|
};
|
|
*counts.entry(name.clone()).or_default() += 1;
|
|
name
|
|
},
|
|
if capture { 9 } else { 3 },
|
|
);
|
|
serde_json::json!({"contract":contract,"ops":counts})
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn mtplx_model_compile_policy_and_run_boundaries() {
|
|
assert_eq!(
|
|
runs([
|
|
(true, false),
|
|
(true, false),
|
|
(true, true),
|
|
(true, false),
|
|
(false, false),
|
|
(true, false)
|
|
]),
|
|
vec![
|
|
Run::Gdn(0..2),
|
|
Run::Eager(2),
|
|
Run::Gdn(3..4),
|
|
Run::Eager(4),
|
|
Run::Gdn(5..6)
|
|
]
|
|
);
|
|
for on in ["1", " true ", "YES", "On"] {
|
|
assert!(CompilePolicy::from_values(Some(on), None).enabled());
|
|
assert!(CompilePolicy::from_values(None, Some(on)).enabled());
|
|
for off in ["0", "", "false", "off", "anything"] {
|
|
for policy in [
|
|
CompilePolicy::from_values(Some(on), Some(off)),
|
|
CompilePolicy::from_values(Some(off), Some(on)),
|
|
] {
|
|
assert!(!policy.enabled());
|
|
}
|
|
}
|
|
}
|
|
let automatic = CompilePolicy::from_values(None, None);
|
|
assert!(!automatic.enabled());
|
|
let environment = CompilePolicy::from_env();
|
|
assert_eq!(
|
|
environment.enabled(),
|
|
CompilePolicy::from_values(
|
|
std::env::var("MTPLX_COMPILED_GDN").ok().as_deref(),
|
|
std::env::var("MTPLX_QWEN4EXP_COMPILE").ok().as_deref(),
|
|
)
|
|
.enabled()
|
|
);
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)]
|
|
fn eager(
|
|
layer: &mut DecoderLayer<'_>,
|
|
hidden: &Array,
|
|
ids: &Array,
|
|
mask: Option<&Array>,
|
|
mut cache: Option<&mut Cache>,
|
|
options: &Options,
|
|
vision: Option<(Option<&Array>, i32)>,
|
|
streams: &Streams,
|
|
stream: Stream,
|
|
) -> Result<Array, String> {
|
|
if options.gdn.capture
|
|
&& layer.ple.is_some()
|
|
&& let Some(c) = cache.as_deref_mut()
|
|
{
|
|
c.capture_ple = Some((hidden.clone(), ids.clone()));
|
|
}
|
|
layer.call(
|
|
hidden,
|
|
Some(ids),
|
|
mask,
|
|
cache,
|
|
options,
|
|
vision,
|
|
streams,
|
|
stream,
|
|
)
|
|
}
|