Skip to content

Optimize Metal prefill and decode on all Apple Silicon chips #555

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ivanfioravanti wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Optimize Metal prefill and decode on all Apple Silicon chips #555

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
de0858c
Optimize M3 Metal prefill and decode
ivanfioravanti Jul 12, 2026
0a413c3
Speed up Metal decode and prefill on M3
ivanfioravanti Jul 17, 2026
97efe60
Optimize M5 TensorOps prefill kernels, bit-identical logits
ivanfioravanti Jul 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
294 changes: 218 additions & 76 deletions ds4.c
View file
Open in desktop

Large diffs are not rendered by default.

47 changes: 42 additions & 5 deletions ds4_cuda.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8040,6 +8040,40 @@ extern "C" int ds4_gpu_matmul_f16_pair_tensor(
return cuda_ok(cudaGetLastError(), "matmul_f16_pair_ordered_chunks launch");
}

extern "C" int ds4_gpu_matmul_f16_pair_compressor_store_tensor(
ds4_gpu_tensor *out_kv,
ds4_gpu_tensor *out_score,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
const void *model_map,
uint64_t model_size,
uint64_t weight_kv_offset,
uint64_t weight_score_offset,
uint64_t ape_offset,
uint32_t ape_type,
uint64_t in_dim,
uint32_t width,
const ds4_gpu_tensor *x,
uint32_t ratio,
uint32_t pos) {
(void)out_kv;
(void)out_score;
(void)state_kv;
(void)state_score;
(void)model_map;
(void)model_size;
(void)weight_kv_offset;
(void)weight_score_offset;
(void)ape_offset;
(void)ape_type;
(void)in_dim;
(void)width;
(void)x;
(void)ratio;
(void)pos;
return 0;
}

extern "C" int ds4_gpu_matmul_f32_tensor(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint64_t in_dim, uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok) {
if (!out || !x || !model_map || in_dim == 0 || out_dim == 0 || n_tok == 0) return 0;
if (weight_offset > model_size || out_dim > UINT64_MAX / in_dim) return 0;
Expand Down Expand Up @@ -8331,7 +8365,8 @@ extern "C" int ds4_gpu_compressor_update_tensor(
float attn_factor,
float beta_fast,
float beta_slow,
float rms_eps) {
float rms_eps,
bool state_already_stored) {
if (!kv_cur || !sc_cur || !state_kv || !state_score || !comp_cache ||
!model_map || head_dim == 0 || ratio == 0 ||
n_rot > head_dim || (n_rot & 1u) != 0 ||
Expand All @@ -8355,10 +8390,12 @@ extern "C" int ds4_gpu_compressor_update_tensor(
(emit && comp_cache->bytes < comp_bytes)) {
return 0;
}
if (!ds4_gpu_compressor_store_batch_tensor(kv_cur, sc_cur, state_kv, state_score,
model_map, model_size, ape_offset, ape_type,
head_dim, ratio, pos, 1)) {
return 0;
if (!state_already_stored) {
if (!ds4_gpu_compressor_store_batch_tensor(kv_cur, sc_cur, state_kv, state_score,
model_map, model_size, ape_offset, ape_type,
head_dim, ratio, pos, 1)) {
return 0;
}
}
if (!emit) return 1;
ds4_gpu_tensor *comp_row_view = ds4_gpu_tensor_view(
Expand Down
67 changes: 66 additions & 1 deletion ds4_gpu.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,9 @@ int ds4_gpu_preload_q4_expert_tables(const void *model_map, uint64_t model_size,
int ds4_gpu_should_use_managed_kv_cache(uint64_t kv_cache_bytes, uint64_t context_bytes);
void ds4_gpu_set_quality(bool quality);
void ds4_gpu_set_ssd_streaming(bool enabled);
#ifdef __APPLE__
void ds4_gpu_release_zero_prefix_prefill_mask_cache(void);
#endif
void ds4_gpu_set_streaming_expert_cache_budget(uint32_t experts);
void ds4_gpu_set_streaming_expert_cache_expert_bytes(uint64_t bytes);
uint64_t ds4_gpu_recommended_working_set_size(void);
Expand Down Expand Up @@ -297,6 +300,26 @@ int ds4_gpu_matmul_f16_pair_tensor(
const ds4_gpu_tensor *x,
uint64_t n_tok);

/* Optional Metal decode fusion. Returns 1 when the paired projection and
* recurrent compressor-state store were encoded, 0 when the optimized path
* is unavailable, and -1 on an attempted-path error. */
int ds4_gpu_matmul_f16_pair_compressor_store_tensor(
ds4_gpu_tensor *out_kv,
ds4_gpu_tensor *out_score,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
const void *model_map,
uint64_t model_size,
uint64_t weight_kv_offset,
uint64_t weight_score_offset,
uint64_t ape_offset,
uint32_t ape_type,
uint64_t in_dim,
uint32_t width,
const ds4_gpu_tensor *x,
uint32_t ratio,
uint32_t pos);

int ds4_gpu_matmul_f32_tensor(
ds4_gpu_tensor *out,
const void *model_map,
Expand Down Expand Up @@ -496,7 +519,8 @@ int ds4_gpu_compressor_update_tensor(
float attn_factor,
float beta_fast,
float beta_slow,
float rms_eps);
float rms_eps,
bool state_already_stored);

int ds4_gpu_compressor_store_batch_tensor(
const ds4_gpu_tensor *kv,
Expand Down Expand Up @@ -595,6 +619,20 @@ int ds4_gpu_attention_decode_heads_tensor(
uint32_t n_head,
uint32_t head_dim);

/* Diagnostic/public form of the dk=512 gathered decode-attention KV staging
* step. The compressed source must be F16; dst writes chronological raw-ring
* rows followed by compressed rows and must not overlap either source. */
int ds4_gpu_flash_kv_stage_f16_tensor(
ds4_gpu_tensor *dst,
const ds4_gpu_tensor *raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_raw,
const ds4_gpu_tensor *comp,
uint32_t comp_is_f16,
uint32_t n_comp,
uint32_t head_dim);

int ds4_gpu_attention_prefill_raw_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
Expand Down Expand Up @@ -895,6 +933,18 @@ int ds4_gpu_hc_weighted_sum_tensor(
uint32_t n_embd,
uint32_t n_hc);

int ds4_gpu_hc_weighted_sum_norm_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *norm_out,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *weights,
const void *model_map,
uint64_t model_size,
uint64_t norm_weight_offset,
uint32_t n_embd,
uint32_t n_hc,
float norm_eps);

int ds4_gpu_hc_weighted_sum_split_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *residual_hc,
Expand Down Expand Up @@ -935,6 +985,21 @@ int ds4_gpu_hc_split_weighted_sum_norm_tensor(
float eps,
float norm_eps);

/* Batched HC RMSNorm followed by its narrow F16 mixer projection. On the
* tuned Metal path, scale_scratch stores one float per row instead of the
* full normalized HC tensor; other shapes retain the established fallback. */
int ds4_gpu_hc_rms_scale_project_f16_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *scale_scratch,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t in_dim,
uint32_t out_dim,
const ds4_gpu_tensor *x,
uint32_t n_rows,
float eps);

int ds4_gpu_output_hc_weights_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *pre,
Expand Down
Loading