hugo/DS4Server
1
0
Fork 0
Code Issues 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

2
migrations/20260724233000_add_session_context/down.sql Normal file
View File

@@ -0,0 +1,2 @@
ALTER TABLE sessions DROP COLUMN context_limit;
ALTER TABLE sessions DROP COLUMN context_used;

2
migrations/20260724233000_add_session_context/up.sql Normal file
View File

@@ -0,0 +1,2 @@
ALTER TABLE sessions ADD COLUMN context_used INTEGER NOT NULL DEFAULT 0 CHECK (context_used >= 0);
ALTER TABLE sessions ADD COLUMN context_limit INTEGER NOT NULL DEFAULT 0 CHECK (context_limit >= 0);

58
src/app.rs
View File

@@ -403,6 +403,8 @@ pub(crate) struct App {
pub(super) composer: String,
pub(super) conversation: Vec<ChatMessage>,
pub(super) generating: bool,
pub(super) context_used: u32,
pub(super) context_limit: u32,
#[cfg(target_os = "macos")]
generation_worker: Option<GenerationWorker>,
error: Option<String>,
@@ -464,6 +466,7 @@ enum GenerationCommand {
enum GenerationEvent {
Loading,
Chunk { reasoning: bool, content: String },
Context { used: u32, limit: u32 },
Finished(Result<(), String>),
}
@@ -572,6 +575,7 @@ impl App {
Ok(draft) => draft,
Err(error) => return Self::failed(error, main_window),
};
let context_limit = preferences.context_tokens.max(0) as u32;
Self {
main_window,
model_manager_window: None,
@@ -593,6 +597,8 @@ impl App {
composer: String::new(),
conversation: Vec::new(),
generating: false,
context_used: 0,
context_limit,
#[cfg(target_os = "macos")]
generation_worker: None,
error: None,
@@ -608,6 +614,7 @@ impl App {
let preferences = AppPreferences::default();
let preference_draft = PreferenceDraft::from_saved(&preferences)
.expect("default preferences must use a supported model");
let context_limit = preferences.context_tokens.max(0) as u32;
Self {
main_window,
model_manager_window: None,
@@ -629,6 +636,8 @@ impl App {
composer: String::new(),
conversation: Vec::new(),
generating: false,
context_used: 0,
context_limit,
#[cfg(target_os = "macos")]
generation_worker: None,
error: Some(format!("Could not open the project database: {error}")),
@@ -953,6 +962,8 @@ impl App {
self.selected_session = None;
self.conversation.clear();
self.composer.clear();
self.context_used = 0;
self.context_limit = self.preferences.context_tokens.max(0) as u32;
self.error = None;
}
Message::DeleteProject(project_id) => {
@@ -969,6 +980,7 @@ impl App {
self.selected_session = None;
self.conversation.clear();
self.composer.clear();
self.context_used = 0;
}
self.reload_projects();
}
@@ -986,6 +998,12 @@ impl App {
if self.selected_session == Some(session_id) {
return Task::none();
}
let saved_context = self
.projects
.iter()
.flat_map(|project| &project.sessions)
.find(|session| session.id == session_id)
.map(|session| (session.context_used, session.context_limit));
let Some(database) = &mut self.database else {
return Task::none();
};
@@ -995,6 +1013,13 @@ impl App {
self.composer.clear();
self.selected_project = Some(project_id);
self.selected_session = Some(session_id);
let (used, limit) = saved_context.unwrap_or_default();
self.context_used = used.max(0) as u32;
self.context_limit = if limit > 0 {
limit as u32
} else {
self.preferences.context_tokens.max(0) as u32
};
self.error = None;
}
Err(error) => {
@@ -1015,6 +1040,7 @@ impl App {
self.selected_session = None;
self.conversation.clear();
self.composer.clear();
self.context_used = 0;
}
self.reload_projects();
}
@@ -1277,6 +1303,8 @@ impl App {
self.selected_session = Some(session.id);
self.conversation.clear();
self.composer.clear();
self.context_used = 0;
self.context_limit = self.preferences.context_tokens.max(0) as u32;
self.error = None;
self.reload_projects();
}
@@ -1473,6 +1501,8 @@ impl App {
#[cfg(target_os = "macos")]
let mut transcript_changed = false;
#[cfg(target_os = "macos")]
let mut context_changed = false;
#[cfg(target_os = "macos")]
loop {
match worker.events.try_recv() {
Ok(GenerationEvent::Loading) => {}
@@ -1484,6 +1514,11 @@ impl App {
transcript_changed = true;
}
}
Ok(GenerationEvent::Context { used, limit }) => {
self.context_used = used;
self.context_limit = limit;
context_changed = true;
}
Ok(GenerationEvent::Finished(result)) => {
self.generating = false;
worker.cancel = None;
@@ -1521,6 +1556,25 @@ impl App {
}
self.error = Some(format!("Could not save generated chat text: {error}"));
}
#[cfg(target_os = "macos")]
if context_changed
&& let Some(session_id) = self.selected_session
&& let Some(database) = &mut self.database
{
if let Err(error) =
database.update_session_context(session_id, self.context_used, self.context_limit)
{
self.error = Some(format!("Could not save context usage: {error}"));
} else if let Some(session) = self
.projects
.iter_mut()
.flat_map(|project| &mut project.sessions)
.find(|session| session.id == session_id)
{
session.context_used = self.context_used as i32;
session.context_limit = self.context_limit as i32;
}
}
}
}
@@ -1567,6 +1621,10 @@ fn spawn_generation_worker() -> Result<GenerationWorker, String> {
let _ = event_sender
.send(GenerationEvent::Chunk { reasoning, content });
},
|used, limit| {
let _ =
event_sender.send(GenerationEvent::Context { used, limit });
},
);
let _ = event_sender.send(GenerationEvent::Finished(result));
last_used = Instant::now();

14
src/app/view.rs
View File

@@ -298,13 +298,27 @@ impl App {
.padding(8)
.on_press(Message::SubmitPrompt)
};
let context_fraction = if self.context_limit == 0 {
0.0
} else {
self.context_used.min(self.context_limit) as f32 / self.context_limit as f32
};
let conversation = column![
scrollable(messages).height(Length::Fill),
container(
column![
composer,
progress_bar(0.0..=1.0, context_fraction).height(3),
row![
icon(ICON_PAPERCLIP, 19),
text(format!(
"{} / {} tokens ({:.0}%)",
self.context_used,
self.context_limit,
context_fraction * 100.0
))
.size(11)
.color(muted_text()),
Space::with_width(Length::Fill),
icon(ICON_MODEL, 16),
text(

26
src/database.rs
View File

@@ -260,6 +260,8 @@ pub struct Session {
pub id: i32,
pub project_id: i32,
pub title: String,
pub context_used: i32,
pub context_limit: i32,
}
#[derive(Insertable)]
@@ -476,6 +478,24 @@ impl Database {
.map_err(|error| error.to_string())
}
pub fn update_session_context(
&mut self,
session_id: i32,
used: u32,
limit: u32,
) -> Result<(), String> {
let used = i32::try_from(used).map_err(|_| "Used context is too large to save")?;
let limit = i32::try_from(limit).map_err(|_| "Context limit is too large to save")?;
diesel::update(sessions::table.find(session_id))
.set((
sessions::context_used.eq(used),
sessions::context_limit.eq(limit),
))
.execute(&mut self.connection)
.map(|_| ())
.map_err(|error| error.to_string())
}
pub fn start_chat_turn(
&mut self,
session_id: i32,
@@ -648,9 +668,15 @@ mod tests {
database
.update_message(assistant.id, Some("Reasoning"), true, "Answer")
.unwrap();
database
.update_session_context(session.id, 1_234, 65_536)
.unwrap();
drop(database);
let mut reopened = Database::open(&path).unwrap();
let projects = reopened.load_projects().unwrap();
assert_eq!(projects[0].sessions[0].context_used, 1_234);
assert_eq!(projects[0].sessions[0].context_limit, 65_536);
let messages = reopened.load_messages(session.id).unwrap();
assert_eq!(messages.len(), 2);
assert!(messages[0].user);

14
src/engine.rs
View File

@@ -324,8 +324,10 @@ impl Generator {
settings: &TurnSettings,
cancelled: &AtomicBool,
mut emit: impl FnMut(bool, String),
mut progress: impl FnMut(u32, u32),
) -> Result<(), String> {
self.executor.reset()?;
progress(self.executor.position(), self.executor.context());
let tokens = self.executor.model().render_conversation(
&settings.system_prompt,
messages,
@@ -337,17 +339,21 @@ impl Generator {
let max_context = self.executor.context() as usize;
if tokens.len() >= max_context {
return Err(format!(
"the conversation uses {} tokens; the current Rust attention port supports fewer than {max_context}",
"the conversation uses {} tokens; the configured context holds fewer than {max_context}",
tokens.len()
));
}
let mut rng = Rng::new(settings.seed.unwrap_or(0x4453_3453_4552_5645));
let mut reasoning = settings.reasoning_mode != ReasoningMode::Direct;
for token in tokens {
let prompt_tokens = tokens.len();
for (index, token) in tokens.into_iter().enumerate() {
if cancelled.load(Ordering::Relaxed) {
return Ok(());
}
self.executor.eval(token)?;
if (index + 1).is_multiple_of(16) || index + 1 == prompt_tokens {
progress(self.executor.position(), self.executor.context());
}
}
for _ in 0..settings
.max_generated_tokens
@@ -380,6 +386,7 @@ impl Generator {
emit(reasoning, String::from_utf8_lossy(&bytes).into_owned());
}
self.executor.eval(token)?;
progress(self.executor.position(), self.executor.context());
}
Ok(())
}
@@ -497,7 +504,8 @@ mod sampling_tests {
ReasoningMode::Direct,
);
assert_eq!(tokens.len(), 10);
let mut executor = metal::Executor::open(model, 128, false).unwrap();
let mut executor = metal::Executor::open(model, 32_768, false).unwrap();
assert_eq!(executor.context(), 32_768);
for token in tokens {
executor.eval(token).unwrap();
}

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,

2
src/schema.rs
View File

@@ -61,6 +61,8 @@ diesel::table! {
id -> Integer,
project_id -> Integer,
title -> Text,
context_used -> Integer,
context_limit -> Integer,
}
}