Reach DS4 Metal performance parity
This commit is contained in:
Georg Bauer
120
docs/METAL_PARITY.md
Normal file
120
docs/METAL_PARITY.md
Normal file
@@ -0,0 +1,120 @@
|
||||
# Apple Metal parity for issue #78
|
||||
|
||||
## Frozen comparison
|
||||
|
||||
- DS4 oracle: `8db89fe083ae4d17c9a2428ccd29803d3ae8f577`
|
||||
- DS4Server pre-issue baseline: `79468c65b68c28e1539acf0a2a14bd42e000e365`
|
||||
- DeepSeek V4 Flash 0731 IQ2XXS/w2Q2K model SHA-256:
|
||||
`ca22ae2f838e14077c22bc1c1417b71b45b5e5a3687bd96c2ac6e17fdb6261c0`
|
||||
- DSpark support-model SHA-256:
|
||||
`7e319924541db3f7a163ed7e11d7532a70d48228ab59d36cb81e1d4511885360`
|
||||
- Machine: MacBook Pro `Mac17,7`, Apple M5 Max, 128 GiB unified memory,
|
||||
macOS 26.6.2, power 100%
|
||||
- Context and prefill allocation: 4096 tokens; quality, steering, speculative
|
||||
decoding, and DSpark disabled for the throughput gate
|
||||
- SSD mode: 4096 cached experts, 4096 popularity-preloaded experts, ordinary
|
||||
(not cold) streaming
|
||||
- Prompt: `Count from one to two hundred, spelling out every number.` rendered
|
||||
to 16 tokens; 128 greedy non-EOS tokens are measured
|
||||
|
||||
The longer prompt is deliberate: all 128 generated tokens are reachable before
|
||||
EOS and are identical between DS4 and DS4Server in each mode. A short-response
|
||||
fixture reached EOS early and made a forced post-EOS speed run compare
|
||||
unreachable, mode-dependent states.
|
||||
|
||||
## DS4 performance-change matrix
|
||||
|
||||
The audited range is
|
||||
`efdadd41e20134af4f3381e1ed90e96fe4faef6f..8db89fe083ae4d17c9a2428ccd29803d3ae8f577`.
|
||||
Documentation-only, parser/server-only, CUDA, ROCm, tensor-parallel, model
|
||||
download, and merge commits do not affect the Rust-owned single-session Metal
|
||||
graph. The applicable performance/runtime commits are accounted for below.
|
||||
|
||||
| DS4 commit(s) | DS4 call site or behavior | DS4Server result |
|
||||
| --- | --- | --- |
|
||||
| `3ce6777`, `8e0f997` | Native streamed-expert cache pinning and safe fallback | Active in the byte-matched native Metal cache; Rust preserves the cache budget and fallback. |
|
||||
| `df04887` | `ds4_gpu_attention_indexed_mixed_batch_heads_tensor` | Active through `encode_layer_with_cache_rows`; the long-context fixture crosses the indexed boundary. |
|
||||
| `74e8f11`, `4591cb1`, `a51e6ec`, `f284165`, `af80694`, `7fb2830`, `0e89a0e`, `769a8ba`, `8d683d6` | DSpark block commit, deterministic scheduling, acceptance replay/direct commit, greedy identity, and exact/opportunistic sampling | Active in `Dspark::propose` and the verifier commit paths; greedy and exact stochastic fixtures cover the frozen support model. |
|
||||
| `96c3ba4`, `d69a017`, `81f8534`, `ac5cb2c`, `48392e3`, `222b2cb`, `532ec8b` | Metal 4 Q4 projection, compact routed-MoE, aligned Q8, indexed-attention, and routed prefill kernels | Active through the existing `encode_batch_layer` Q4/indexed/routed entry points; no second prefill graph was added. |
|
||||
| `24903f3` | `ds4_gpu_head_rms_norm_rope_tail_tensor` | Existing Rust Q-head norm/RoPE fusion is retained and used. |
|
||||
| `4893e0c`, `d14ce35` | Mapped-prefill expert seeding and consistent expert/non-expert maps | Rust seeds with `ds4_gpu_stream_expert_cache_seed_experts_gpu_copy`, falls back to CPU/pread, and restores one cached static decode map with the same non-routed spans. |
|
||||
| `0dd0d36` | Release speed regression gate | Rust has an ignored release gate with retained raw CSV, DS4 reference threshold, baseline threshold, p50/p95, cache metrics, and token hashes. |
|
||||
| `4e401a1`, `0e125e6` | Automatic pre-M5/M5 admission for exact decode fusions | Rust binds both device predicates and mirrors per-generation rollback variables without a UI preference. |
|
||||
| `71d8c2a` | Q/KV pair plus ratio-4/128 compressor projection/store | `try_qkv_pair_compressor_fusion` calls the native combined entry point only for DS4's resident full-decode shapes. |
|
||||
| `b64d3e0` | Emit-path compressor finalizer | `update_compression_pair` defers both updates only when the resident finalizer is eligible, then requires exact success. |
|
||||
| `3826f52`, `ad4d05c` | Exact compressor softmax/product and ratio-4 pool | Native kernels are active; Rust's compressor ABI now passes the upstream `decode_one_token` and `defer_finalize` flags exactly. |
|
||||
| `92d83e6` | HC norm/mix producer clustering | `decode_hc_mix` uses split-norm producer fusion, with the established norm/mix/split sequence as fallback. |
|
||||
| `1be41c8` | Packed exact split-K attention reduction | Active in the vendored native attention kernel reached by Rust's decode attention calls. |
|
||||
| `0e8fb54`, `e6a6ab5` | Router/select, shared gate/up, QKV norm/KV RoPE/store, parallel full FFN, inverse-attention RoPE, and pre-M5 ports | All applicable entry points are called with DS4 shape, tensor, quality, steering, profile, SSD, and device guards. Parallel FFN aborts on partial failure; inverse RoPE falls back when the attention kernel does not consume it. |
|
||||
| `0658dac` | Long-context prefill/decode correctness | The 4096-crossing ignored fixture covers compressed/indexed state and the next decode token. |
|
||||
| `023614e`, `0ad494e`, `42033ee` | Remove stream-512 experiment; make exact paths automatic; verify DFlash pipeline | No stream-512 or opt-in experiment remains. Feature rollbacks retain identical token hashes. |
|
||||
| `24fa85e` | Keep successfully loaded disk KV checkpoints | Existing Rust checkpoint load/round-trip behavior is verified; only invalid checkpoints follow deletion policy. |
|
||||
| `7694112` | Reuse a shorter exact GLM live prefix | `live_prefix_rewind_target` rewinds to `prompt_len - 1`, clears pending speculation, and reevaluates the last prompt token. |
|
||||
| `5186e2c` | Keep Metal decode changes portable | Device checks and unavailable returns preserve non-M5 and unfused execution. |
|
||||
|
||||
The Metal MXFP4 commits (`725b084`, `1f862bd`, `7bec128`, `bcd1e57`,
|
||||
`f75a065`, and their merge/format work) are intentionally excluded: MXFP4 is
|
||||
issue #79. `ds4_gpu_set_decode_pipeline_fast_lookup` is reset to disabled,
|
||||
matching DS4 because that cache is enabled only for the pre-M5 MXFP4 path.
|
||||
|
||||
## Results
|
||||
|
||||
Three implementations were run sequentially, three times per mode, with a
|
||||
five-second cooldown. Medians are from
|
||||
[`metal-parity-issue-78.csv`](metal-parity-issue-78.csv).
|
||||
|
||||
| Implementation | Mode | Prefill tok/s | TTFT ms | First decode ms | Stable tok/s | p50 ms | p95 ms |
|
||||
| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: |
|
||||
| DS4Server baseline | resident | 69.63 | 256.07 | 26.37 | 39.53 | 24.88 | 26.61 |
|
||||
| DS4 | resident | 56.49 | 305.94 | 22.71 | 45.60 | 21.91 | 22.32 |
|
||||
| DS4Server issue #78 | resident | 69.76 | 252.01 | 22.66 | 45.60 | 21.94 | 22.33 |
|
||||
| DS4Server baseline | SSD | 6.81 | 2476.94 | 125.30 | 8.63 | 114.93 | 124.28 |
|
||||
| DS4 | SSD | 11.88 | 1393.54 | 46.75 | 26.11 | 35.62 | 46.86 |
|
||||
| DS4Server issue #78 | SSD | 11.49 | 1440.79 | 48.84 | 25.07 | 37.70 | 48.85 |
|
||||
|
||||
TTFT is the per-run prompt prefill duration plus first decode latency; its
|
||||
median is reported separately from the first decode kernel latency.
|
||||
|
||||
Resident stable decode improves 15.4% over the baseline and matches DS4.
|
||||
SSD stable decode improves 190.3% over the baseline and is 4.0% below
|
||||
DS4, inside the required 5% boundary. SSD cache accounting is identical to
|
||||
DS4 for the measured token sequence: 4096 entries, 32,958 hits, 4,194 misses,
|
||||
and 29,684,662,272 pread bytes. The resident and SSD token SHA-256 values are,
|
||||
respectively, `ed5461476a82824c3e0440e943a0b2f6b8a8f19dc26a550a1bc255003851ab8d`
|
||||
and `2bc33a04bb88f9f3c56f8bae0124b7bf25fd2f32aee81a1bdaf6e979448da052`;
|
||||
DS4 emitted the same 128-token sequences in the corresponding modes.
|
||||
|
||||
Representative `/usr/bin/time -l` measurements report:
|
||||
|
||||
| Implementation | Mode | Maximum RSS bytes | Peak footprint bytes | Process swaps |
|
||||
| --- | --- | ---: | ---: | ---: |
|
||||
| DS4 | resident | 78,790,656 | 4,746,511,440 | 0 |
|
||||
| DS4Server issue #78 | resident | 129,024,000 | 104,071,816 | 0 |
|
||||
| DS4 | SSD | 29,066,280,960 | 29,595,620,664 | 0 |
|
||||
| DS4Server issue #78 | SSD | 29,098,393,600 | 104,350,344 | 0 |
|
||||
|
||||
The post-run system memory-pressure reading was 92% free. Every measured
|
||||
process reported zero swaps; the post-run host counters were 0 swap-ins and 20
|
||||
historical swap-outs.
|
||||
|
||||
## Reproduction
|
||||
|
||||
DS4 uses the exact 16 Rust-rendered prompt token IDs with its speed-benchmark
|
||||
format. DS4Server uses the installed-model fixture:
|
||||
|
||||
```sh
|
||||
DS4_REFERENCE_TPS=45.60 \
|
||||
DS4SERVER_BASELINE_TPS=39.526849 \
|
||||
cargo test --release --all-features flash_0731_m5_decode_performance_gate \
|
||||
-- --ignored --nocapture --test-threads=1
|
||||
|
||||
DS4SERVER_BENCH_SSD=1 \
|
||||
DS4_REFERENCE_TPS=26.11 \
|
||||
DS4SERVER_BASELINE_TPS=8.634659 \
|
||||
cargo test --release --all-features flash_0731_m5_decode_performance_gate \
|
||||
-- --ignored --nocapture --test-threads=1
|
||||
```
|
||||
|
||||
The gate fails unless throughput improves over the frozen DS4Server baseline
|
||||
and remains at least 95% of DS4. `DS4SERVER_BENCH_TOKENS=1` prints the prompt
|
||||
and generated sequences for direct DS4 comparison.
|
||||
19
docs/metal-parity-issue-78.csv
Normal file
19
docs/metal-parity-issue-78.csv
Normal file
@@ -0,0 +1,19 @@
|
||||
implementation,mode,run,prompt_tokens,decode_tokens,prefill_tps,first_token_ms,steady_tps,p50_ms,p95_ms,cache_entries,cache_hits,cache_misses,pread_bytes,token_sha256
|
||||
ds4server-baseline-79468c65,resident,1,16,128,69.627264,26.270667,39.534586,24.875959,26.607167,0,0,0,0,
|
||||
ds4server-baseline-79468c65,resident,2,16,128,69.799133,26.365458,39.505867,24.905750,26.624667,0,0,0,0,
|
||||
ds4server-baseline-79468c65,resident,3,16,128,69.618099,26.454750,39.526849,24.877584,26.585875,0,0,0,0,
|
||||
ds4server-baseline-79468c65,ssd,1,16,128,6.824429,122.398208,8.638560,114.928000,124.375958,4096,34040,7208,51017416704,
|
||||
ds4server-baseline-79468c65,ssd,2,16,128,6.805883,126.033250,8.634659,114.667833,124.220583,4096,34040,7208,51017416704,
|
||||
ds4server-baseline-79468c65,ssd,3,16,128,6.726701,125.304417,8.532966,116.431375,124.278709,4096,34040,7208,51017416704,
|
||||
ds4-8db89fe0,resident,1,16,128,56.49,22.706,45.67,21.905,22.297,0,0,0,0,ed5461476a82824c3e0440e943a0b2f6b8a8f19dc26a550a1bc255003851ab8d
|
||||
ds4-8db89fe0,resident,2,16,128,56.49,22.719,45.59,21.908,22.320,0,0,0,0,ed5461476a82824c3e0440e943a0b2f6b8a8f19dc26a550a1bc255003851ab8d
|
||||
ds4-8db89fe0,resident,3,16,128,56.51,22.638,45.60,21.890,22.342,0,0,0,0,ed5461476a82824c3e0440e943a0b2f6b8a8f19dc26a550a1bc255003851ab8d
|
||||
ds4-8db89fe0,ssd,1,16,128,12.07,46.746,25.73,36.461,46.861,4096,32958,4194,29684662272,2bc33a04bb88f9f3c56f8bae0124b7bf25fd2f32aee81a1bdaf6e979448da052
|
||||
ds4-8db89fe0,ssd,2,16,128,11.88,46.737,26.11,35.622,46.982,4096,32958,4194,29684662272,2bc33a04bb88f9f3c56f8bae0124b7bf25fd2f32aee81a1bdaf6e979448da052
|
||||
ds4-8db89fe0,ssd,3,16,128,11.16,47.092,26.20,35.513,46.554,4096,32958,4194,29684662272,2bc33a04bb88f9f3c56f8bae0124b7bf25fd2f32aee81a1bdaf6e979448da052
|
||||
ds4server-issue-78,resident,1,16,128,69.930426,22.547583,45.683038,21.889833,22.275042,0,0,0,0,ed5461476a82824c3e0440e943a0b2f6b8a8f19dc26a550a1bc255003851ab8d
|
||||
ds4server-issue-78,resident,2,16,128,69.760230,22.657125,45.602401,21.936042,22.331375,0,0,0,0,ed5461476a82824c3e0440e943a0b2f6b8a8f19dc26a550a1bc255003851ab8d
|
||||
ds4server-issue-78,resident,3,16,128,69.727346,22.663208,45.361642,22.052833,22.459084,0,0,0,0,ed5461476a82824c3e0440e943a0b2f6b8a8f19dc26a550a1bc255003851ab8d
|
||||
ds4server-issue-78,ssd,1,16,128,11.279689,48.890291,24.849723,37.847292,48.853791,4096,32958,4194,29684662272,2bc33a04bb88f9f3c56f8bae0124b7bf25fd2f32aee81a1bdaf6e979448da052
|
||||
ds4server-issue-78,ssd,2,16,128,11.520460,48.567417,25.076916,37.696334,48.779625,4096,32958,4194,29684662272,2bc33a04bb88f9f3c56f8bae0124b7bf25fd2f32aee81a1bdaf6e979448da052
|
||||
ds4server-issue-78,ssd,3,16,128,11.494703,48.844875,25.066058,37.616542,48.911333,4096,32958,4194,29684662272,2bc33a04bb88f9f3c56f8bae0124b7bf25fd2f32aee81a1bdaf6e979448da052
|
||||
|
1955
src/engine/metal.rs
1955
src/engine/metal.rs
File diff suppressed because it is too large
Load Diff
@@ -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;
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user