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

chore: refactorings for more code reorganizing

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

154
src/engine/metal.rs
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;

22
src/engine/metal/gpu.rs
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()
}