1342 lines
44 KiB
Rust
1342 lines
44 KiB
Rust
use super::checkpoint::{read_buffer, read_u32, read_u64, write_buffer, write_u32, write_u64};
|
|
use super::*;
|
|
use crate::settings::EngineSsdSettings;
|
|
|
|
const CHECKPOINT_MAGIC: &[u8; 8] = b"DS4GLM01";
|
|
const CHECKPOINT_VERSION: u32 = 1;
|
|
const CACHE_F16: bool = true;
|
|
const DECODE_FLUSH_LAYERS: usize = 4;
|
|
|
|
#[derive(Clone, Copy)]
|
|
struct SparseWeights {
|
|
router: Weight,
|
|
bias: Weight,
|
|
gate: Weight,
|
|
up: Weight,
|
|
down: Weight,
|
|
shared_gate: Weight,
|
|
shared_up: Weight,
|
|
shared_down: Weight,
|
|
}
|
|
|
|
#[derive(Clone, Copy)]
|
|
struct DenseWeights {
|
|
gate: Weight,
|
|
up: Weight,
|
|
down: Weight,
|
|
}
|
|
|
|
struct GlmLayer {
|
|
attn_norm: Weight,
|
|
q_a: Weight,
|
|
q_a_norm: Weight,
|
|
q_b: Weight,
|
|
kv_a: Weight,
|
|
kv_norm: Weight,
|
|
k_b: Weight,
|
|
v_b: Weight,
|
|
output: Weight,
|
|
indexer_k: Weight,
|
|
indexer_q: Weight,
|
|
indexer_k_norm: Weight,
|
|
indexer_k_bias: Weight,
|
|
indexer_proj: Weight,
|
|
ffn_norm: Weight,
|
|
dense: Option<DenseWeights>,
|
|
sparse: Option<SparseWeights>,
|
|
}
|
|
|
|
struct GlmWeights {
|
|
embedding: Weight,
|
|
output_norm: Weight,
|
|
output: Weight,
|
|
layers: Vec<GlmLayer>,
|
|
}
|
|
|
|
impl GlmWeights {
|
|
fn bind(model: &Model) -> Result<Self, String> {
|
|
let main = &model.main;
|
|
let normal_layers = model.shape.layers - model.shape.nextn;
|
|
let layers = (0..normal_layers)
|
|
.map(|index| {
|
|
let required = |suffix: &str| Weight::bind(main, &format!("blk.{index}.{suffix}"));
|
|
let dense = (index < model.shape.leading_dense)
|
|
.then(|| {
|
|
Ok::<_, String>(DenseWeights {
|
|
gate: required("ffn_gate.weight")?,
|
|
up: required("ffn_up.weight")?,
|
|
down: required("ffn_down.weight")?,
|
|
})
|
|
})
|
|
.transpose()?;
|
|
let sparse = (index >= model.shape.leading_dense)
|
|
.then(|| {
|
|
Ok::<_, String>(SparseWeights {
|
|
router: required("ffn_gate_inp.weight")?,
|
|
bias: required("exp_probs_b.bias")?,
|
|
gate: required("ffn_gate_exps.weight")?,
|
|
up: required("ffn_up_exps.weight")?,
|
|
down: required("ffn_down_exps.weight")?,
|
|
shared_gate: required("ffn_gate_shexp.weight")?,
|
|
shared_up: required("ffn_up_shexp.weight")?,
|
|
shared_down: required("ffn_down_shexp.weight")?,
|
|
})
|
|
})
|
|
.transpose()?;
|
|
Ok(GlmLayer {
|
|
attn_norm: required("attn_norm.weight")?,
|
|
q_a: required("attn_q_a.weight")?,
|
|
q_a_norm: required("attn_q_a_norm.weight")?,
|
|
q_b: required("attn_q_b.weight")?,
|
|
kv_a: required("attn_kv_a_mqa.weight")?,
|
|
kv_norm: required("attn_kv_a_norm.weight")?,
|
|
k_b: required("attn_k_b.weight")?,
|
|
v_b: required("attn_v_b.weight")?,
|
|
output: required("attn_output.weight")?,
|
|
indexer_k: required("indexer.attn_k.weight")?,
|
|
indexer_q: required("indexer.attn_q_b.weight")?,
|
|
indexer_k_norm: required("indexer.k_norm.weight")?,
|
|
indexer_k_bias: required("indexer.k_norm.bias")?,
|
|
indexer_proj: required("indexer.proj.weight")?,
|
|
ffn_norm: required("ffn_norm.weight")?,
|
|
dense,
|
|
sparse,
|
|
})
|
|
})
|
|
.collect::<Result<_, String>>()?;
|
|
Ok(Self {
|
|
embedding: Weight::bind(main, "token_embd.weight")?,
|
|
output_norm: Weight::bind(main, "output_norm.weight")?,
|
|
output: Weight::bind(main, "output.weight")?,
|
|
layers,
|
|
})
|
|
}
|
|
}
|
|
|
|
struct LayerCache {
|
|
kv: Buffer,
|
|
rope: Buffer,
|
|
indexer: Option<Buffer>,
|
|
}
|
|
|
|
impl LayerCache {
|
|
fn allocate(shape: super::super::Shape, layer: usize, context: u32) -> Result<Self, String> {
|
|
Ok(Self {
|
|
kv: Buffer::bytes(u64::from(context) * shape.kv_lora * 2)?,
|
|
rope: Buffer::bytes(u64::from(context) * shape.rot * 2)?,
|
|
indexer: full_indexer_layer(shape, layer)
|
|
.then(|| Buffer::bytes(u64::from(context) * shape.indexer_head_dim * 2))
|
|
.transpose()?,
|
|
})
|
|
}
|
|
}
|
|
|
|
struct GlmScratch {
|
|
current: Buffer,
|
|
next: Buffer,
|
|
attn_norm: Buffer,
|
|
q_rank: Buffer,
|
|
q_rank_norm: Buffer,
|
|
q: Buffer,
|
|
kv_raw: Buffer,
|
|
indexer_k: Buffer,
|
|
indexer_q: Buffer,
|
|
indexer_weights: Buffer,
|
|
indexer_scores: Buffer,
|
|
indexer_selected: Buffer,
|
|
qk_low: Buffer,
|
|
heads: Buffer,
|
|
attn_out: Buffer,
|
|
after_attn: Buffer,
|
|
ffn_norm: Buffer,
|
|
ffn_gate: Buffer,
|
|
ffn_up: Buffer,
|
|
ffn_mid: Buffer,
|
|
ffn_out: Buffer,
|
|
ffn_sum: Buffer,
|
|
router_logits: Buffer,
|
|
router_probs: Buffer,
|
|
router_selected: Buffer,
|
|
router_weights: Buffer,
|
|
output_norm: Buffer,
|
|
logits: Buffer,
|
|
}
|
|
|
|
impl GlmScratch {
|
|
fn allocate(model: &Model, context: u32) -> Result<Self, String> {
|
|
let shape = model.shape;
|
|
let q = shape.heads * shape.key_mla;
|
|
let heads = shape.heads * shape.value_mla;
|
|
Ok(Self {
|
|
current: Buffer::floats(shape.embd)?,
|
|
next: Buffer::floats(shape.embd)?,
|
|
attn_norm: Buffer::floats(shape.embd)?,
|
|
q_rank: Buffer::floats(shape.lora_q)?,
|
|
q_rank_norm: Buffer::floats(shape.lora_q)?,
|
|
q: Buffer::floats(q)?,
|
|
kv_raw: Buffer::floats(shape.head_dim)?,
|
|
indexer_k: Buffer::floats(shape.indexer_head_dim)?,
|
|
indexer_q: Buffer::floats(shape.indexer_heads * shape.indexer_head_dim)?,
|
|
indexer_weights: Buffer::floats(shape.indexer_heads)?,
|
|
indexer_scores: Buffer::floats(u64::from(context))?,
|
|
indexer_selected: Buffer::bytes(shape.indexer_top_k * 4)?,
|
|
qk_low: Buffer::floats(shape.heads * shape.kv_lora)?,
|
|
heads: Buffer::floats(heads)?,
|
|
attn_out: Buffer::floats(shape.embd)?,
|
|
after_attn: Buffer::floats(shape.embd)?,
|
|
ffn_norm: Buffer::floats(shape.embd)?,
|
|
ffn_gate: Buffer::floats(shape.ff_dense.max(shape.experts_used * shape.ff_expert))?,
|
|
ffn_up: Buffer::floats(shape.ff_dense.max(shape.experts_used * shape.ff_expert))?,
|
|
ffn_mid: Buffer::floats(shape.ff_dense.max(shape.experts_used * shape.ff_expert))?,
|
|
ffn_out: Buffer::floats(shape.embd)?,
|
|
ffn_sum: Buffer::floats(shape.embd)?,
|
|
router_logits: Buffer::floats(shape.experts)?,
|
|
router_probs: Buffer::floats(shape.experts)?,
|
|
router_selected: Buffer::bytes(shape.experts_used * 4)?,
|
|
router_weights: Buffer::floats(shape.experts_used)?,
|
|
output_norm: Buffer::floats(shape.embd)?,
|
|
logits: Buffer::floats(shape.vocab)?,
|
|
})
|
|
}
|
|
}
|
|
|
|
pub(in crate::engine) struct GlmExecutor {
|
|
weights: GlmWeights,
|
|
scratch: GlmScratch,
|
|
caches: Vec<LayerCache>,
|
|
logits: Vec<f32>,
|
|
tokens: Vec<i32>,
|
|
context: u32,
|
|
quality: bool,
|
|
ssd: EngineSsdSettings,
|
|
checkpoint_tag: [u8; 32],
|
|
model_modified: (u64, u32),
|
|
model_identity: [u8; 32],
|
|
_context: Context,
|
|
model: Model,
|
|
}
|
|
|
|
impl GlmExecutor {
|
|
pub(super) fn open(
|
|
model: Model,
|
|
context: u32,
|
|
quality: bool,
|
|
ssd: EngineSsdSettings,
|
|
) -> Result<Self, String> {
|
|
if context == 0 || u64::from(context) > model.shape.original_context {
|
|
return Err(format!(
|
|
"GLM context must be between 1 and {} tokens",
|
|
model.shape.original_context
|
|
));
|
|
}
|
|
let weights = GlmWeights::bind(&model)?;
|
|
let admission = admission_bytes(&model, &weights, context, ssd)?;
|
|
let context_handle = Context::open(&model, quality, ssd.enabled, admission)?;
|
|
configure_streaming(&model, &weights, ssd)?;
|
|
let scratch = GlmScratch::allocate(&model, context)?;
|
|
let caches = (0..weights.layers.len())
|
|
.map(|layer| LayerCache::allocate(model.shape, layer, context))
|
|
.collect::<Result<_, _>>()?;
|
|
let model_modified = fs::metadata(model.main.path())
|
|
.and_then(|metadata| metadata.modified())
|
|
.ok()
|
|
.and_then(|modified| modified.duration_since(UNIX_EPOCH).ok())
|
|
.map(|duration| (duration.as_secs(), duration.subsec_nanos()))
|
|
.unwrap_or_default();
|
|
let model_identity = model.checkpoint_identity();
|
|
Ok(Self {
|
|
weights,
|
|
scratch,
|
|
caches,
|
|
logits: vec![0.0; model.shape.vocab as usize],
|
|
tokens: Vec::new(),
|
|
context,
|
|
quality,
|
|
ssd,
|
|
checkpoint_tag: [0; 32],
|
|
model_modified,
|
|
model_identity,
|
|
_context: context_handle,
|
|
model,
|
|
})
|
|
}
|
|
|
|
pub(super) fn eval(&mut self, token: i32) -> Result<(), String> {
|
|
let shape = self.model.shape;
|
|
if token < 0 || token as u64 >= shape.vocab {
|
|
return Err(format!("token {token} is outside the vocabulary"));
|
|
}
|
|
let pos = u32::try_from(self.tokens.len()).map_err(|_| "GLM position overflow")?;
|
|
if pos >= self.context {
|
|
return Err(format!(
|
|
"the GLM Metal executor supports {} tokens",
|
|
self.context
|
|
));
|
|
}
|
|
let map = self.model.main.map_ptr().cast();
|
|
let size = self.model.main.len();
|
|
let commands = Commands::begin()?;
|
|
call(
|
|
unsafe {
|
|
ds4_gpu_embed_token_quant_tensor(
|
|
self.scratch.current.raw(),
|
|
map,
|
|
size,
|
|
self.weights.embedding.offset,
|
|
self.weights.embedding.kind,
|
|
shape.vocab as u32,
|
|
token as u32,
|
|
shape.embd as u32,
|
|
)
|
|
},
|
|
"GLM token embedding",
|
|
)?;
|
|
let mut selected_count = 0;
|
|
for (layer_index, ((layer, cache), ordinal)) in self
|
|
.weights
|
|
.layers
|
|
.iter()
|
|
.zip(&self.caches)
|
|
.zip(0_u32..)
|
|
.enumerate()
|
|
{
|
|
self.encode_layer(layer, cache, layer_index, ordinal, pos, &mut selected_count)?;
|
|
std::mem::swap(&mut self.scratch.current, &mut self.scratch.next);
|
|
if (layer_index + 1).is_multiple_of(DECODE_FLUSH_LAYERS)
|
|
&& layer_index + 1 < self.weights.layers.len()
|
|
{
|
|
call(
|
|
unsafe { ds4_gpu_flush_commands() },
|
|
"flushing the GLM decode graph",
|
|
)?;
|
|
}
|
|
}
|
|
norm(
|
|
&self.scratch.output_norm,
|
|
&self.scratch.current,
|
|
self.weights.output_norm,
|
|
shape.embd as u32,
|
|
shape.rms_epsilon,
|
|
map,
|
|
size,
|
|
)?;
|
|
project(
|
|
&self.scratch.logits,
|
|
self.weights.output,
|
|
shape.embd,
|
|
shape.vocab,
|
|
&self.scratch.output_norm,
|
|
map,
|
|
size,
|
|
self.ssd.enabled,
|
|
)?;
|
|
commands.finish()?;
|
|
self.scratch.logits.read_f32(&mut self.logits)?;
|
|
self.tokens.push(token);
|
|
Ok(())
|
|
}
|
|
|
|
fn encode_layer(
|
|
&self,
|
|
layer: &GlmLayer,
|
|
cache: &LayerCache,
|
|
_layer_index: usize,
|
|
ordinal: u32,
|
|
pos: u32,
|
|
selected_count: &mut u32,
|
|
) -> Result<(), String> {
|
|
let shape = self.model.shape;
|
|
let map = self.model.main.map_ptr().cast();
|
|
let size = self.model.main.len();
|
|
let q_dim = shape.heads * shape.key_mla;
|
|
let q_nope = shape.key_mla - shape.rot;
|
|
norm(
|
|
&self.scratch.attn_norm,
|
|
&self.scratch.current,
|
|
layer.attn_norm,
|
|
shape.embd as u32,
|
|
shape.rms_epsilon,
|
|
map,
|
|
size,
|
|
)?;
|
|
project(
|
|
&self.scratch.q_rank,
|
|
layer.q_a,
|
|
shape.embd,
|
|
shape.lora_q,
|
|
&self.scratch.attn_norm,
|
|
map,
|
|
size,
|
|
self.ssd.enabled,
|
|
)?;
|
|
project(
|
|
&self.scratch.kv_raw,
|
|
layer.kv_a,
|
|
shape.embd,
|
|
shape.head_dim,
|
|
&self.scratch.attn_norm,
|
|
map,
|
|
size,
|
|
self.ssd.enabled,
|
|
)?;
|
|
call(
|
|
unsafe {
|
|
ds4_gpu_glm_qkv_norm_store_compact_kv_tensor(
|
|
self.scratch.q_rank_norm.raw(),
|
|
self.scratch.q_rank.raw(),
|
|
map,
|
|
size,
|
|
layer.q_a_norm.offset,
|
|
shape.lora_q as u32,
|
|
cache.kv.raw(),
|
|
cache.rope.raw(),
|
|
self.scratch.kv_raw.raw(),
|
|
layer.kv_norm.offset,
|
|
pos,
|
|
1,
|
|
self.context,
|
|
shape.head_dim as u32,
|
|
shape.kv_lora as u32,
|
|
shape.rot as u32,
|
|
CACHE_F16,
|
|
shape.rms_epsilon,
|
|
)
|
|
},
|
|
"storing GLM compact KV",
|
|
)?;
|
|
project(
|
|
&self.scratch.q,
|
|
layer.q_b,
|
|
shape.lora_q,
|
|
q_dim,
|
|
&self.scratch.q_rank_norm,
|
|
map,
|
|
size,
|
|
self.ssd.enabled,
|
|
)?;
|
|
call(
|
|
unsafe {
|
|
ds4_gpu_glm_rope_tail_tensor(
|
|
self.scratch.q.raw(),
|
|
1,
|
|
shape.heads as u32,
|
|
shape.key_mla as u32,
|
|
shape.rot as u32,
|
|
pos,
|
|
0,
|
|
shape.rope_base,
|
|
1.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
)
|
|
},
|
|
"applying GLM query RoPE",
|
|
)?;
|
|
|
|
if let Some(indexer_cache) = &cache.indexer {
|
|
project(
|
|
&self.scratch.indexer_k,
|
|
layer.indexer_k,
|
|
shape.embd,
|
|
shape.indexer_head_dim,
|
|
&self.scratch.current,
|
|
map,
|
|
size,
|
|
self.ssd.enabled,
|
|
)?;
|
|
call(
|
|
unsafe {
|
|
ds4_gpu_glm_store_indexer_k_tensor(
|
|
indexer_cache.raw(),
|
|
self.scratch.indexer_k.raw(),
|
|
map,
|
|
size,
|
|
layer.indexer_k_norm.offset,
|
|
layer.indexer_k_bias.offset,
|
|
pos,
|
|
1,
|
|
self.context,
|
|
shape.indexer_head_dim as u32,
|
|
shape.rot as u32,
|
|
0,
|
|
1.0e-6,
|
|
shape.rope_base,
|
|
1.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
CACHE_F16,
|
|
)
|
|
},
|
|
"storing the GLM indexer key",
|
|
)?;
|
|
let visible = pos + 1;
|
|
*selected_count = visible.min(shape.indexer_top_k as u32);
|
|
if visible <= shape.indexer_top_k as u32 {
|
|
call(
|
|
unsafe {
|
|
ds4_gpu_glm_fill_selected_range_tensor(
|
|
self.scratch.indexer_selected.raw(),
|
|
*selected_count,
|
|
)
|
|
},
|
|
"selecting the visible GLM context",
|
|
)?;
|
|
} else {
|
|
project(
|
|
&self.scratch.indexer_q,
|
|
layer.indexer_q,
|
|
shape.lora_q,
|
|
shape.indexer_heads * shape.indexer_head_dim,
|
|
&self.scratch.q_rank_norm,
|
|
map,
|
|
size,
|
|
self.ssd.enabled,
|
|
)?;
|
|
call(
|
|
unsafe {
|
|
ds4_gpu_glm_indexer_rope_tail_tensor(
|
|
self.scratch.indexer_q.raw(),
|
|
1,
|
|
shape.indexer_heads as u32,
|
|
shape.indexer_head_dim as u32,
|
|
shape.rot as u32,
|
|
pos,
|
|
0,
|
|
shape.rope_base,
|
|
1.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
)
|
|
},
|
|
"applying GLM indexer RoPE",
|
|
)?;
|
|
f32_project(
|
|
&self.scratch.indexer_weights,
|
|
layer.indexer_proj,
|
|
shape.embd,
|
|
shape.indexer_heads,
|
|
&self.scratch.current,
|
|
map,
|
|
size,
|
|
)?;
|
|
let scale = 1.0 / ((shape.indexer_heads * shape.indexer_head_dim) as f32).sqrt();
|
|
call(
|
|
unsafe {
|
|
ds4_gpu_glm_indexer_score_one_tensor(
|
|
self.scratch.indexer_scores.raw(),
|
|
self.scratch.indexer_q.raw(),
|
|
self.scratch.indexer_weights.raw(),
|
|
indexer_cache.raw(),
|
|
visible,
|
|
shape.indexer_heads as u32,
|
|
shape.indexer_head_dim as u32,
|
|
scale,
|
|
CACHE_F16,
|
|
)
|
|
},
|
|
"scoring the GLM indexer",
|
|
)?;
|
|
call(
|
|
unsafe {
|
|
ds4_gpu_indexer_topk_tensor(
|
|
self.scratch.indexer_selected.raw(),
|
|
self.scratch.indexer_scores.raw(),
|
|
visible,
|
|
1,
|
|
*selected_count,
|
|
)
|
|
},
|
|
"selecting GLM indexed attention rows",
|
|
)?;
|
|
}
|
|
}
|
|
if *selected_count == 0 {
|
|
return Err("GLM indexer did not select an attention context".into());
|
|
}
|
|
call(
|
|
unsafe {
|
|
ds4_gpu_glm_qk_lowrank_typed_tensor(
|
|
self.scratch.qk_low.raw(),
|
|
self.scratch.q.raw(),
|
|
map,
|
|
size,
|
|
layer.k_b.offset,
|
|
layer.k_b.kind,
|
|
shape.heads as u32,
|
|
shape.kv_lora as u32,
|
|
q_nope as u32,
|
|
shape.key_mla as u32,
|
|
)
|
|
},
|
|
"projecting the GLM low-rank query",
|
|
)?;
|
|
call(
|
|
unsafe {
|
|
ds4_gpu_glm_attention_indexed_decode_typed_tensor(
|
|
self.scratch.heads.raw(),
|
|
self.scratch.q.raw(),
|
|
self.scratch.qk_low.raw(),
|
|
cache.kv.raw(),
|
|
cache.rope.raw(),
|
|
map,
|
|
size,
|
|
layer.v_b.offset,
|
|
layer.v_b.kind,
|
|
self.scratch.indexer_selected.raw(),
|
|
*selected_count,
|
|
self.context,
|
|
CACHE_F16,
|
|
shape.heads as u32,
|
|
shape.kv_lora as u32,
|
|
q_nope as u32,
|
|
shape.rot as u32,
|
|
shape.value_mla as u32,
|
|
0,
|
|
shape.rope_base,
|
|
1.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
)
|
|
},
|
|
"running GLM indexed attention",
|
|
)?;
|
|
project(
|
|
&self.scratch.attn_out,
|
|
layer.output,
|
|
shape.heads * shape.value_mla,
|
|
shape.embd,
|
|
&self.scratch.heads,
|
|
map,
|
|
size,
|
|
self.ssd.enabled,
|
|
)?;
|
|
call(
|
|
unsafe {
|
|
ds4_gpu_add_rms_norm_weight_tensor(
|
|
self.scratch.ffn_norm.raw(),
|
|
self.scratch.after_attn.raw(),
|
|
self.scratch.current.raw(),
|
|
self.scratch.attn_out.raw(),
|
|
map,
|
|
size,
|
|
layer.ffn_norm.offset,
|
|
shape.embd as u32,
|
|
shape.rms_epsilon,
|
|
)
|
|
},
|
|
"normalizing the GLM FFN input",
|
|
)?;
|
|
|
|
if let Some(dense) = layer.dense {
|
|
project(
|
|
&self.scratch.ffn_gate,
|
|
dense.gate,
|
|
shape.embd,
|
|
shape.ff_dense,
|
|
&self.scratch.ffn_norm,
|
|
map,
|
|
size,
|
|
self.ssd.enabled,
|
|
)?;
|
|
project(
|
|
&self.scratch.ffn_up,
|
|
dense.up,
|
|
shape.embd,
|
|
shape.ff_dense,
|
|
&self.scratch.ffn_norm,
|
|
map,
|
|
size,
|
|
self.ssd.enabled,
|
|
)?;
|
|
call(
|
|
unsafe {
|
|
ds4_gpu_swiglu_tensor(
|
|
self.scratch.ffn_mid.raw(),
|
|
self.scratch.ffn_gate.raw(),
|
|
self.scratch.ffn_up.raw(),
|
|
shape.ff_dense as u32,
|
|
0.0,
|
|
1.0,
|
|
)
|
|
},
|
|
"activating the dense GLM FFN",
|
|
)?;
|
|
project(
|
|
&self.scratch.ffn_out,
|
|
dense.down,
|
|
shape.ff_dense,
|
|
shape.embd,
|
|
&self.scratch.ffn_mid,
|
|
map,
|
|
size,
|
|
self.ssd.enabled,
|
|
)?;
|
|
call(
|
|
unsafe {
|
|
ds4_gpu_add_tensor(
|
|
self.scratch.next.raw(),
|
|
self.scratch.after_attn.raw(),
|
|
self.scratch.ffn_out.raw(),
|
|
shape.embd as u32,
|
|
)
|
|
},
|
|
"adding the dense GLM residual",
|
|
)?;
|
|
} else if let Some(sparse) = layer.sparse {
|
|
f32_project(
|
|
&self.scratch.router_logits,
|
|
sparse.router,
|
|
shape.embd,
|
|
shape.experts,
|
|
&self.scratch.ffn_norm,
|
|
map,
|
|
size,
|
|
)?;
|
|
call(
|
|
unsafe {
|
|
ds4_gpu_glm_router_select_tensor(
|
|
self.scratch.router_selected.raw(),
|
|
self.scratch.router_weights.raw(),
|
|
self.scratch.router_probs.raw(),
|
|
map,
|
|
size,
|
|
sparse.bias.offset,
|
|
self.scratch.router_logits.raw(),
|
|
shape.experts as u32,
|
|
shape.experts_used as u32,
|
|
shape.expert_weight_scale,
|
|
)
|
|
},
|
|
"routing the GLM experts",
|
|
)?;
|
|
let force_resident = self.ssd.enabled
|
|
&& ordinal.saturating_sub(shape.leading_dense) < self.ssd.full_layers;
|
|
if self.ssd.enabled && !force_resident {
|
|
let table = expert_table(map, size, ordinal, shape, sparse);
|
|
call(
|
|
unsafe {
|
|
ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor(
|
|
&table,
|
|
self.scratch.router_selected.raw(),
|
|
shape.experts_used as u32,
|
|
)
|
|
},
|
|
"loading selected GLM experts",
|
|
)?;
|
|
}
|
|
let (gate_expert, gate_row) = expert_layout(sparse.gate, shape.experts);
|
|
let (up_expert, up_row) = expert_layout(sparse.up, shape.experts);
|
|
let (down_expert, down_row) = expert_layout(sparse.down, shape.experts);
|
|
call(
|
|
unsafe {
|
|
ds4_gpu_glm_routed_moe_one_tensor(
|
|
self.scratch.ffn_out.raw(),
|
|
self.scratch.ffn_mid.raw(),
|
|
map,
|
|
size,
|
|
sparse.gate.offset,
|
|
sparse.up.offset,
|
|
sparse.down.offset,
|
|
sparse.gate.kind,
|
|
sparse.up.kind,
|
|
sparse.down.kind,
|
|
gate_expert,
|
|
gate_row,
|
|
up_expert,
|
|
up_row,
|
|
down_expert,
|
|
down_row,
|
|
shape.embd as u32,
|
|
shape.ff_expert as u32,
|
|
shape.embd as u32,
|
|
self.scratch.router_selected.raw(),
|
|
self.scratch.router_weights.raw(),
|
|
shape.experts as u32,
|
|
shape.experts_used as u32,
|
|
ordinal,
|
|
self.scratch.ffn_norm.raw(),
|
|
force_resident,
|
|
)
|
|
},
|
|
"running the routed GLM experts",
|
|
)?;
|
|
project(
|
|
&self.scratch.ffn_gate,
|
|
sparse.shared_gate,
|
|
shape.embd,
|
|
shape.ff_expert,
|
|
&self.scratch.ffn_norm,
|
|
map,
|
|
size,
|
|
self.ssd.enabled,
|
|
)?;
|
|
project(
|
|
&self.scratch.ffn_up,
|
|
sparse.shared_up,
|
|
shape.embd,
|
|
shape.ff_expert,
|
|
&self.scratch.ffn_norm,
|
|
map,
|
|
size,
|
|
self.ssd.enabled,
|
|
)?;
|
|
call(
|
|
unsafe {
|
|
ds4_gpu_swiglu_tensor(
|
|
self.scratch.ffn_mid.raw(),
|
|
self.scratch.ffn_gate.raw(),
|
|
self.scratch.ffn_up.raw(),
|
|
shape.ff_expert as u32,
|
|
0.0,
|
|
1.0,
|
|
)
|
|
},
|
|
"activating the shared GLM expert",
|
|
)?;
|
|
project(
|
|
&self.scratch.ffn_sum,
|
|
sparse.shared_down,
|
|
shape.ff_expert,
|
|
shape.embd,
|
|
&self.scratch.ffn_mid,
|
|
map,
|
|
size,
|
|
self.ssd.enabled,
|
|
)?;
|
|
call(
|
|
unsafe {
|
|
ds4_gpu_add3_tensor(
|
|
self.scratch.next.raw(),
|
|
self.scratch.after_attn.raw(),
|
|
self.scratch.ffn_out.raw(),
|
|
self.scratch.ffn_sum.raw(),
|
|
shape.embd as u32,
|
|
)
|
|
},
|
|
"adding the sparse GLM residual",
|
|
)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
pub(super) fn prefill(
|
|
&mut self,
|
|
tokens: &[i32],
|
|
mut progress: impl FnMut(u32) -> bool,
|
|
) -> Result<usize, String> {
|
|
let mut completed = 0;
|
|
for &token in tokens {
|
|
if !progress(self.position()) {
|
|
break;
|
|
}
|
|
self.eval(token)?;
|
|
completed += 1;
|
|
}
|
|
Ok(completed)
|
|
}
|
|
|
|
pub(super) fn logits(&self) -> &[f32] {
|
|
&self.logits
|
|
}
|
|
pub(super) fn model(&self) -> &Model {
|
|
&self.model
|
|
}
|
|
pub(super) fn context(&self) -> u32 {
|
|
self.context
|
|
}
|
|
pub(super) fn position(&self) -> u32 {
|
|
self.tokens.len() as u32
|
|
}
|
|
pub(super) fn tokens(&self) -> &[i32] {
|
|
&self.tokens
|
|
}
|
|
pub(super) fn checkpoint_tag(&self) -> [u8; 32] {
|
|
self.checkpoint_tag
|
|
}
|
|
pub(super) fn note_checkpoint_tag(&mut self, tag: [u8; 32]) {
|
|
self.checkpoint_tag = tag;
|
|
}
|
|
|
|
pub(super) fn reset(&mut self) -> Result<(), String> {
|
|
self.scratch = GlmScratch::allocate(&self.model, self.context)?;
|
|
self.caches = (0..self.weights.layers.len())
|
|
.map(|layer| LayerCache::allocate(self.model.shape, layer, self.context))
|
|
.collect::<Result<_, _>>()?;
|
|
self.tokens.clear();
|
|
self.checkpoint_tag = [0; 32];
|
|
Ok(())
|
|
}
|
|
|
|
pub(super) fn align_prompt(&mut self, tokens: &[i32]) -> Result<usize, String> {
|
|
if !tokens.starts_with(&self.tokens) {
|
|
self.reset()?;
|
|
}
|
|
Ok(self.tokens.len())
|
|
}
|
|
|
|
pub(super) fn save_checkpoint(
|
|
&mut self,
|
|
path: &Path,
|
|
tag: [u8; 32],
|
|
progress: &mut impl FnMut(u64),
|
|
) -> Result<(), String> {
|
|
if let Some(parent) = path.parent() {
|
|
fs::create_dir_all(parent).map_err(|e| e.to_string())?;
|
|
}
|
|
let temporary = path.with_extension("tmp");
|
|
let mut file = File::create(&temporary).map_err(|e| e.to_string())?;
|
|
file.write_all(CHECKPOINT_MAGIC)
|
|
.map_err(|e| e.to_string())?;
|
|
let shape = self.model.shape;
|
|
for value in [
|
|
CHECKPOINT_VERSION,
|
|
self.context,
|
|
self.position(),
|
|
self.quality as u32,
|
|
self.weights.layers.len() as u32,
|
|
shape.kv_lora as u32,
|
|
shape.rot as u32,
|
|
shape.indexer_head_dim as u32,
|
|
shape.vocab as u32,
|
|
] {
|
|
write_u32(&mut file, value)?;
|
|
}
|
|
write_u64(&mut file, self.model.main.len())?;
|
|
write_u64(&mut file, self.model_modified.0)?;
|
|
write_u32(&mut file, self.model_modified.1)?;
|
|
file.write_all(&self.model_identity)
|
|
.map_err(|e| e.to_string())?;
|
|
file.write_all(&tag).map_err(|e| e.to_string())?;
|
|
for &token in &self.tokens {
|
|
write_u32(&mut file, token as u32)?;
|
|
}
|
|
for &logit in &self.logits {
|
|
write_u32(&mut file, logit.to_bits())?;
|
|
}
|
|
let mut chunk = vec![0; CHECKPOINT_IO_CHUNK];
|
|
let rows = u64::from(self.position());
|
|
for cache in &self.caches {
|
|
write_buffer(
|
|
&mut file,
|
|
&cache.kv,
|
|
0,
|
|
rows * shape.kv_lora * 2,
|
|
&mut chunk,
|
|
progress,
|
|
)?;
|
|
write_buffer(
|
|
&mut file,
|
|
&cache.rope,
|
|
0,
|
|
rows * shape.rot * 2,
|
|
&mut chunk,
|
|
progress,
|
|
)?;
|
|
if let Some(indexer) = &cache.indexer {
|
|
write_buffer(
|
|
&mut file,
|
|
indexer,
|
|
0,
|
|
rows * shape.indexer_head_dim * 2,
|
|
&mut chunk,
|
|
progress,
|
|
)?;
|
|
}
|
|
}
|
|
file.sync_all().map_err(|e| e.to_string())?;
|
|
fs::rename(&temporary, path).map_err(|e| e.to_string())?;
|
|
self.checkpoint_tag = tag;
|
|
Ok(())
|
|
}
|
|
|
|
pub(super) fn load_checkpoint(
|
|
&mut self,
|
|
path: &Path,
|
|
progress: &mut impl FnMut(u64),
|
|
) -> Result<bool, String> {
|
|
let mut file = match File::open(path) {
|
|
Ok(file) => file,
|
|
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false),
|
|
Err(error) => return Err(error.to_string()),
|
|
};
|
|
let mut magic = [0; 8];
|
|
file.read_exact(&mut magic).map_err(|e| e.to_string())?;
|
|
if &magic != CHECKPOINT_MAGIC
|
|
|| read_u32(&mut file)? != CHECKPOINT_VERSION
|
|
|| read_u32(&mut file)? != self.context
|
|
{
|
|
return Err("GLM checkpoint does not match the current executor".into());
|
|
}
|
|
let token_count = read_u32(&mut file)?;
|
|
let shape = self.model.shape;
|
|
if token_count > self.context
|
|
|| read_u32(&mut file)? != self.quality as u32
|
|
|| read_u32(&mut file)? != self.weights.layers.len() as u32
|
|
|| read_u32(&mut file)? != shape.kv_lora as u32
|
|
|| read_u32(&mut file)? != shape.rot as u32
|
|
|| read_u32(&mut file)? != shape.indexer_head_dim as u32
|
|
|| read_u32(&mut file)? != shape.vocab as u32
|
|
|| read_u64(&mut file)? != self.model.main.len()
|
|
|| read_u64(&mut file)? != self.model_modified.0
|
|
|| read_u32(&mut file)? != self.model_modified.1
|
|
{
|
|
return Err("GLM checkpoint was written for a different model or configuration".into());
|
|
}
|
|
let mut identity = [0; 32];
|
|
file.read_exact(&mut identity).map_err(|e| e.to_string())?;
|
|
if identity != self.model_identity {
|
|
return Err("GLM checkpoint model identity changed".into());
|
|
}
|
|
let mut tag = [0; 32];
|
|
file.read_exact(&mut tag).map_err(|e| e.to_string())?;
|
|
let mut tokens = Vec::with_capacity(token_count as usize);
|
|
for _ in 0..token_count {
|
|
let token = read_u32(&mut file)?;
|
|
if u64::from(token) >= self.model.shape.vocab {
|
|
return Err("GLM checkpoint token is invalid".into());
|
|
}
|
|
tokens.push(token as i32);
|
|
}
|
|
let mut logits = Vec::with_capacity(shape.vocab as usize);
|
|
for _ in 0..shape.vocab {
|
|
logits.push(f32::from_bits(read_u32(&mut file)?));
|
|
}
|
|
self.reset()?;
|
|
let mut chunk = vec![0; CHECKPOINT_IO_CHUNK];
|
|
let rows = u64::from(token_count);
|
|
for cache in &self.caches {
|
|
read_buffer(
|
|
&mut file,
|
|
&cache.kv,
|
|
0,
|
|
rows * shape.kv_lora * 2,
|
|
&mut chunk,
|
|
progress,
|
|
)?;
|
|
read_buffer(
|
|
&mut file,
|
|
&cache.rope,
|
|
0,
|
|
rows * shape.rot * 2,
|
|
&mut chunk,
|
|
progress,
|
|
)?;
|
|
if let Some(indexer) = &cache.indexer {
|
|
read_buffer(
|
|
&mut file,
|
|
indexer,
|
|
0,
|
|
rows * shape.indexer_head_dim * 2,
|
|
&mut chunk,
|
|
progress,
|
|
)?;
|
|
}
|
|
}
|
|
let mut trailing = [0];
|
|
if file.read(&mut trailing).map_err(|e| e.to_string())? != 0 {
|
|
return Err("GLM checkpoint has trailing data".into());
|
|
}
|
|
self.tokens = tokens;
|
|
self.logits = logits;
|
|
self.checkpoint_tag = tag;
|
|
Ok(true)
|
|
}
|
|
}
|
|
|
|
fn full_indexer_layer(shape: super::super::Shape, layer: usize) -> bool {
|
|
layer < (shape.layers - shape.nextn) as usize
|
|
&& (layer < shape.leading_dense as usize || (layer >= 6 && (layer - 6).is_multiple_of(4)))
|
|
}
|
|
|
|
fn expert_layout(weight: Weight, experts: u64) -> (u64, u64) {
|
|
let expert = weight.bytes / experts;
|
|
(expert, expert / weight.dims[1])
|
|
}
|
|
|
|
fn expert_table(
|
|
map: *const c_void,
|
|
size: u64,
|
|
layer: u32,
|
|
shape: super::super::Shape,
|
|
weights: SparseWeights,
|
|
) -> StreamExpertTable {
|
|
StreamExpertTable {
|
|
model_map: map,
|
|
model_size: size,
|
|
layer,
|
|
total_experts: shape.experts as u32,
|
|
gate_offset: weights.gate.offset,
|
|
up_offset: weights.up.offset,
|
|
down_offset: weights.down.offset,
|
|
gate_expert_bytes: weights.gate.bytes / shape.experts,
|
|
down_expert_bytes: weights.down.bytes / shape.experts,
|
|
}
|
|
}
|
|
|
|
fn configure_streaming(
|
|
model: &Model,
|
|
weights: &GlmWeights,
|
|
ssd: EngineSsdSettings,
|
|
) -> Result<(), String> {
|
|
if !ssd.enabled {
|
|
return Ok(());
|
|
}
|
|
let sparse = weights
|
|
.layers
|
|
.iter()
|
|
.find_map(|layer| layer.sparse)
|
|
.ok_or("GLM model has no routed expert layers")?;
|
|
let gate = sparse.gate.bytes / model.shape.experts;
|
|
let up = sparse.up.bytes / model.shape.experts;
|
|
let down = sparse.down.bytes / model.shape.experts;
|
|
let per_expert = gate
|
|
.checked_add(up)
|
|
.and_then(|bytes| bytes.checked_add(down))
|
|
.ok_or("GLM expert size overflow")?;
|
|
let budget = if ssd.cache_experts != 0 {
|
|
ssd.cache_experts
|
|
} else if ssd.cache_bytes != 0 {
|
|
u32::try_from((ssd.cache_bytes / per_expert).max(1)).unwrap_or(u32::MAX)
|
|
} else {
|
|
unsafe { ds4_gpu_stream_expert_cache_budget_for_expert_size(gate, down) }
|
|
};
|
|
if budget == 0 {
|
|
return Err("GLM SSD streaming has no memory for an expert cache".into());
|
|
}
|
|
unsafe {
|
|
ds4_gpu_set_streaming_expert_cache_expert_bytes(per_expert);
|
|
ds4_gpu_set_streaming_expert_cache_budget(budget);
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn admission_bytes(
|
|
model: &Model,
|
|
weights: &GlmWeights,
|
|
context: u32,
|
|
ssd: EngineSsdSettings,
|
|
) -> Result<u64, String> {
|
|
let shape = model.shape;
|
|
let normal_layers = u64::from(shape.layers - shape.nextn);
|
|
let indexer_layers = (0..normal_layers as usize)
|
|
.filter(|layer| full_indexer_layer(shape, *layer))
|
|
.count() as u64;
|
|
let per_layer = u64::from(context)
|
|
.checked_mul((shape.kv_lora + shape.rot) * 2)
|
|
.ok_or("GLM compact-cache size overflow")?;
|
|
let kv = normal_layers
|
|
.checked_mul(per_layer)
|
|
.and_then(|bytes| {
|
|
bytes.checked_add(indexer_layers * u64::from(context) * shape.indexer_head_dim * 2)
|
|
})
|
|
.ok_or("GLM compact-cache size overflow")?;
|
|
let resident = if ssd.enabled {
|
|
model
|
|
.main
|
|
.tensors
|
|
.iter()
|
|
.filter(|(name, _)| {
|
|
!name.ends_with("ffn_gate_exps.weight")
|
|
&& !name.ends_with("ffn_up_exps.weight")
|
|
&& !name.ends_with("ffn_down_exps.weight")
|
|
})
|
|
.try_fold(0_u64, |total, (_, tensor)| total.checked_add(tensor.bytes))
|
|
.ok_or("GLM resident tensor size overflow")?
|
|
} else {
|
|
model.main.len() - model.main.data_offset()
|
|
};
|
|
let sparse_expert = weights
|
|
.layers
|
|
.iter()
|
|
.find_map(|layer| layer.sparse)
|
|
.map(|weights| {
|
|
weights.gate.bytes / shape.experts
|
|
+ weights.up.bytes / shape.experts
|
|
+ weights.down.bytes / shape.experts
|
|
})
|
|
.unwrap_or(0);
|
|
let cache = if ssd.enabled {
|
|
if ssd.cache_bytes != 0 {
|
|
ssd.cache_bytes
|
|
} else if ssd.cache_experts != 0 {
|
|
sparse_expert.saturating_mul(u64::from(ssd.cache_experts))
|
|
} else {
|
|
// Reference auto policy reserves roughly one seventh of the available
|
|
// budget for full layers and uses the remainder for selected experts.
|
|
16 * 1024 * 1024 * 1024_u64
|
|
}
|
|
} else {
|
|
0
|
|
};
|
|
let full_layers = sparse_expert
|
|
.saturating_mul(shape.experts)
|
|
.saturating_mul(u64::from(ssd.full_layers));
|
|
resident
|
|
.checked_add(kv)
|
|
.and_then(|bytes| bytes.checked_add(cache))
|
|
.and_then(|bytes| bytes.checked_add(full_layers))
|
|
.and_then(|bytes| bytes.checked_add(512 * 1024 * 1024))
|
|
.ok_or_else(|| "GLM runtime memory size overflow".into())
|
|
}
|
|
|
|
fn norm(
|
|
out: &Buffer,
|
|
input: &Buffer,
|
|
weight: Weight,
|
|
width: u32,
|
|
epsilon: f32,
|
|
map: *const c_void,
|
|
size: u64,
|
|
) -> Result<(), String> {
|
|
call(
|
|
unsafe {
|
|
ds4_gpu_rms_norm_weight_tensor(
|
|
out.raw(),
|
|
input.raw(),
|
|
map,
|
|
size,
|
|
weight.offset,
|
|
width,
|
|
epsilon,
|
|
)
|
|
},
|
|
"normalizing GLM activations",
|
|
)
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)]
|
|
fn project(
|
|
out: &Buffer,
|
|
weight: Weight,
|
|
input: u64,
|
|
output: u64,
|
|
x: &Buffer,
|
|
map: *const c_void,
|
|
size: u64,
|
|
streaming: bool,
|
|
) -> Result<(), String> {
|
|
let result = unsafe {
|
|
if streaming {
|
|
ds4_gpu_matmul_quant_tensor(
|
|
out.raw(),
|
|
map,
|
|
size,
|
|
weight.offset,
|
|
weight.kind,
|
|
input,
|
|
output,
|
|
x.raw(),
|
|
1,
|
|
)
|
|
} else {
|
|
ds4_gpu_matmul_quant_decode_mpp_model_view_tensor(
|
|
out.raw(),
|
|
map,
|
|
size,
|
|
weight.offset,
|
|
weight.kind,
|
|
input,
|
|
output,
|
|
x.raw(),
|
|
1,
|
|
)
|
|
}
|
|
};
|
|
call(result, "projecting GLM activations")
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)]
|
|
fn f32_project(
|
|
out: &Buffer,
|
|
weight: Weight,
|
|
input: u64,
|
|
output: u64,
|
|
x: &Buffer,
|
|
map: *const c_void,
|
|
size: u64,
|
|
) -> Result<(), String> {
|
|
call(
|
|
unsafe {
|
|
ds4_gpu_matmul_f32_tensor(
|
|
out.raw(),
|
|
map,
|
|
size,
|
|
weight.offset,
|
|
input,
|
|
output,
|
|
x.raw(),
|
|
1,
|
|
)
|
|
},
|
|
"projecting GLM F32 activations",
|
|
)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::{GlmExecutor, full_indexer_layer};
|
|
use crate::engine::{GLM, Model};
|
|
use crate::model::ModelChoice;
|
|
use crate::settings::EngineSsdSettings;
|
|
use std::path::Path;
|
|
|
|
#[test]
|
|
fn dsa_indexer_schedule_matches_the_reference() {
|
|
assert!(full_indexer_layer(GLM, 0));
|
|
assert!(full_indexer_layer(GLM, 2));
|
|
assert!(!full_indexer_layer(GLM, 3));
|
|
assert!(full_indexer_layer(GLM, 6));
|
|
assert!(full_indexer_layer(GLM, 10));
|
|
assert!(!full_indexer_layer(GLM, 11));
|
|
assert!(!full_indexer_layer(GLM, 78));
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "requires the 197 GiB GLM 5.2 checkpoint and Apple Metal"]
|
|
fn resident_and_streamed_glm_match_the_short_code_fixture() {
|
|
let path = std::env::var("DS4_GLM_MODEL").unwrap_or_else(|_| {
|
|
"../ds4/models/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf".into()
|
|
});
|
|
let prompt = b"Complete the C statement with the next exact token only:\nreturn snprintf(buf, sizeof(buf), \"%d\", value";
|
|
for streamed in [false, true] {
|
|
let model = Model::open_main(Path::new(&path), ModelChoice::Glm52).unwrap();
|
|
let tokens = model.tokenize(std::str::from_utf8(prompt).unwrap());
|
|
let mut executor = GlmExecutor::open(
|
|
model,
|
|
4096,
|
|
true,
|
|
EngineSsdSettings {
|
|
enabled: streamed,
|
|
cold: false,
|
|
cache_experts: 0,
|
|
cache_bytes: 0,
|
|
full_layers: 0,
|
|
full_layers_set: false,
|
|
preload_experts: 0,
|
|
},
|
|
)
|
|
.unwrap();
|
|
assert_eq!(executor.prefill(&tokens, |_| true).unwrap(), tokens.len());
|
|
let token = executor
|
|
.logits()
|
|
.iter()
|
|
.enumerate()
|
|
.max_by(|a, b| a.1.total_cmp(b.1))
|
|
.unwrap()
|
|
.0 as i32;
|
|
assert_eq!(
|
|
executor.model().token_bytes(token).as_deref(),
|
|
Some(b")".as_slice())
|
|
);
|
|
}
|
|
}
|
|
}
|