diff --git a/src/app.rs b/src/app.rs index 1f46f84..b3d72df 100644 --- a/src/app.rs +++ b/src/app.rs @@ -24,7 +24,7 @@ use crate::engine::ChatTurn; #[cfg(target_os = "macos")] use crate::metrics::WorkSource; use crate::metrics::{KvCacheReport, Metrics, MetricsSnapshot}; -use crate::model::{self, DownloadOutcome, DownloadProgress, ManagedArtifactId, ModelChoice}; +use crate::model::{DownloadOutcome, DownloadProgress, ManagedArtifactId, ModelChoice}; #[cfg(target_os = "macos")] use crate::runtime::{ ActiveGeneration, CheckpointTarget, CompactionInput, GenerationEvent, GenerationService, @@ -950,6 +950,10 @@ impl App { Ok(message) => message, Err(task) => return task, }; + let message = match self.update_model_manager_message(message) { + Ok(message) => message, + Err(task) => return task, + }; match message { Message::Noop => { #[cfg(target_os = "macos")] @@ -973,7 +977,6 @@ impl App { } } Message::ScrollPreferences(section) => return scroll_preferences_to(section), - Message::OpenModelManager => return self.open_model_manager(), Message::OpenHelp => return self.open_help(), Message::HelpOpened(id) => { if self.help_window == Some(id) { @@ -1012,11 +1015,6 @@ impl App { self.error = Some(format!("Could not export the chat: {error}")); } } - Message::ModelManagerOpened(id) => { - if self.model_manager_window == Some(id) { - return window::gain_focus(id); - } - } Message::WindowOpened(id) => { if id == self.main_window { self.refresh_git_state(); @@ -1144,34 +1142,6 @@ impl App { self.sync_native_menu(); } Message::MetricsTick => self.sample_metrics(), - Message::DownloadArtifact(artifact) => { - self.start_model_operation(artifact, ModelOperation::Download) - } - Message::ValidateArtifact(artifact) => { - self.start_model_operation(artifact, ModelOperation::Validate) - } - Message::DeleteArtifact(artifact) => self.pending_model_delete = Some(artifact), - Message::ConfirmDeleteArtifact => { - if let Some(artifact) = self.pending_model_delete.take() { - match model::delete_managed_artifact(artifact, &models_path()) { - Ok(()) => { - self.model_download = ModelDownload::Idle; - self.error = None; - } - Err(error) => { - self.error = Some(format!("Could not delete {artifact}: {error}")) - } - } - } - } - Message::CancelDeleteArtifact => self.pending_model_delete = None, - Message::StopModelDownload => { - if let ModelDownload::Active(download) = &mut self.model_download { - download.stopping = true; - download.cancel.store(true, Ordering::Relaxed); - } - } - Message::DownloadProgressTick => self.update_download_progress(), Message::ComposerAction(action) => self.composer.perform(action), Message::TranscriptAction(index, action) => { if !action.is_edit() @@ -1833,7 +1803,7 @@ impl App { } } } - _ => unreachable!("preference messages are dispatched before the main update"), + _ => unreachable!("domain messages are dispatched before the main update"), } Task::none() } diff --git a/src/app/model_manager.rs b/src/app/model_manager.rs index 72d3fef..a1e4fca 100644 --- a/src/app/model_manager.rs +++ b/src/app/model_manager.rs @@ -1,4 +1,5 @@ use super::*; +use crate::model; #[derive(Debug)] pub(super) enum ModelDownload { @@ -29,6 +30,50 @@ pub(super) struct ActiveDownload { } impl App { + pub(super) fn update_model_manager_message( + &mut self, + message: Message, + ) -> Result> { + match message { + Message::OpenModelManager => return Err(self.open_model_manager()), + Message::ModelManagerOpened(id) => { + if self.model_manager_window == Some(id) { + return Err(window::gain_focus(id)); + } + } + Message::DownloadArtifact(artifact) => { + self.start_model_operation(artifact, ModelOperation::Download) + } + Message::ValidateArtifact(artifact) => { + self.start_model_operation(artifact, ModelOperation::Validate) + } + Message::DeleteArtifact(artifact) => self.pending_model_delete = Some(artifact), + Message::ConfirmDeleteArtifact => { + if let Some(artifact) = self.pending_model_delete.take() { + match model::delete_managed_artifact(artifact, &models_path()) { + Ok(()) => { + self.model_download = ModelDownload::Idle; + self.error = None; + } + Err(error) => { + self.error = Some(format!("Could not delete {artifact}: {error}")) + } + } + } + } + Message::CancelDeleteArtifact => self.pending_model_delete = None, + Message::StopModelDownload => { + if let ModelDownload::Active(download) = &mut self.model_download { + download.stopping = true; + download.cancel.store(true, Ordering::Relaxed); + } + } + Message::DownloadProgressTick => self.update_download_progress(), + message => return Ok(message), + } + Err(Task::none()) + } + pub(super) fn open_model_manager(&mut self) -> Task { if let Some(id) = self.model_manager_window { return window::gain_focus(id); diff --git a/src/engine.rs b/src/engine.rs index 00f3bf1..af6dcfc 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -7,7 +7,7 @@ mod tokenizer; mod validation; #[cfg(target_os = "macos")] -use crate::metrics::{KvLookup, Metrics}; +use crate::metrics::{KvLookup, Metrics, SsdStats}; use crate::model::ModelChoice; #[cfg(target_os = "macos")] use crate::settings::TurnSettings; @@ -509,37 +509,7 @@ impl Generator { settings.steering.clone(), settings.diagnostics.expert_profile_path.as_deref(), )?; - let stats = executor.execution_stats(); - metrics.speculative_stats( - stats.speculative_mode, - stats.speculative_cycles, - stats.drafted_tokens, - stats.accepted_draft_tokens, - stats.verifier_passes, - stats.verifier_ms, - ); - metrics.ssd_stats( - stats.ssd_enabled, - stats.ssd_resident_bytes, - stats.ssd_cache_bytes, - stats.ssd_cache_experts, - stats.ssd_cache_entries, - stats.ssd_preloaded_experts, - stats.ssd_cache_hits, - stats.ssd_cache_misses, - stats.ssd_cache_evictions, - stats.ssd_cache_wraps, - stats.ssd_buffer_allocs, - stats.ssd_buffer_reuses, - stats.ssd_pread_bytes, - stats.ssd_pread_ms, - stats.ssd_evict_advise_bytes, - stats.ssd_willneed_advise_bytes, - stats.ssd_selected_requests, - stats.ssd_requested_bytes, - stats.ssd_wait_ms, - ); - Ok(Self { + let generator = Self { executor, _simulated_memory: simulated_memory, checkpoint: None, @@ -553,7 +523,9 @@ impl Generator { .filter(|limit| *limit > 0) .unwrap_or(1) .saturating_sub(1), - }) + }; + generator.publish_execution_stats(); + Ok(generator) } pub(crate) fn summary(&self) -> ModelSummary { @@ -622,27 +594,27 @@ impl Generator { stats.verifier_passes, stats.verifier_ms, ); - self.metrics.ssd_stats( - stats.ssd_enabled, - stats.ssd_resident_bytes, - stats.ssd_cache_bytes, - stats.ssd_cache_experts, - stats.ssd_cache_entries, - stats.ssd_preloaded_experts, - stats.ssd_cache_hits, - stats.ssd_cache_misses, - stats.ssd_cache_evictions, - stats.ssd_cache_wraps, - stats.ssd_buffer_allocs, - stats.ssd_buffer_reuses, - stats.ssd_pread_bytes, - stats.ssd_pread_ms, - stats.ssd_evict_advise_bytes, - stats.ssd_willneed_advise_bytes, - stats.ssd_selected_requests, - stats.ssd_requested_bytes, - stats.ssd_wait_ms, - ); + self.metrics.ssd_stats(SsdStats { + enabled: stats.ssd_enabled, + resident_bytes: stats.ssd_resident_bytes, + cache_bytes: stats.ssd_cache_bytes, + cache_experts: stats.ssd_cache_experts, + cache_entries: stats.ssd_cache_entries, + preloaded_experts: stats.ssd_preloaded_experts, + cache_hits: stats.ssd_cache_hits, + cache_misses: stats.ssd_cache_misses, + cache_evictions: stats.ssd_cache_evictions, + cache_wraps: stats.ssd_cache_wraps, + buffer_allocs: stats.ssd_buffer_allocs, + buffer_reuses: stats.ssd_buffer_reuses, + pread_bytes: stats.ssd_pread_bytes, + pread_ms: stats.ssd_pread_ms, + evict_advise_bytes: stats.ssd_evict_advise_bytes, + willneed_advise_bytes: stats.ssd_willneed_advise_bytes, + selected_requests: stats.ssd_selected_requests, + requested_bytes: stats.ssd_requested_bytes, + wait_ms: stats.ssd_wait_ms, + }); } pub(crate) fn generate_transient( diff --git a/src/engine/metal.rs b/src/engine/metal.rs index e1da086..652f6f9 100644 --- a/src/engine/metal.rs +++ b/src/engine/metal.rs @@ -490,16 +490,11 @@ impl Dspark { }, "capturing batched DSpark target hidden states", )?; - call( - unsafe { - ds4_gpu_tensor_copy( - self.target_hidden.raw(), - u64::from(slot) * shape.embd * 4, - batch.raw(), - u64::from(rows - 1) * shape.embd * 4, - shape.embd * 4, - ) - }, + self.target_hidden.copy_from( + u64::from(slot) * shape.embd * 4, + &batch, + u64::from(rows - 1) * shape.embd * 4, + shape.embd * 4, "capturing final DSpark target hidden state", )?; self.capture_mask |= 1 << slot; @@ -881,32 +876,22 @@ impl Dspark { self.eval_stage(support, stage, pos, raw_cap, shape)?; if stage + 1 < self.weights.len() { let commands = Commands::begin()?; - call( - unsafe { - ds4_gpu_tensor_copy( - self.stage_input_hc.raw(), - shape.hc * shape.embd * 4, - self.scratch.next_hc.raw(), - 0, - u64::from(self.config.block_size) * shape.hc * shape.embd * 4, - ) - }, + self.stage_input_hc.copy_from( + shape.hc * shape.embd * 4, + &self.scratch.next_hc, + 0, + u64::from(self.config.block_size) * shape.hc * shape.embd * 4, "feeding the next DSpark stage", )?; commands.finish()?; } } let commands = Commands::begin()?; - call( - unsafe { - ds4_gpu_tensor_copy( - self.stage_output_hc.raw(), - 0, - self.scratch.next_hc.raw(), - 0, - u64::from(self.config.block_size) * shape.hc * shape.embd * 4, - ) - }, + self.stage_output_hc.copy_from( + 0, + &self.scratch.next_hc, + 0, + u64::from(self.config.block_size) * shape.hc * shape.embd * 4, "capturing DSpark stage output", )?; commands.finish()?; @@ -2952,14 +2937,18 @@ impl DeepSeekExecutor { fn snapshot(state: &CompressionState, bytes: u64) -> Result { let state_kv = Buffer::bytes(bytes)?; let state_score = Buffer::bytes(bytes)?; - call( - unsafe { ds4_gpu_tensor_copy(state_kv.raw(), 0, state.state_kv.raw(), 0, bytes) }, + state_kv.copy_from( + 0, + &state.state_kv, + 0, + bytes, "saving speculative compressor KV state", )?; - call( - unsafe { - ds4_gpu_tensor_copy(state_score.raw(), 0, state.state_score.raw(), 0, bytes) - }, + state_score.copy_from( + 0, + &state.state_score, + 0, + bytes, "saving speculative compressor score state", )?; Ok(CompressionFrontier { @@ -3007,10 +2996,11 @@ impl DeepSeekExecutor { .map(|dspark| { let bytes = dspark.config.target_layers.len() as u64 * shape.embd * 4; let saved = Buffer::bytes(bytes)?; - call( - unsafe { - ds4_gpu_tensor_copy(saved.raw(), 0, dspark.target_hidden.raw(), 0, bytes) - }, + saved.copy_from( + 0, + &dspark.target_hidden, + 0, + bytes, "saving speculative DSpark target state", )?; Ok::<_, String>(saved) @@ -3034,28 +3024,18 @@ impl DeepSeekExecutor { state: &mut CompressionState, saved: &CompressionFrontier, ) -> Result<(), String> { - call( - unsafe { - ds4_gpu_tensor_copy( - state.state_kv.raw(), - 0, - saved.state_kv.raw(), - 0, - saved.bytes, - ) - }, + state.state_kv.copy_from( + 0, + &saved.state_kv, + 0, + saved.bytes, "restoring speculative compressor KV state", )?; - call( - unsafe { - ds4_gpu_tensor_copy( - state.state_score.raw(), - 0, - saved.state_score.raw(), - 0, - saved.bytes, - ) - }, + state.state_score.copy_from( + 0, + &saved.state_score, + 0, + saved.bytes, "restoring speculative compressor score state", )?; state.rows = saved.rows; @@ -3080,10 +3060,11 @@ impl DeepSeekExecutor { } if let (Some(dspark), Some(saved)) = (&mut self.dspark, &frontier.dspark_target_hidden) { let bytes = dspark.config.target_layers.len() as u64 * self.model.shape.embd * 4; - call( - unsafe { - ds4_gpu_tensor_copy(dspark.target_hidden.raw(), 0, saved.raw(), 0, bytes) - }, + dspark.target_hidden.copy_from( + 0, + saved, + 0, + bytes, "restoring speculative DSpark target state", )?; dspark.capture_mask = frontier.dspark_capture_mask; @@ -3418,16 +3399,11 @@ impl DeepSeekExecutor { None, Some(mtp.raw_rows), )?; - call( - unsafe { - ds4_gpu_tensor_copy( - mtp.next_hc.raw(), - 0, - scratch.current_hc.raw(), - 0, - shape.hc * shape.embd * 4, - ) - }, + mtp.next_hc.copy_from( + 0, + &scratch.current_hc, + 0, + shape.hc * shape.embd * 4, "MTP hidden-state capture", )?; encode_mtp_output( @@ -3683,16 +3659,11 @@ impl DeepSeekExecutor { let mut tops = Vec::with_capacity(output_rows as usize); for row in first_output..rows { let commands = Commands::begin()?; - call( - unsafe { - ds4_gpu_tensor_copy( - self.session.scratch.current_hc.raw(), - 0, - batch.current_hc.raw(), - u64::from(row) * shape.hc * shape.embd * 4, - shape.hc * shape.embd * 4, - ) - }, + self.session.scratch.current_hc.copy_from( + 0, + &batch.current_hc, + u64::from(row) * shape.hc * shape.embd * 4, + shape.hc * shape.embd * 4, "selecting a prefill output row", )?; encode_output(&self.session.scratch, &self.weights, shape, map, size)?; @@ -6391,16 +6362,11 @@ fn update_indexer_compression( }, "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, - ) - }, + state.cache.copy_from( + state.rows as u64 * shape.indexer_head_dim * 4, + &s.compressed_stage, + 0, + shape.indexer_head_dim * 4, "compressed index cache write", )?; state.rows += 1; diff --git a/src/engine/metal/gpu.rs b/src/engine/metal/gpu.rs index 7000241..1828eab 100644 --- a/src/engine/metal/gpu.rs +++ b/src/engine/metal/gpu.rs @@ -1543,6 +1543,28 @@ impl Buffer { ) } + pub(super) fn copy_from( + &self, + destination_offset: u64, + source: &Self, + source_offset: u64, + bytes: u64, + operation: &str, + ) -> Result<(), String> { + call( + unsafe { + ds4_gpu_tensor_copy( + self.raw(), + destination_offset, + source.raw(), + source_offset, + bytes, + ) + }, + operation, + ) + } + pub(super) fn raw(&self) -> *mut GpuTensor { self.0.as_ptr() } diff --git a/src/metrics.rs b/src/metrics.rs index b6a968d..dfc3324 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -56,6 +56,29 @@ pub(crate) enum KvLookup { Invalid, } +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub(crate) struct SsdStats { + pub(crate) enabled: bool, + pub(crate) resident_bytes: u64, + pub(crate) cache_bytes: u64, + pub(crate) cache_experts: u64, + pub(crate) cache_entries: u64, + pub(crate) preloaded_experts: u64, + pub(crate) cache_hits: u64, + pub(crate) cache_misses: u64, + pub(crate) cache_evictions: u64, + pub(crate) cache_wraps: u64, + pub(crate) buffer_allocs: u64, + pub(crate) buffer_reuses: u64, + pub(crate) pread_bytes: u64, + pub(crate) pread_ms: u64, + pub(crate) evict_advise_bytes: u64, + pub(crate) willneed_advise_bytes: u64, + pub(crate) selected_requests: u64, + pub(crate) requested_bytes: u64, + pub(crate) wait_ms: u64, +} + #[derive(Clone, Debug, Default)] pub(crate) struct MetricsSnapshot { pub(crate) uptime_seconds: u64, @@ -468,59 +491,42 @@ impl Metrics { self.verifier_ms.store(verifier_ms, Ordering::Relaxed); } - #[allow(clippy::too_many_arguments)] - pub(crate) fn ssd_stats( - &self, - enabled: bool, - resident_bytes: u64, - cache_bytes: u64, - cache_experts: u64, - cache_entries: u64, - preloaded_experts: u64, - cache_hits: u64, - cache_misses: u64, - cache_evictions: u64, - cache_wraps: u64, - buffer_allocs: u64, - buffer_reuses: u64, - pread_bytes: u64, - pread_ms: u64, - evict_advise_bytes: u64, - willneed_advise_bytes: u64, - selected_requests: u64, - requested_bytes: u64, - wait_ms: u64, - ) { - self.ssd_enabled.store(enabled, Ordering::Relaxed); + pub(crate) fn ssd_stats(&self, stats: SsdStats) { + self.ssd_enabled.store(stats.enabled, Ordering::Relaxed); self.ssd_resident_bytes - .store(resident_bytes, Ordering::Relaxed); - self.ssd_cache_bytes.store(cache_bytes, Ordering::Relaxed); + .store(stats.resident_bytes, Ordering::Relaxed); + self.ssd_cache_bytes + .store(stats.cache_bytes, Ordering::Relaxed); self.ssd_cache_experts - .store(cache_experts, Ordering::Relaxed); + .store(stats.cache_experts, Ordering::Relaxed); self.ssd_cache_entries - .store(cache_entries, Ordering::Relaxed); + .store(stats.cache_entries, Ordering::Relaxed); self.ssd_preloaded_experts - .store(preloaded_experts, Ordering::Relaxed); - self.ssd_cache_hits.store(cache_hits, Ordering::Relaxed); - self.ssd_cache_misses.store(cache_misses, Ordering::Relaxed); + .store(stats.preloaded_experts, Ordering::Relaxed); + self.ssd_cache_hits + .store(stats.cache_hits, Ordering::Relaxed); + self.ssd_cache_misses + .store(stats.cache_misses, Ordering::Relaxed); self.ssd_cache_evictions - .store(cache_evictions, Ordering::Relaxed); - self.ssd_cache_wraps.store(cache_wraps, Ordering::Relaxed); + .store(stats.cache_evictions, Ordering::Relaxed); + self.ssd_cache_wraps + .store(stats.cache_wraps, Ordering::Relaxed); self.ssd_buffer_allocs - .store(buffer_allocs, Ordering::Relaxed); + .store(stats.buffer_allocs, Ordering::Relaxed); self.ssd_buffer_reuses - .store(buffer_reuses, Ordering::Relaxed); - self.ssd_pread_bytes.store(pread_bytes, Ordering::Relaxed); - self.ssd_pread_ms.store(pread_ms, Ordering::Relaxed); + .store(stats.buffer_reuses, Ordering::Relaxed); + self.ssd_pread_bytes + .store(stats.pread_bytes, Ordering::Relaxed); + self.ssd_pread_ms.store(stats.pread_ms, Ordering::Relaxed); self.ssd_evict_advise_bytes - .store(evict_advise_bytes, Ordering::Relaxed); + .store(stats.evict_advise_bytes, Ordering::Relaxed); self.ssd_willneed_advise_bytes - .store(willneed_advise_bytes, Ordering::Relaxed); + .store(stats.willneed_advise_bytes, Ordering::Relaxed); self.ssd_selected_requests - .store(selected_requests, Ordering::Relaxed); + .store(stats.selected_requests, Ordering::Relaxed); self.ssd_requested_bytes - .store(requested_bytes, Ordering::Relaxed); - self.ssd_wait_ms.store(wait_ms, Ordering::Relaxed); + .store(stats.requested_bytes, Ordering::Relaxed); + self.ssd_wait_ms.store(stats.wait_ms, Ordering::Relaxed); } pub(crate) fn unloaded(&self) { @@ -535,7 +541,7 @@ impl Metrics { self.prefill_tps.store(0, Ordering::Relaxed); self.prefill_sample.store(0, Ordering::Relaxed); self.speculative_stats(0, 0, 0, 0, 0, 0); - self.ssd_stats(false, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0); + self.ssd_stats(SsdStats::default()); self.model_unloads.fetch_add(1, Ordering::Relaxed); } @@ -1042,10 +1048,27 @@ mod tests { metrics.kv_write_finished(Duration::from_millis(20), false); assert_eq!(metrics.take_kv_io_sample(), (2_048, 4_096)); assert_eq!(metrics.take_kv_io_sample(), (0, 0)); - metrics.ssd_stats( - true, 1_024, 2_048, 16, 10, 4, 90, 10, 3, 2, 8, 12, 4_096, 25, 1_024, 2_048, 12, 8_192, - 30, - ); + metrics.ssd_stats(SsdStats { + enabled: true, + resident_bytes: 1_024, + cache_bytes: 2_048, + cache_experts: 16, + cache_entries: 10, + preloaded_experts: 4, + cache_hits: 90, + cache_misses: 10, + cache_evictions: 3, + cache_wraps: 2, + buffer_allocs: 8, + buffer_reuses: 12, + pread_bytes: 4_096, + pread_ms: 25, + evict_advise_bytes: 1_024, + willneed_advise_bytes: 2_048, + selected_requests: 12, + requested_bytes: 8_192, + wait_ms: 30, + }); metrics.request_finished( WorkSource::LocalChat, Duration::from_millis(250),