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

Reach DS4 Metal performance parity

Browse Source
This commit is contained in:
Georg Bauer
2026-08-30 15:26:27 +02:00
parent 79468c65b6
commit 1ac559bbd0
5 changed files with 1958 additions and 433 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
src/engine/metal/glm.rs
View File

@@ -8,6 +8,11 @@ const CACHE_F16: bool = true;
const DECODE_FLUSH_LAYERS: usize = 4;
const AUTO_CACHE_BYTES: u64 = 12 * 1024 * 1024 * 1024;
const STREAMING_TOKEN_PREFILL_MAX: u32 = 64;
fn live_prefix_rewind_target(live: &[i32], incoming: &[i32]) -> Option<usize> {
(incoming.len() > 1 && incoming.len() < live.len() && live.starts_with(incoming))
.then_some(incoming.len() - 1)
}
const STREAMING_FULL_ATTN_CONTEXT: u32 = 8192;
const LONG_CONTEXT_THRESHOLD: u32 = 65_536;
const LONG_CONTEXT_FULL_ATTN_CONTEXT: u32 = 4096;
@@ -2365,6 +2370,14 @@ impl GlmExecutor {
}
pub(super) fn align_prompt(&mut self, tokens: &[i32]) -> Result<usize, String> {
if let Some(rewind) = live_prefix_rewind_target(&self.tokens, tokens) {
self.tokens.truncate(rewind);
if let Some(mtp) = &mut self.mtp {
mtp.pending = None;
mtp.min_pos = None;
}
return Ok(rewind);
}
if !tokens.starts_with(&self.tokens) {
self.reset()?;
}
@@ -3217,7 +3230,7 @@ fn f32_project_rows(
mod tests {
use super::{
GlmExecutor, argmax, dynamic_expert_budget, full_indexer_layer, indexed_prefill_rows,
streaming_token_prefill_eligible,
live_prefix_rewind_target, streaming_token_prefill_eligible,
};
use crate::engine::{GLM, Model, ReasoningMode};
use crate::model::ModelChoice;
@@ -3250,6 +3263,15 @@ mod tests {
assert_eq!(indexed_prefill_rows(0, 17, 2048), 17);
}
#[test]
fn repeated_glm_prompt_rewinds_one_token_for_logits() {
let live = [1, 2, 3, 4, 5, 6];
assert_eq!(live_prefix_rewind_target(&live, &[1, 2, 3, 4]), Some(3));
assert_eq!(live_prefix_rewind_target(&live, &[1]), None);
assert_eq!(live_prefix_rewind_target(&live, &live), None);
assert_eq!(live_prefix_rewind_target(&live, &[1, 2, 9]), None);
}
#[test]
fn glm_byte_budget_reserves_prefill_before_dynamic_experts() {
let mib = 1024 * 1024;

273
src/engine/metal/gpu.rs
View File

@@ -67,6 +67,12 @@ unsafe extern "C" {
expert_priorities: *const u32,
experts: u32,
) -> i32;
pub(super) fn ds4_gpu_stream_expert_cache_seed_experts_gpu_copy(
table: *const StreamExpertTable,
expert_ids: *const i32,
expert_priorities: *const u32,
experts: u32,
) -> i32;
pub(super) fn ds4_gpu_stream_expert_cache_begin_selected_load(
table: *const StreamExpertTable,
selected_ids: *const i32,
@@ -87,6 +93,26 @@ unsafe extern "C" {
count: u32,
) -> i32;
pub(super) fn ds4_gpu_flush_commands() -> i32;
pub(super) fn ds4_gpu_device_is_pre_m5_apple_silicon() -> i32;
pub(super) fn ds4_gpu_device_is_m5_apple_silicon() -> i32;
pub(super) fn ds4_gpu_set_decode_pipeline_fast_lookup(enabled: i32) -> i32;
pub(super) fn ds4_gpu_parallel_ffn_start(
gate: *mut GpuTensor,
up: *mut GpuTensor,
mid: *mut GpuTensor,
shared_out: *mut GpuTensor,
map: *const c_void,
size: u64,
gate_offset: u64,
up_offset: u64,
down_offset: u64,
model_dim: u32,
shared_dim: u32,
x: *const GpuTensor,
clamp: f32,
) -> i32;
pub(super) fn ds4_gpu_parallel_ffn_finish() -> i32;
pub(super) fn ds4_gpu_parallel_ffn_abort();
pub(super) fn ds4_gpu_tensor_alloc(bytes: u64) -> *mut GpuTensor;
pub(super) fn ds4_gpu_tensor_view(
base: *const GpuTensor,
@@ -664,6 +690,64 @@ unsafe extern "C" {
ratio: u32,
pos: u32,
) -> i32;
pub(super) fn ds4_gpu_matmul_f16_quad_compressor_store_tensor(
out0_kv: *mut GpuTensor,
out0_score: *mut GpuTensor,
out1_kv: *mut GpuTensor,
out1_score: *mut GpuTensor,
state0_kv: *mut GpuTensor,
state0_score: *mut GpuTensor,
state1_kv: *mut GpuTensor,
state1_score: *mut GpuTensor,
map: *const c_void,
size: u64,
weight0_kv: u64,
weight0_score: u64,
weight1_kv: u64,
weight1_score: u64,
ape0: u64,
ape0_type: u32,
ape1: u64,
ape1_type: u32,
input: u64,
width0: u32,
width1: u32,
x: *const GpuTensor,
ratio: u32,
pos: u32,
) -> i32;
pub(super) fn ds4_gpu_qkv_pair_quad_compressor_store_tensor(
q_rank: *mut GpuTensor,
kv_raw: *mut GpuTensor,
out0_kv: *mut GpuTensor,
out0_score: *mut GpuTensor,
out1_kv: *mut GpuTensor,
out1_score: *mut GpuTensor,
state0_kv: *mut GpuTensor,
state0_score: *mut GpuTensor,
state1_kv: *mut GpuTensor,
state1_score: *mut GpuTensor,
map: *const c_void,
size: u64,
q_a: u64,
kv: u64,
weight0_kv: u64,
weight0_score: u64,
weight1_kv: u64,
weight1_score: u64,
ape0: u64,
ape0_type: u32,
ape1: u64,
ape1_type: u32,
input: u32,
q_rank_width: u32,
kv_width: u32,
width0: u32,
width1: u32,
x: *const GpuTensor,
ratio: u32,
pos: u32,
) -> i32;
pub(super) fn ds4_gpu_hc_split_weighted_sum_norm_tensor(
out: *mut GpuTensor,
norm: *mut GpuTensor,
@@ -681,6 +765,38 @@ unsafe extern "C" {
eps: f32,
norm_eps: f32,
) -> i32;
pub(super) fn ds4_gpu_hc_rms_norm_mix_f16_available() -> i32;
pub(super) fn ds4_gpu_hc_rms_norm_mix_f16_tensor(
out: *mut GpuTensor,
x: *const GpuTensor,
map: *const c_void,
size: u64,
weight: u64,
input: u32,
output: u32,
eps: f32,
) -> i32;
pub(super) fn ds4_gpu_hc_rms_norm_mix_split_norm_f16_tensor(
mix: *mut GpuTensor,
out: *mut GpuTensor,
norm: *mut GpuTensor,
split: *mut GpuTensor,
residual: *const GpuTensor,
map: *const c_void,
size: u64,
mix_weight: u64,
scale: u64,
base: u64,
norm_weight: u64,
input: u32,
mix_width: u32,
embd: u32,
hc: u32,
iterations: u32,
eps: f32,
hc_eps: f32,
norm_eps: f32,
) -> i32;
pub(super) fn ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(
q_out: *mut GpuTensor,
q: *const GpuTensor,
@@ -695,6 +811,47 @@ unsafe extern "C" {
rows: u32,
eps: f32,
) -> i32;
pub(super) fn ds4_gpu_dsv4_qkv_rms_norm_kv_rope_fp8_store_tensor(
q_out: *mut GpuTensor,
q: *const GpuTensor,
map: *const c_void,
size: u64,
q_weight: u64,
q_width: u32,
kv_out: *mut GpuTensor,
kv: *const GpuTensor,
kv_weight: u64,
kv_width: u32,
raw_cache: *mut GpuTensor,
raw_cap: u64,
raw_row: u32,
rot: u32,
pos: u32,
original_context: u32,
freq_base: f32,
freq_scale: f32,
ext_factor: f32,
attn_factor: f32,
beta_fast: f32,
beta_slow: f32,
eps: f32,
) -> i32;
pub(super) fn ds4_gpu_kv_rope_fp8_fuse_available() -> i32;
pub(super) fn ds4_gpu_decode_attn_rope_fuse_available() -> i32;
pub(super) fn ds4_gpu_decode_attn_rope_fuse_used() -> i32;
pub(super) fn ds4_gpu_set_decode_attn_rope_fuse(
head_dim: u32,
n_rot: u32,
pos0: u32,
n_ctx_orig: u32,
inverse: bool,
freq_base: f32,
freq_scale: f32,
ext_factor: f32,
attn_factor: f32,
beta_fast: f32,
beta_slow: f32,
);
pub(super) fn ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor(
out: *mut GpuTensor,
half: *mut GpuTensor,
@@ -898,6 +1055,33 @@ unsafe extern "C" {
beta_slow: f32,
rms_eps: f32,
state_already_stored: bool,
decode_one_token: bool,
defer_finalize: bool,
) -> i32;
pub(super) fn ds4_gpu_dsv4_comp_row_finalize_tensor(
attn_stage: *mut GpuTensor,
attn_cache: *mut GpuTensor,
attn_row: u32,
attn_norm: u64,
index_cache: *mut GpuTensor,
index_row: u32,
index_norm: u64,
attn_state_kv: *mut GpuTensor,
attn_state_score: *mut GpuTensor,
index_state_kv: *mut GpuTensor,
index_state_score: *mut GpuTensor,
map: *const c_void,
size: u64,
pos: u32,
rot: u32,
original_context: u32,
freq_base: f32,
freq_scale: f32,
ext_factor: f32,
attn_factor: f32,
beta_fast: f32,
beta_slow: f32,
rms_eps: f32,
) -> i32;
pub(super) fn ds4_gpu_compressor_prefill_state_ratio4_tensor(
state_kv: *mut GpuTensor,
@@ -1258,6 +1442,35 @@ unsafe extern "C" {
x: *const GpuTensor,
clamp: f32,
) -> i32;
pub(super) fn ds4_gpu_router_shared_gate_up_q8_0_tensor(
router_logits: *mut GpuTensor,
gate: *mut GpuTensor,
up: *mut GpuTensor,
mid: *mut GpuTensor,
map: *const c_void,
size: u64,
router_weight: u64,
gate_weight: u64,
up_weight: u64,
input: u64,
experts: u64,
shared_width: u64,
x: *const GpuTensor,
clamp: f32,
router_only: bool,
) -> i32;
pub(super) fn ds4_gpu_router_project_select_fused_tensor(
router_logits: *mut GpuTensor,
probs: *mut GpuTensor,
selected: *mut GpuTensor,
weights: *mut GpuTensor,
map: *const c_void,
size: u64,
router_weight: u64,
bias: u64,
has_bias: bool,
x: *const GpuTensor,
) -> i32;
pub(super) fn ds4_gpu_shared_down_hc_expand_q8_0_tensor(
out_hc: *mut GpuTensor,
shared_out: *mut GpuTensor,
@@ -1308,6 +1521,10 @@ impl Context {
unsafe {
ds4_gpu_set_glm_model(model.shape.family == ModelFamily::Glm);
ds4_gpu_set_ssd_streaming(ssd_streaming);
// DS4 only enables this cache for the pre-M5 MXFP4 decode path.
// Rust does not accept MXFP4 weights yet, so keep the global
// native switch explicitly disabled until that path is admitted.
ds4_gpu_set_decode_pipeline_fast_lookup(0);
}
let recommended = unsafe { ds4_gpu_recommended_working_set_size() };
if admission_bytes != 0 && recommended != 0 && admission_bytes > recommended {
@@ -1409,6 +1626,62 @@ impl Commands {
}
}
pub(super) struct ParallelFfn(bool);
impl ParallelFfn {
#[allow(clippy::too_many_arguments)]
pub(super) fn start(
gate: &Buffer,
up: &Buffer,
mid: &Buffer,
shared_out: &Buffer,
map: *const c_void,
size: u64,
gate_weight: u64,
up_weight: u64,
down_weight: u64,
model_dim: u32,
shared_dim: u32,
x: &Buffer,
clamp: f32,
) -> Option<Self> {
(unsafe {
ds4_gpu_parallel_ffn_start(
gate.raw(),
up.raw(),
mid.raw(),
shared_out.raw(),
map,
size,
gate_weight,
up_weight,
down_weight,
model_dim,
shared_dim,
x.raw(),
clamp,
)
} != 0)
.then_some(Self(true))
}
pub(super) fn finish(mut self) -> Result<(), String> {
self.0 = false;
check(
unsafe { ds4_gpu_parallel_ffn_finish() },
"joining parallel Metal FFN work",
)
}
}
impl Drop for ParallelFfn {
fn drop(&mut self) {
if self.0 {
unsafe { ds4_gpu_parallel_ffn_abort() };
}
}
}
impl Drop for Commands {
fn drop(&mut self) {
if self.0 {