chore: refactorings for more code reorganizing

This commit is contained in:
Georg Bauer
2026-07-29 19:28:28 +02:00
parent 95ba6c3c76
commit 92f351ae0b
6 changed files with 229 additions and 231 deletions

View File

@@ -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()
}

View File

@@ -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<Message, Task<Message>> {
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<Message> {
if let Some(id) = self.model_manager_window {
return window::gain_focus(id);

View File

@@ -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(

View File

@@ -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<CompressionFrontier, String> {
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;

View File

@@ -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()
}

View File

@@ -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),