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

Complete long-context chat support

Browse Source
This commit is contained in:
Georg Bauer
2026-07-24 18:16:34 +02:00
parent 5ab110ed0b
commit 36946e6e5f
8 changed files with 515 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

460
src/engine/metal.rs
View File

@@ -1,4 +1,4 @@
use super::gguf::{Gguf, Q8_0, Tensor as GgufTensor};
use super::gguf::{F16, Gguf, Q8_0, Tensor as GgufTensor};
use super::{Model, ModelFamily};
use std::env;
use std::ffi::c_void;
@@ -78,6 +78,13 @@ unsafe extern "C" {
data: *mut c_void,
bytes: u64,
) -> i32;
fn ds4_gpu_tensor_copy(
dst: *mut GpuTensor,
dst_offset: u64,
src: *const GpuTensor,
src_offset: u64,
bytes: u64,
) -> i32;
fn ds4_gpu_tensor_copy_f32_to_f16(
dst: *mut GpuTensor,
dst_offset: u64,
@@ -264,6 +271,46 @@ unsafe extern "C" {
heads: u32,
head_dim: u32,
) -> i32;
fn ds4_gpu_attention_indexed_mixed_batch_heads_tensor(
heads: *mut GpuTensor,
map: *const c_void,
size: u64,
sinks: u64,
q: *const GpuTensor,
raw_kv: *const GpuTensor,
comp_kv: *const GpuTensor,
comp_f16: u32,
topk: *const GpuTensor,
tokens: u32,
pos: u32,
n_raw: u32,
raw_cap: u32,
raw_start: u32,
n_comp: u32,
top_k: u32,
window: u32,
ratio: u32,
heads: u32,
head_dim: u32,
) -> i32;
fn ds4_gpu_indexer_score_one_tensor(
scores: *mut GpuTensor,
q: *const GpuTensor,
weights: *const GpuTensor,
index_comp: *const GpuTensor,
n_comp: u32,
heads: u32,
head_dim: u32,
scale: f32,
) -> i32;
fn ds4_gpu_indexer_topk_tensor(
selected: *mut GpuTensor,
scores: *const GpuTensor,
n_comp: u32,
tokens: u32,
top_k: u32,
) -> i32;
fn ds4_gpu_dsv4_indexer_qat_tensor(x: *mut GpuTensor, rows: u32, head_dim: u32) -> i32;
fn ds4_gpu_compressor_update_tensor(
kv: *const GpuTensor,
score: *const GpuTensor,
@@ -579,6 +626,7 @@ struct Layer {
attn_output_a: Weight,
attn_output_b: Weight,
attn_compressor: Option<CompressorWeights>,
indexer: Option<IndexerWeights>,
hc_ffn_fn: Weight,
hc_ffn_scale: Weight,
hc_ffn_base: Weight,
@@ -602,6 +650,13 @@ struct CompressorWeights {
norm: Weight,
}
#[derive(Clone, Copy)]
struct IndexerWeights {
q: Weight,
proj: Weight,
compressor: CompressorWeights,
}
impl Layer {
fn bind(model: &Gguf, index: u32) -> Result<Self, String> {
let required = |suffix: &str| Weight::bind(model, &format!("blk.{index}.{suffix}"));
@@ -616,6 +671,20 @@ impl Layer {
})
})
.transpose()?;
let indexer = (compression_ratio(index) == 4)
.then(|| {
Ok::<_, String>(IndexerWeights {
q: required("indexer.attn_q_b.weight")?,
proj: required("indexer.proj.weight")?,
compressor: CompressorWeights {
ape: required("indexer_compressor_ape.weight")?,
kv: required("indexer_compressor_kv.weight")?,
gate: required("indexer_compressor_gate.weight")?,
norm: required("indexer_compressor_norm.weight")?,
},
})
})
.transpose()?;
let layer = Self {
hc_attn_fn: required("hc_attn_fn.weight")?,
hc_attn_scale: required("hc_attn_scale.weight")?,
@@ -630,6 +699,7 @@ impl Layer {
attn_output_a: required("attn_output_a.weight")?,
attn_output_b: required("attn_output_b.weight")?,
attn_compressor,
indexer,
hc_ffn_fn: required("hc_ffn_fn.weight")?,
hc_ffn_scale: required("hc_ffn_scale.weight")?,
hc_ffn_base: required("hc_ffn_base.weight")?,
@@ -656,6 +726,11 @@ impl Layer {
return Err(format!("Flash Metal path requires Q8_0 {name} weights"));
}
}
if let Some(indexer) = layer.indexer
&& (!matches!(indexer.q.kind, F16 | Q8_0) || indexer.proj.kind != F16)
{
return Err("Flash Metal path requires F16/Q8 indexer weights".into());
}
Ok(layer)
}
}
@@ -711,6 +786,10 @@ struct Scratch {
compressed_kv: Buffer,
compressed_score: Buffer,
compressed_stage: Buffer,
indexer_q: Buffer,
indexer_weights: Buffer,
indexer_scores: Buffer,
indexer_selected: Buffer,
heads: Buffer,
attention_low: Buffer,
attention_out: Buffer,
@@ -735,7 +814,7 @@ struct Scratch {
}
impl Scratch {
fn allocate(model: &Model) -> Result<Self, String> {
fn allocate(model: &Model, context: u32) -> Result<Self, String> {
let shape = model.shape;
let hc_dim = shape.hc * shape.embd;
let mix_hc = 2 * shape.hc + shape.hc * shape.hc;
@@ -758,6 +837,10 @@ impl Scratch {
compressed_kv: Buffer::floats(2 * shape.head_dim)?,
compressed_score: Buffer::floats(2 * shape.head_dim)?,
compressed_stage: Buffer::floats(shape.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 / 4 + 2))?,
indexer_selected: Buffer::bytes(shape.indexer_top_k * 4)?,
heads: Buffer::floats(q_dim)?,
attention_low: Buffer::floats(low_dim)?,
attention_out: Buffer::floats(shape.embd)?,
@@ -796,6 +879,7 @@ pub(super) struct Session {
struct LayerState {
raw_cache: Buffer,
compression: Option<CompressionState>,
indexer: Option<CompressionState>,
}
struct CompressionState {
@@ -828,9 +912,27 @@ impl LayerState {
rows: 0,
})
};
let indexer = if ratio == 4 {
let width = 2 * shape.indexer_head_dim;
let state_rows = 2 * u64::from(ratio);
let state_kv = Buffer::floats(width * state_rows)?;
let state_score = Buffer::floats(width * state_rows)?;
state_kv.fill(0.0, width * state_rows)?;
state_score.fill(f32::NEG_INFINITY, width * state_rows)?;
Some(CompressionState {
ratio,
cache: Buffer::floats(u64::from(context / ratio + 2) * shape.indexer_head_dim)?,
state_kv,
state_score,
rows: 0,
})
} else {
None
};
Ok(Self {
raw_cache: Buffer::floats(raw_cap as u64 * shape.head_dim)?,
compression,
indexer,
})
}
}
@@ -840,15 +942,12 @@ impl Session {
if context == 0 {
return Err("context must contain at least one token".into());
}
// ponytail: the ratio-4 indexer is needed above 2K; cap here until that
// sparse selection path is ported.
let context = context.min(2_048);
let raw_cap = model.shape.sliding_window as u32;
let raw_cap = (model.shape.sliding_window as u32).min(context);
let layers = (0..model.shape.layers)
.map(|index| LayerState::allocate(model, index, context, raw_cap))
.collect::<Result<_, _>>()?;
Ok(Self {
scratch: Scratch::allocate(model)?,
scratch: Scratch::allocate(model, context)?,
layers,
context,
raw_cap,
@@ -1202,35 +1301,169 @@ fn encode_layer(
attn_factor,
)?;
}
if let (Some(weights), Some(indexer)) = (w.indexer, state.indexer.as_mut()) {
update_indexer_compression(
s,
indexer,
weights.compressor,
shape,
map,
size,
pos,
original,
freq_base,
freq_scale,
ext,
attn_factor,
)?;
}
let (comp_cache, n_comp) = state
.compression
.as_ref()
.map_or((std::ptr::null(), 0), |compression| {
(compression.cache.raw().cast_const(), compression.rows)
});
call(
unsafe {
ds4_gpu_attention_decode_heads_tensor(
s.heads.raw(),
map,
size,
w.attn_sinks.offset,
s.q.raw(),
raw_cache,
n_raw,
raw_cap,
raw_start,
comp_cache,
1,
n_comp,
std::ptr::null(),
0,
shape.heads as u32,
shape.head_dim as u32,
)
},
"attention",
)?;
let indexed = if n_comp > 1_024 {
match (w.indexer, state.indexer.as_ref()) {
(Some(weights), Some(indexer)) if indexer.rows > shape.indexer_top_k as u32 => {
matmul(
&s.indexer_q,
weights.q,
shape.lora_q,
shape.indexer_heads * shape.indexer_head_dim,
&s.q_rank_norm,
map,
size,
)?;
call(
unsafe {
ds4_gpu_rope_tail_tensor(
s.indexer_q.raw(),
1,
shape.indexer_heads as u32,
shape.indexer_head_dim as u32,
shape.rot as u32,
pos,
original,
false,
freq_base,
freq_scale,
ext,
attn_factor,
shape.rope_beta_fast,
shape.rope_beta_slow,
)
},
"indexer Q RoPE",
)?;
call(
unsafe {
ds4_gpu_dsv4_indexer_qat_tensor(
s.indexer_q.raw(),
shape.indexer_heads as u32,
shape.indexer_head_dim as u32,
)
},
"indexer Q quantization",
)?;
f16(
&s.indexer_weights,
weights.proj,
shape.embd,
shape.indexer_heads,
&s.norm,
map,
size,
)?;
call(
unsafe {
ds4_gpu_indexer_score_one_tensor(
s.indexer_scores.raw(),
s.indexer_q.raw(),
s.indexer_weights.raw(),
indexer.cache.raw(),
indexer.rows,
shape.indexer_heads as u32,
shape.indexer_head_dim as u32,
((shape.indexer_heads * shape.indexer_head_dim) as f32)
.sqrt()
.recip(),
)
},
"indexer scoring",
)?;
call(
unsafe {
ds4_gpu_indexer_topk_tensor(
s.indexer_selected.raw(),
s.indexer_scores.raw(),
indexer.rows,
1,
shape.indexer_top_k as u32,
)
},
"indexer top-k selection",
)?;
Some(indexer)
}
_ => None,
}
} else {
None
};
if indexed.is_some() {
call(
unsafe {
ds4_gpu_attention_indexed_mixed_batch_heads_tensor(
s.heads.raw(),
map,
size,
w.attn_sinks.offset,
s.q.raw(),
raw_cache,
comp_cache,
1,
s.indexer_selected.raw(),
1,
pos,
n_raw,
raw_cap,
raw_start,
n_comp,
shape.indexer_top_k as u32,
shape.sliding_window as u32,
compression_ratio(layer),
shape.heads as u32,
shape.head_dim as u32,
)
},
"indexed attention",
)?;
} else {
call(
unsafe {
ds4_gpu_attention_decode_heads_tensor(
s.heads.raw(),
map,
size,
w.attn_sinks.offset,
s.q.raw(),
raw_cache,
n_raw,
raw_cap,
raw_start,
comp_cache,
1,
n_comp,
std::ptr::null(),
0,
shape.heads as u32,
shape.head_dim as u32,
)
},
"attention",
)?;
}
call(
unsafe {
ds4_gpu_rope_tail_tensor(
@@ -1454,7 +1687,125 @@ fn update_compression(
ext: f32,
attn_factor: f32,
) -> Result<(), String> {
let width = if state.ratio == 4 { 2 } else { 1 } * shape.head_dim as u32;
let emit = update_compressor_stage(
s,
state,
weights,
shape,
map,
size,
pos,
original,
freq_base,
freq_scale,
ext,
attn_factor,
shape.head_dim as u32,
)?;
if emit {
call(
unsafe {
ds4_gpu_dsv4_fp8_kv_quantize_tensor(
s.compressed_stage.raw(),
1,
shape.head_dim as u32,
shape.rot as u32,
)
},
"compressed KV quantization",
)?;
call(
unsafe {
ds4_gpu_tensor_copy_f32_to_f16(
state.cache.raw(),
state.rows as u64 * shape.head_dim * 2,
s.compressed_stage.raw(),
0,
shape.head_dim,
)
},
"compressed KV cache write",
)?;
state.rows += 1;
}
Ok(())
}
#[allow(clippy::too_many_arguments)]
fn update_indexer_compression(
s: &Scratch,
state: &mut CompressionState,
weights: CompressorWeights,
shape: super::Shape,
map: *const c_void,
size: u64,
pos: u32,
original: u32,
freq_base: f32,
freq_scale: f32,
ext: f32,
attn_factor: f32,
) -> Result<(), String> {
let emit = update_compressor_stage(
s,
state,
weights,
shape,
map,
size,
pos,
original,
freq_base,
freq_scale,
ext,
attn_factor,
shape.indexer_head_dim as u32,
)?;
if emit {
call(
unsafe {
ds4_gpu_dsv4_indexer_qat_tensor(
s.compressed_stage.raw(),
1,
shape.indexer_head_dim as u32,
)
},
"compressed index quantization",
)?;
call(
unsafe {
ds4_gpu_tensor_copy(
state.cache.raw(),
state.rows as u64 * shape.indexer_head_dim * 4,
s.compressed_stage.raw(),
0,
shape.indexer_head_dim * 4,
)
},
"compressed index cache write",
)?;
state.rows += 1;
}
Ok(())
}
#[allow(clippy::too_many_arguments)]
fn update_compressor_stage(
s: &Scratch,
state: &CompressionState,
weights: CompressorWeights,
shape: super::Shape,
map: *const c_void,
size: u64,
pos: u32,
original: u32,
freq_base: f32,
freq_scale: f32,
ext: f32,
attn_factor: f32,
head_dim: u32,
) -> Result<bool, String> {
let width = if state.ratio == 4 { 2 } else { 1 } * head_dim;
let fused = unsafe {
ds4_gpu_matmul_f16_pair_compressor_store_tensor(
s.compressed_kv.raw(),
@@ -1496,7 +1847,6 @@ fn update_compression(
"compressor projection",
)?;
}
let emit = (pos + 1).is_multiple_of(state.ratio);
call(
unsafe {
ds4_gpu_compressor_update_tensor(
@@ -1511,7 +1861,7 @@ fn update_compression(
weights.ape.kind,
weights.norm.offset,
weights.norm.kind,
shape.head_dim as u32,
head_dim,
state.ratio,
pos,
0,
@@ -1529,33 +1879,7 @@ fn update_compression(
},
"compressor update",
)?;
if emit {
call(
unsafe {
ds4_gpu_dsv4_fp8_kv_quantize_tensor(
s.compressed_stage.raw(),
1,
shape.head_dim as u32,
shape.rot as u32,
)
},
"compressed KV quantization",
)?;
call(
unsafe {
ds4_gpu_tensor_copy_f32_to_f16(
state.cache.raw(),
state.rows as u64 * shape.head_dim * 2,
s.compressed_stage.raw(),
0,
shape.head_dim,
)
},
"compressed KV cache write",
)?;
state.rows += 1;
}
Ok(())
Ok((pos + 1).is_multiple_of(state.ratio))
}
fn compression_ratio(layer: u32) -> u32 {
@@ -1690,6 +2014,22 @@ fn f16(
)
}
fn matmul(
out: &Buffer,
weight: Weight,
input: u64,
output: u64,
x: &Buffer,
map: *const c_void,
size: u64,
) -> Result<(), String> {
match weight.kind {
F16 => f16(out, weight, input, output, x, map, size),
Q8_0 => q8(out, weight, input, output, x, map, size),
kind => Err(format!("unsupported Metal projection type {kind}")),
}
}
fn q8(
out: &Buffer,
weight: Weight,