From 36946e6e5f3b8d14f6e129fea077ce32ec10e7cb Mon Sep 17 00:00:00 2001 From: Georg Bauer Date: Fri, 24 Jul 2026 18:16:34 +0200 Subject: [PATCH] Complete long-context chat support --- .../down.sql | 2 + .../20260724233000_add_session_context/up.sql | 2 + src/app.rs | 58 +++ src/app/view.rs | 14 + src/database.rs | 26 + src/engine.rs | 14 +- src/engine/metal.rs | 460 +++++++++++++++--- src/schema.rs | 2 + 8 files changed, 515 insertions(+), 63 deletions(-) create mode 100644 migrations/20260724233000_add_session_context/down.sql create mode 100644 migrations/20260724233000_add_session_context/up.sql diff --git a/migrations/20260724233000_add_session_context/down.sql b/migrations/20260724233000_add_session_context/down.sql new file mode 100644 index 0000000..45b20ba --- /dev/null +++ b/migrations/20260724233000_add_session_context/down.sql @@ -0,0 +1,2 @@ +ALTER TABLE sessions DROP COLUMN context_limit; +ALTER TABLE sessions DROP COLUMN context_used; diff --git a/migrations/20260724233000_add_session_context/up.sql b/migrations/20260724233000_add_session_context/up.sql new file mode 100644 index 0000000..7e4e5e0 --- /dev/null +++ b/migrations/20260724233000_add_session_context/up.sql @@ -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); diff --git a/src/app.rs b/src/app.rs index b5e0f13..8eb26d4 100644 --- a/src/app.rs +++ b/src/app.rs @@ -403,6 +403,8 @@ pub(crate) struct App { pub(super) composer: String, pub(super) conversation: Vec, pub(super) generating: bool, + pub(super) context_used: u32, + pub(super) context_limit: u32, #[cfg(target_os = "macos")] generation_worker: Option, error: Option, @@ -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 { 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(); diff --git a/src/app/view.rs b/src/app/view.rs index da40eff..bc28891 100644 --- a/src/app/view.rs +++ b/src/app/view.rs @@ -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( diff --git a/src/database.rs b/src/database.rs index f4b8a1e..46d0a80 100644 --- a/src/database.rs +++ b/src/database.rs @@ -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); diff --git a/src/engine.rs b/src/engine.rs index 1967612..f18dbab 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -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(); } diff --git a/src/engine/metal.rs b/src/engine/metal.rs index f1409ef..9ae2d7e 100644 --- a/src/engine/metal.rs +++ b/src/engine/metal.rs @@ -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, + indexer: Option, 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 { 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 { + fn allocate(model: &Model, context: u32) -> Result { 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, + indexer: Option, } 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::>()?; 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 { + 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, diff --git a/src/schema.rs b/src/schema.rs index 431789c..6c901e3 100644 --- a/src/schema.rs +++ b/src/schema.rs @@ -61,6 +61,8 @@ diesel::table! { id -> Integer, project_id -> Integer, title -> Text, + context_used -> Integer, + context_limit -> Integer, } }