diff --git a/ds4.c b/ds4.c index 640511eb0..19f7e8fae 100644 --- a/ds4.c +++ b/ds4.c @@ -10818,7 +10818,16 @@ static ds4_gpu_tensor *metal_graph_alloc_kv_cache_tensor(bool managed, uint64_t */ static bool metal_graph_debug_wants(const char *name, uint32_t il, uint32_t pos) { - const char *prefix = getenv("DS4_METAL_GRAPH_DUMP_PREFIX"); + /* + * The dump prefix gates every debug hook in the hot path; read it once. + * The remaining filter envs are only consulted in diagnostic runs. + */ + static const char *prefix; + static int prefix_loaded; + if (!prefix_loaded) { + prefix = getenv("DS4_METAL_GRAPH_DUMP_PREFIX"); + prefix_loaded = 1; + } if (!prefix || !prefix[0]) return false; const char *name_env = getenv("DS4_METAL_GRAPH_DUMP_NAME"); @@ -10840,9 +10849,9 @@ static void metal_graph_debug_dump_tensor( uint64_t n_f32, uint32_t il, uint32_t pos) { - const char *prefix = getenv("DS4_METAL_GRAPH_DUMP_PREFIX"); if (!t || n_f32 == 0 || !metal_graph_debug_wants(name, il, pos)) return; + const char *prefix = getenv("DS4_METAL_GRAPH_DUMP_PREFIX"); if (ds4_gpu_synchronize() == 0) { fprintf(stderr, "ds4: failed to synchronize before dumping %s layer %u pos %u\n", name, il, pos); return; @@ -13271,6 +13280,11 @@ static bool metal_graph_use_reference_qkv_norm(void) { return metal_graph_env_flag("DS4_METAL_DISABLE_QKV_NORM_FUSION", &cache); } +static bool metal_graph_use_reference_qkv_pair_proj(void) { + static int cache = -1; + return metal_graph_env_flag("DS4_METAL_DISABLE_QKV_PAIR_PROJ", &cache); +} + static bool metal_graph_use_reference_compressor_pair_proj(void) { static int cache = -1; return metal_graph_env_flag("DS4_METAL_DISABLE_COMPRESSOR_PAIR_PROJ", &cache); @@ -13283,8 +13297,15 @@ static bool metal_graph_use_reference_hc_norm_decode(void) { static bool metal_graph_enable_batch_hc_norm_fusion(void) { static int cache = -1; +#if defined(__APPLE__) + return !metal_graph_use_reference_hc_norm_decode() && + !metal_graph_env_flag("DS4_METAL_DISABLE_BATCH_HC_NORM_FUSION", &cache); +#else + /* This M3-tuned default has only been validated on Metal. Preserve the + * existing opt-in policy for CUDA/ROCm builds. */ return !metal_graph_use_reference_hc_norm_decode() && metal_graph_env_flag("DS4_METAL_ENABLE_BATCH_HC_NORM_FUSION", &cache); +#endif } static bool metal_graph_use_reference_shared_down_hc(void) { @@ -14945,7 +14966,7 @@ static bool metal_graph_encode_decode_layer( metal_graph_debug_dump_tensor("attn_norm", g->attn_norm, DS4_N_EMBD, il, pos); } bool qkv_pair_projected = false; - if (ok && qkv_rms_fused) { + if (ok && qkv_rms_fused && !metal_graph_use_reference_qkv_pair_proj()) { qkv_pair_projected = ds4_gpu_matmul_q8_0_pair_tensor(g->qr, g->kv_raw, model->map, @@ -15102,17 +15123,41 @@ static bool metal_graph_encode_decode_layer( fprintf(stderr, "ds4: Metal graph compressed KV cache capacity exceeded at layer %u\n", il); ok = false; } + bool comp_state_already_stored = false; if (ok && !metal_graph_use_reference_compressor_pair_proj()) { - ok = ds4_gpu_matmul_f16_pair_tensor(g->comp_kv_cur, - g->comp_sc_cur, - model->map, - model->size, - layer->attn_compressor_kv->abs_offset, - layer->attn_compressor_gate->abs_offset, - DS4_N_EMBD, - comp_width, - g->attn_norm, - 1) != 0; + const int fused_store = + ds4_gpu_matmul_f16_pair_compressor_store_tensor( + g->comp_kv_cur, + g->comp_sc_cur, + g->layer_attn_state_kv[il], + g->layer_attn_state_score[il], + model->map, + model->size, + layer->attn_compressor_kv->abs_offset, + layer->attn_compressor_gate->abs_offset, + layer->attn_compressor_ape->abs_offset, + layer->attn_compressor_ape->type, + DS4_N_EMBD, + comp_width, + g->attn_norm, + ratio, + pos); + if (fused_store < 0) { + ok = false; + } else if (fused_store > 0) { + comp_state_already_stored = true; + } else { + ok = ds4_gpu_matmul_f16_pair_tensor(g->comp_kv_cur, + g->comp_sc_cur, + model->map, + model->size, + layer->attn_compressor_kv->abs_offset, + layer->attn_compressor_gate->abs_offset, + DS4_N_EMBD, + comp_width, + g->attn_norm, + 1) != 0; + } } else { if (ok) ok = ds4_gpu_matmul_f16_tensor(g->comp_kv_cur, model->map, model->size, layer->attn_compressor_kv->abs_offset, @@ -15147,7 +15192,8 @@ static bool metal_graph_encode_decode_layer( attn_factor, DS4_ROPE_YARN_BETA_FAST, DS4_ROPE_YARN_BETA_SLOW, - DS4_RMS_EPS) != 0; + DS4_RMS_EPS, + comp_state_already_stored) != 0; if (ok && emit) { ds4_gpu_tensor *comp_row_view = metal_graph_attn_comp_row_view(g, il, comp_row); if (!comp_row_view) { @@ -15180,17 +15226,41 @@ static bool metal_graph_encode_decode_layer( fprintf(stderr, "ds4: Metal graph indexer compressed KV cache capacity exceeded at layer %u\n", il); ok = false; } + bool index_state_already_stored = false; if (ok && !metal_graph_use_reference_compressor_pair_proj()) { - ok = ds4_gpu_matmul_f16_pair_tensor(g->comp_kv_cur, - g->comp_sc_cur, - model->map, - model->size, - layer->indexer_compressor_kv->abs_offset, - layer->indexer_compressor_gate->abs_offset, - DS4_N_EMBD, - index_width, - g->attn_norm, - 1) != 0; + const int fused_store = + ds4_gpu_matmul_f16_pair_compressor_store_tensor( + g->comp_kv_cur, + g->comp_sc_cur, + g->layer_index_state_kv[il], + g->layer_index_state_score[il], + model->map, + model->size, + layer->indexer_compressor_kv->abs_offset, + layer->indexer_compressor_gate->abs_offset, + layer->indexer_compressor_ape->abs_offset, + layer->indexer_compressor_ape->type, + DS4_N_EMBD, + index_width, + g->attn_norm, + ratio, + pos); + if (fused_store < 0) { + ok = false; + } else if (fused_store > 0) { + index_state_already_stored = true; + } else { + ok = ds4_gpu_matmul_f16_pair_tensor(g->comp_kv_cur, + g->comp_sc_cur, + model->map, + model->size, + layer->indexer_compressor_kv->abs_offset, + layer->indexer_compressor_gate->abs_offset, + DS4_N_EMBD, + index_width, + g->attn_norm, + 1) != 0; + } } else { if (ok) ok = ds4_gpu_matmul_f16_tensor(g->comp_kv_cur, model->map, model->size, layer->indexer_compressor_kv->abs_offset, @@ -15225,7 +15295,8 @@ static bool metal_graph_encode_decode_layer( attn_factor, DS4_ROPE_YARN_BETA_FAST, DS4_ROPE_YARN_BETA_SLOW, - DS4_RMS_EPS) != 0; + DS4_RMS_EPS, + index_state_already_stored) != 0; if (ok && emit) { ds4_gpu_tensor *index_row_view = ds4_gpu_tensor_view( g->layer_index_comp_cache[il], @@ -16097,21 +16168,46 @@ static bool metal_graph_encode_output_head( if (ok) { metal_graph_debug_dump_tensor("result_hc_weights", g->output_weights, DS4_N_HC, DS4_N_LAYER, 0); } - if (ok) ok = ds4_gpu_hc_weighted_sum_tensor(g->output_embd, - g->cur_hc, - g->output_weights, - DS4_N_EMBD, - DS4_N_HC) != 0; + bool output_sum_norm_fused = false; +#if defined(__APPLE__) + if (ok) { + output_sum_norm_fused = + ds4_gpu_hc_weighted_sum_norm_tensor( + g->output_embd, + g->output_norm, + g->cur_hc, + g->output_weights, + model->map, + model->size, + weights->output_norm->abs_offset, + DS4_N_EMBD, + DS4_N_HC, + DS4_RMS_EPS) != 0; + if (!output_sum_norm_fused && + getenv("DS4_METAL_REQUIRE_OUTPUT_HC_SUM_NORM_FUSION") != NULL) { + ok = false; + } + } +#endif + if (ok && !output_sum_norm_fused) { + ok = ds4_gpu_hc_weighted_sum_tensor(g->output_embd, + g->cur_hc, + g->output_weights, + DS4_N_EMBD, + DS4_N_HC) != 0; + } if (ok) { metal_graph_debug_dump_tensor("result_hc", g->output_embd, DS4_N_EMBD, DS4_N_LAYER, 0); } - if (ok) ok = ds4_gpu_rms_norm_weight_tensor(g->output_norm, - g->output_embd, - model->map, - model->size, - weights->output_norm->abs_offset, - DS4_N_EMBD, - DS4_RMS_EPS) != 0; + if (ok && !output_sum_norm_fused) { + ok = ds4_gpu_rms_norm_weight_tensor(g->output_norm, + g->output_embd, + model->map, + model->size, + weights->output_norm->abs_offset, + DS4_N_EMBD, + DS4_RMS_EPS) != 0; + } if (ok) { metal_graph_debug_dump_tensor("result_norm", g->output_norm, DS4_N_EMBD, DS4_N_LAYER, 0); } @@ -17176,6 +17272,52 @@ static bool metal_graph_upload_prompt_embeddings_hc( n_tokens); } +static bool metal_graph_hc_rms_scale_project( + ds4_gpu_tensor *out, + ds4_gpu_tensor *norm_scratch, + const ds4_model *model, + const ds4_tensor *weight, + const ds4_gpu_tensor *x, + uint64_t in_dim, + uint32_t n_tokens) { + if (!out || !norm_scratch || !model || !weight || !x || + in_dim > UINT32_MAX) { + return false; + } +#if defined(__APPLE__) + return ds4_gpu_hc_rms_scale_project_f16_tensor( + out, + norm_scratch, + model->map, + model->size, + weight->abs_offset, + (uint32_t)in_dim, + 2u * DS4_N_HC + DS4_N_HC * DS4_N_HC, + x, + n_tokens, + DS4_RMS_EPS) != 0; +#else + bool ok = ds4_gpu_rms_norm_plain_rows_tensor( + norm_scratch, + x, + (uint32_t)in_dim, + n_tokens, + DS4_RMS_EPS) != 0; + if (ok) { + ok = ds4_gpu_matmul_f16_tensor( + out, + model->map, + model->size, + weight->abs_offset, + in_dim, + 2u * DS4_N_HC + DS4_N_HC * DS4_N_HC, + norm_scratch, + n_tokens) != 0; + } + return ok; +#endif +} + static bool metal_graph_warmup_prefill_kernels( ds4_gpu_graph *g, const ds4_model *model, @@ -17196,18 +17338,17 @@ static bool metal_graph_warmup_prefill_kernels( if (n_tokens <= 8) return true; const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; bool ok = ds4_gpu_begin_commands() != 0; if (ok) { - ok = ds4_gpu_matmul_f16_tensor(g->batch_hc_mix, - model->map, - model->size, - weights->layer[0].hc_attn_fn->abs_offset, - hc_dim, - mix_hc, - g->batch_flat_hc, - n_tokens) != 0; + ok = metal_graph_hc_rms_scale_project( + g->batch_hc_mix, + g->batch_flat_hc, + model, + weights->layer[0].hc_attn_fn, + g->batch_cur_hc, + hc_dim, + n_tokens); } if (ok) ok = ds4_gpu_end_commands() != 0; if (!ok) { @@ -17364,22 +17505,17 @@ static bool metal_graph_encode_layer_attention_batch( ds4_gpu_tensor *after_attn_hc_view = ds4_gpu_tensor_view( g->batch_after_attn_hc, 0, (uint64_t)n_tokens * hc_dim * sizeof(float)); bool ok = hc_mix_view && hc_split_view && attn_cur_view && after_attn_hc_view; - const bool fuse_hc_norm = DS4_N_HC == 4 && + const bool fuse_hc_norm = n_tokens > 1 && + DS4_N_HC == 4 && !metal_graph_use_reference_hc_decode() && metal_graph_enable_batch_hc_norm_fusion(); - if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(g->batch_flat_hc, - g->batch_cur_hc, - (uint32_t)hc_dim, - n_tokens, - DS4_RMS_EPS) != 0; - if (ok) ok = ds4_gpu_matmul_f16_tensor(hc_mix_view, - model->map, - model->size, - layer->hc_attn_fn->abs_offset, - hc_dim, - mix_hc, - g->batch_flat_hc, - n_tokens) != 0; + if (ok) ok = metal_graph_hc_rms_scale_project(hc_mix_view, + g->batch_flat_hc, + model, + layer->hc_attn_fn, + g->batch_cur_hc, + hc_dim, + n_tokens); if (metal_graph_use_reference_hc_decode()) { if (ok) ok = ds4_gpu_hc_split_sinkhorn_tensor(hc_split_view, hc_mix_view, @@ -17982,7 +18118,8 @@ static bool metal_graph_encode_layer_attention_batch( attn_factor, DS4_ROPE_YARN_BETA_FAST, DS4_ROPE_YARN_BETA_SLOW, - DS4_RMS_EPS) != 0; + DS4_RMS_EPS, + false) != 0; if (ok && emit) { ds4_gpu_tensor *comp_row_view = metal_graph_attn_comp_row_view(g, il, comp_row); ok = comp_row_view && @@ -18274,7 +18411,8 @@ static bool metal_graph_encode_layer_attention_batch( attn_factor, DS4_ROPE_YARN_BETA_FAST, DS4_ROPE_YARN_BETA_SLOW, - DS4_RMS_EPS) != 0; + DS4_RMS_EPS, + false) != 0; if (ok && emit) { ds4_gpu_tensor *index_row_view = ds4_gpu_tensor_view( g->layer_index_comp_cache[il], @@ -18814,22 +18952,17 @@ static bool metal_graph_encode_layer_ffn_batch( ds4_gpu_tensor *next_hc_view = ds4_gpu_tensor_view( g->batch_next_hc, 0, (uint64_t)n_tokens * hc_dim * sizeof(float)); bool ok = hc_mix_view && hc_split_view && ffn_cur_view && next_hc_view; - const bool fuse_hc_norm = DS4_N_HC == 4 && + const bool fuse_hc_norm = n_tokens > 1 && + DS4_N_HC == 4 && !metal_graph_use_reference_hc_decode() && metal_graph_enable_batch_hc_norm_fusion(); - if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(g->batch_flat_hc, - g->batch_after_attn_hc, - (uint32_t)hc_dim, - n_tokens, - DS4_RMS_EPS) != 0; - if (ok) ok = ds4_gpu_matmul_f16_tensor(hc_mix_view, - model->map, - model->size, - layer->hc_ffn_fn->abs_offset, - hc_dim, - mix_hc, - g->batch_flat_hc, - n_tokens) != 0; + if (ok) ok = metal_graph_hc_rms_scale_project(hc_mix_view, + g->batch_flat_hc, + model, + layer->hc_ffn_fn, + g->batch_after_attn_hc, + hc_dim, + n_tokens); if (metal_graph_use_reference_hc_decode()) { if (ok) ok = ds4_gpu_hc_split_sinkhorn_tensor(hc_split_view, hc_mix_view, @@ -20419,6 +20552,9 @@ static bool metal_graph_prefill_layer_major( } return false; } +#ifdef __APPLE__ + ds4_gpu_release_zero_prefix_prefill_mask_cache(); +#endif const double t_before_read = profile ? now_sec() : 0.0; if (logits) { @@ -20805,6 +20941,12 @@ static bool metal_graph_prefill_layer_major( } return false; } +#ifdef __APPLE__ + /* Zero-prefix masks are shared across the 43 per-layer command batches, + * then become dead weight. Release them before the output head and later + * replay/decode chunks so the prefill win does not add residency pressure. */ + ds4_gpu_release_zero_prefix_prefill_mask_cache(); +#endif if (show_progress) fputc('\n', stderr); metal_graph_stream_prefill_selected_profile_summary(g); #ifdef DS4_ROCM_BUILD diff --git a/ds4_cuda.cu b/ds4_cuda.cu index 188b341ad..c7bcfae43 100644 --- a/ds4_cuda.cu +++ b/ds4_cuda.cu @@ -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; @@ -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 || @@ -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( diff --git a/ds4_gpu.h b/ds4_gpu.h index b58aca9bd..14e88b104 100644 --- a/ds4_gpu.h +++ b/ds4_gpu.h @@ -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); @@ -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, @@ -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, @@ -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, @@ -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, @@ -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, diff --git a/ds4_metal.m b/ds4_metal.m index 7e3f8bd5c..81e3aee00 100644 --- a/ds4_metal.m +++ b/ds4_metal.m @@ -57,18 +57,25 @@ static id g_concat_pipeline; static id g_cpy_f32_f32_pipeline; static id g_cpy_f32_f16_pipeline; +static id g_cpy_contig_f32_f16_pipeline; static id g_cpy_f16_f32_pipeline; +static id g_cpy_contig_f16_f32_pipeline; +static id g_cpy_contig_f16_f16_pipeline; +static id g_flash_kv_stage_f16_pipeline; static id g_swiglu_pipeline; static id g_add_pipeline; static id g_moe_sum6_pipeline; static id g_mul_pipeline; static id g_rms_norm_pipeline; static id g_rms_norm_plain_pipeline; +static id g_rms_norm_scale_pipeline; static id g_dsv4_qkv_rms_norm_pipeline; static id g_hc_split_sinkhorn_pipeline; static id g_hc_split_weighted_sum_pipeline; static id g_hc_split_weighted_sum_norm_pipeline; static id g_hc_weighted_sum_pipeline; +static id g_hc_weighted_sum_norm_pipeline; +static id g_output_hc_weights4_pipeline; static id g_hc_expand_pipeline; static id g_unary_sigmoid_pipeline; static id g_unary_silu_pipeline; @@ -115,10 +122,15 @@ static id g_moe_table_q4_pair_up_encoder; static id g_moe_table_q4_sum_down_encoder; static id g_rope_tail_batch_pipeline; +static id g_rope_tail_inplace_pair_pipeline; +static id g_rope_tail_inplace_pair_shared4_pipeline; +static id g_rope_tail_inplace_pair_affine_pipeline; static id g_dsv4_fp8_kv_quantize_pipeline; static id g_dsv4_indexer_qat_pipeline; static id g_dsv4_kv_fp8_store_pipeline; static id g_dsv4_ratio4_shift_pipeline; +static id g_dsv4_compressor_pack_ratio4_pipeline; +static id g_dsv4_softmax_pool_ratio4_direct_pipeline; static id g_dsv4_softmax_pool_pipeline; static id g_soft_max_f32_pipeline; static id g_soft_max_f32_4_pipeline; @@ -135,7 +147,10 @@ static id g_dsv4_indexed_attention_heads8_rb16_pipeline; static id g_dsv4_softplus_sqrt_pipeline; static id g_dsv4_router_finalize_one_pipeline; +static id g_dsv4_router_finalize_one_simd_pipeline; +static id g_dsv4_router_finalize_weights_one_simd_pipeline; static id g_dsv4_router_weights_one_pipeline; +static id g_dsv4_router_weights_batch_pipeline; static id g_dsv4_hc_expand4_pipeline; static NSMutableDictionary> *g_pipeline_cache; static NSMutableDictionary> *g_model_buffer_cache; @@ -143,7 +158,38 @@ static NSMutableDictionary *g_q4_expert_layer_residency_cache; static NSMutableArray> *g_transient_buffers; static id g_model_residency_set; + +typedef struct { + id __strong mask; + id __strong blk; + NSUInteger mask_bytes; + NSUInteger blk_bytes; + uint32_t kind; + uint32_t n_tokens; + uint32_t n_comp; + uint32_t n_keys; + uint32_t window; + uint32_t ratio; + uint32_t nqptg; + uint32_t ncpsg; + bool has_kvpad; + bool bc_mask; + bool valid; + bool blk_ready; +} ds4_gpu_zero_prefix_prefill_mask_cache_entry; + +enum { + DS4_GPU_PREFILL_MASK_CACHE_RAW = 1, + DS4_GPU_PREFILL_MASK_CACHE_RATIO4 = 2, + DS4_GPU_PREFILL_MASK_CACHE_RATIO128 = 3, + DS4_GPU_PREFILL_MASK_CACHE_SLOTS = 3, +}; + +static ds4_gpu_zero_prefix_prefill_mask_cache_entry + g_zero_prefix_prefill_mask_cache[DS4_GPU_PREFILL_MASK_CACHE_SLOTS]; +static void ds4_gpu_invalidate_zero_prefix_prefill_block_maps(void); static id g_flash_attn_mask_buffer; +static id g_flash_attn_zero_mask_buffer; static id g_flash_attn_pad_buffer; static id g_flash_attn_tmp_buffer; static id g_flash_attn_blk_buffer; @@ -290,6 +336,7 @@ static int ds4_gpu_stream_expert_cache_entry_protected( const int32_t *protect_ids, uint32_t n_protect); static NSUInteger g_flash_attn_mask_bytes; +static NSUInteger g_flash_attn_zero_mask_bytes; static NSUInteger g_flash_attn_pad_bytes; static NSUInteger g_flash_attn_tmp_bytes; static NSUInteger g_flash_attn_blk_bytes; @@ -920,6 +967,7 @@ static int ds4_gpu_wait_pending_command_buffers(const char *label) { } [g_pending_cbs removeAllObjects]; ds4_gpu_stream_expert_cache_note_pending_completed(); + if (!ok) ds4_gpu_invalidate_zero_prefix_prefill_block_maps(); return ok; } @@ -928,7 +976,10 @@ static int ds4_gpu_finish_command_buffer(id cb, int owned, con [cb commit]; int ok = ds4_gpu_wait_pending_command_buffers(label); - if (!ds4_gpu_wait_command_buffer(cb, label)) ok = 0; + if (!ds4_gpu_wait_command_buffer(cb, label)) { + ok = 0; + ds4_gpu_invalidate_zero_prefix_prefill_block_maps(); + } ds4_gpu_stream_expert_cache_note_owned_completed(); [g_transient_buffers removeAllObjects]; ds4_gpu_model_buffer_cache_maybe_evict(label); @@ -999,6 +1050,25 @@ static int ds4_gpu_ensure_scratch_buffer( return 1; } +static int ds4_gpu_ensure_zero_attention_mask(NSUInteger bytes) { + const NSUInteger capacity = 8192u * sizeof(uint16_t); + if (bytes > capacity) return 0; + if (g_flash_attn_zero_mask_buffer && + g_flash_attn_zero_mask_bytes >= capacity) { + return 1; + } + if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_zero_mask_buffer, + &g_flash_attn_zero_mask_bytes, + capacity, + "ds4_flash_attn_zero_mask")) { + return 0; + } + void *contents = [g_flash_attn_zero_mask_buffer contents]; + if (!contents) return 0; + memset(contents, 0, g_flash_attn_zero_mask_bytes); + return 1; +} + static uint64_t round_up_u64(uint64_t v, uint64_t align) { return (v + align - 1) & ~(align - 1); } @@ -1687,6 +1757,136 @@ static int ds4_gpu_map_model_views( return buffer; } +static int ds4_gpu_zero_prefix_prefill_mask_cache_enabled(void) { + if (getenv("DS4_METAL_DISABLE_M3_ZERO_PREFIX_PREFILL_MASK_CACHE") != NULL || + getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL) { + return 0; + } + return ds4_gpu_device_name_contains("M3") || + getenv("DS4_METAL_ENABLE_ZERO_PREFIX_PREFILL_MASK_CACHE") != NULL; +} + +static ds4_gpu_zero_prefix_prefill_mask_cache_entry * +ds4_gpu_get_zero_prefix_prefill_mask_cache( + uint32_t kind, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t n_keys, + uint32_t window, + uint32_t ratio, + uint32_t nqptg, + uint32_t ncpsg, + bool has_kvpad, + bool bc_mask, + NSUInteger mask_bytes, + NSUInteger blk_bytes, + bool *created) { + if (created) *created = false; + if (!created || !ds4_gpu_zero_prefix_prefill_mask_cache_enabled() || + mask_bytes == 0 || blk_bytes == 0) { + return NULL; + } + + uint32_t slot = UINT32_MAX; + switch (kind) { + case DS4_GPU_PREFILL_MASK_CACHE_RAW: slot = 0; break; + case DS4_GPU_PREFILL_MASK_CACHE_RATIO4: slot = 1; break; + case DS4_GPU_PREFILL_MASK_CACHE_RATIO128: slot = 2; break; + default: return NULL; + } + + ds4_gpu_zero_prefix_prefill_mask_cache_entry *entry = + &g_zero_prefix_prefill_mask_cache[slot]; + if (entry->valid && entry->mask && entry->blk && + entry->mask_bytes == mask_bytes && entry->blk_bytes == blk_bytes && + entry->kind == kind && entry->n_tokens == n_tokens && + entry->n_comp == n_comp && entry->n_keys == n_keys && + entry->window == window && entry->ratio == ratio && + entry->nqptg == nqptg && entry->ncpsg == ncpsg && + entry->has_kvpad == has_kvpad && entry->bc_mask == bc_mask) { + return entry; + } + + id mask = [g_device newBufferWithLength:mask_bytes + options:MTLResourceStorageModeShared]; + id blk = [g_device newBufferWithLength:blk_bytes + options:MTLResourceStorageModePrivate]; + if (!blk) { + blk = [g_device newBufferWithLength:blk_bytes + options:MTLResourceStorageModeShared]; + } + if (!mask || !blk) return NULL; + + mask.label = [NSString stringWithFormat:@"ds4_prefill_mask_cache_%u", kind]; + blk.label = [NSString stringWithFormat:@"ds4_prefill_blk_cache_%u", kind]; + + /* A replaced entry may still be referenced by an uncommitted or in-flight + * command buffer. Keep its resources alive with the other batch-scoped + * buffers until command completion instead of mutating or releasing them. */ + if (entry->mask) [g_transient_buffers addObject:entry->mask]; + if (entry->blk) [g_transient_buffers addObject:entry->blk]; + + entry->mask = mask; + entry->blk = blk; + entry->mask_bytes = mask_bytes; + entry->blk_bytes = blk_bytes; + entry->kind = kind; + entry->n_tokens = n_tokens; + entry->n_comp = n_comp; + entry->n_keys = n_keys; + entry->window = window; + entry->ratio = ratio; + entry->nqptg = nqptg; + entry->ncpsg = ncpsg; + entry->has_kvpad = has_kvpad; + entry->bc_mask = bc_mask; + /* The caller publishes the entry only after synchronously filling mask. */ + entry->valid = false; + entry->blk_ready = false; + *created = true; + return entry; +} + +static void ds4_gpu_invalidate_zero_prefix_prefill_block_maps(void) { + for (uint32_t i = 0; i < DS4_GPU_PREFILL_MASK_CACHE_SLOTS; i++) { + g_zero_prefix_prefill_mask_cache[i].blk_ready = false; + } +} + +static void ds4_gpu_clear_zero_prefix_prefill_mask_cache(void) { + for (uint32_t i = 0; i < DS4_GPU_PREFILL_MASK_CACHE_SLOTS; i++) { + ds4_gpu_zero_prefix_prefill_mask_cache_entry *entry = + &g_zero_prefix_prefill_mask_cache[i]; + entry->mask = nil; + entry->blk = nil; + entry->mask_bytes = 0; + entry->blk_bytes = 0; + entry->kind = 0; + entry->n_tokens = 0; + entry->n_comp = 0; + entry->n_keys = 0; + entry->window = 0; + entry->ratio = 0; + entry->nqptg = 0; + entry->ncpsg = 0; + entry->has_kvpad = false; + entry->bc_mask = false; + entry->valid = false; + entry->blk_ready = false; + } +} + +void ds4_gpu_release_zero_prefix_prefill_mask_cache(void) { + /* Layer-major prefill waits each layer before advancing, so its final + * release point has no outstanding cache users. Keep the guard here for + * diagnostic callers that may have an open or asynchronously flushed CB. */ + if (!g_initialized || g_batch_cb || + (g_pending_cbs && [g_pending_cbs count] != 0)) { + return; + } + ds4_gpu_clear_zero_prefix_prefill_mask_cache(); +} + static id ds4_gpu_get_mul_mm_pipeline( const char *function_name, bool bc_inp, @@ -2291,23 +2491,48 @@ static int ds4_gpu_warm_model_views(void) { bool has_bias, bool has_scap, bool has_kvpad, + bool shared_kvpad, int32_t ns10, int32_t ns20, int32_t nsg, int32_t nwg) { - NSString *key = [NSString stringWithFormat:@"%s_mask=%d_sinks=%d_bias=%d_scap=%d_kvpad=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d", + /* + * Decode calls this once per layer with identical arguments, so memoize + * the last hit and skip the NSString key + dictionary lookup on the hot + * path. The generic cache below remains the fallback for new variants. + */ + static struct { + const char *fn; + bool m, s, b, c, k, sp; + int32_t n10, n20, sg, wg; + id pipeline; + } memo; + if (memo.pipeline && memo.fn != NULL && strcmp(memo.fn, function_name) == 0 && + memo.m == has_mask && memo.s == has_sinks && memo.b == has_bias && + memo.c == has_scap && memo.k == has_kvpad && memo.sp == shared_kvpad && + memo.n10 == ns10 && memo.n20 == ns20 && memo.sg == nsg && memo.wg == nwg) { + return memo.pipeline; + } + + NSString *key = [NSString stringWithFormat:@"%s_mask=%d_sinks=%d_bias=%d_scap=%d_kvpad=%d_sharedpad=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d", function_name, has_mask ? 1 : 0, has_sinks ? 1 : 0, has_bias ? 1 : 0, has_scap ? 1 : 0, has_kvpad ? 1 : 0, + shared_kvpad ? 1 : 0, (int)ns10, (int)ns20, (int)nsg, (int)nwg]; id cached = [g_pipeline_cache objectForKey:key]; - if (cached) return cached; + if (cached) { + memo = (typeof(memo)){ function_name, has_mask, has_sinks, has_bias, + has_scap, has_kvpad, shared_kvpad, ns10, ns20, + nsg, nwg, cached }; + return cached; + } MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; [constants setConstantValue:&has_mask type:MTLDataTypeBool atIndex:400]; @@ -2315,6 +2540,7 @@ static int ds4_gpu_warm_model_views(void) { [constants setConstantValue:&has_bias type:MTLDataTypeBool atIndex:402]; [constants setConstantValue:&has_scap type:MTLDataTypeBool atIndex:403]; [constants setConstantValue:&has_kvpad type:MTLDataTypeBool atIndex:404]; + [constants setConstantValue:&shared_kvpad type:MTLDataTypeBool atIndex:405]; [constants setConstantValue:&ns10 type:MTLDataTypeInt atIndex:420]; [constants setConstantValue:&ns20 type:MTLDataTypeInt atIndex:421]; [constants setConstantValue:&nsg type:MTLDataTypeInt atIndex:422]; @@ -2340,16 +2566,29 @@ static int ds4_gpu_warm_model_views(void) { } [g_pipeline_cache setObject:pipeline forKey:key]; + memo = (typeof(memo)){ function_name, has_mask, has_sinks, has_bias, + has_scap, has_kvpad, shared_kvpad, ns10, ns20, + nsg, nwg, pipeline }; return pipeline; } static id ds4_gpu_get_flash_attn_reduce_pipeline( int32_t dv, int32_t nwg) { + /* Same per-layer memo pattern as the vec getter above. */ + static int32_t memo_dv, memo_nwg; + static id memo_pipeline; + if (memo_pipeline && memo_dv == dv && memo_nwg == nwg) { + return memo_pipeline; + } + NSString *key = [NSString stringWithFormat:@"kernel_flash_attn_ext_vec_reduce_dv=%d_nwg=%d", (int)dv, (int)nwg]; id cached = [g_pipeline_cache objectForKey:key]; - if (cached) return cached; + if (cached) { + memo_dv = dv; memo_nwg = nwg; memo_pipeline = cached; + return cached; + } MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; [constants setConstantValue:&dv type:MTLDataTypeInt atIndex:500]; @@ -2374,6 +2613,7 @@ static int ds4_gpu_warm_model_views(void) { } [g_pipeline_cache setObject:pipeline forKey:key]; + memo_dv = dv; memo_nwg = nwg; memo_pipeline = pipeline; return pipeline; } @@ -2654,11 +2894,22 @@ static void ds4_gpu_print_task_memory_report(void) { } void ds4_gpu_print_memory_report(const char *label) { + uint64_t cached_prefill_mask_bytes = 0; + uint64_t cached_prefill_blk_bytes = 0; + for (uint32_t i = 0; i < DS4_GPU_PREFILL_MASK_CACHE_SLOTS; i++) { + const ds4_gpu_zero_prefix_prefill_mask_cache_entry *entry = + &g_zero_prefix_prefill_mask_cache[i]; + if (entry->mask) cached_prefill_mask_bytes += entry->mask_bytes; + if (entry->blk) cached_prefill_blk_bytes += entry->blk_bytes; + } const uint64_t scratch = (uint64_t)g_flash_attn_mask_bytes + + (uint64_t)g_flash_attn_zero_mask_bytes + + cached_prefill_mask_bytes + (uint64_t)g_flash_attn_pad_bytes + (uint64_t)g_flash_attn_tmp_bytes + (uint64_t)g_flash_attn_blk_bytes + + cached_prefill_blk_bytes + (uint64_t)g_flash_attn_ring_bytes + (uint64_t)g_flash_attn_kv_bytes + (uint64_t)g_compressor_pool_kv_bytes + @@ -2906,10 +3157,13 @@ void ds4_gpu_print_memory_report(const char *label) { fprintf(stderr, "ds4: scratch %.2f MiB (flash mask %.2f, pad %.2f, tmp %.2f, blk %.2f, ring %.2f, kv %.2f, compressor %.2f, router %.2f, indexer %.2f, moe %.2f, f16 %.2f, raw-store %.2f)\n", ds4_gpu_mib(scratch), - ds4_gpu_mib((uint64_t)g_flash_attn_mask_bytes), + ds4_gpu_mib((uint64_t)g_flash_attn_mask_bytes + + (uint64_t)g_flash_attn_zero_mask_bytes + + cached_prefill_mask_bytes), ds4_gpu_mib((uint64_t)g_flash_attn_pad_bytes), ds4_gpu_mib((uint64_t)g_flash_attn_tmp_bytes), - ds4_gpu_mib((uint64_t)g_flash_attn_blk_bytes), + ds4_gpu_mib((uint64_t)g_flash_attn_blk_bytes + + cached_prefill_blk_bytes), ds4_gpu_mib((uint64_t)g_flash_attn_ring_bytes), ds4_gpu_mib((uint64_t)g_flash_attn_kv_bytes), ds4_gpu_mib((uint64_t)g_compressor_pool_kv_bytes + @@ -3894,6 +4148,26 @@ static NSUInteger ds4_gpu_rms_norm_pipeline_threads( uint64_t nb1; } ds4_gpu_hc_weighted_sum_args; +typedef struct { + int64_t n_embd; + int64_t n_hc; + int64_t n_tokens; + uint64_t nb_x0; + uint64_t nb_x1; + uint64_t nb_x2; + uint64_t nb_w0; + uint64_t nb_w1; + uint64_t nb0; + uint64_t nb1; + uint64_t nb_norm1; + float norm_eps; +} ds4_gpu_hc_weighted_sum_norm_args; + +typedef struct { + float post_scale; + float eps; +} ds4_gpu_output_hc_weights4_args; + typedef struct { int64_t n_embd; int32_t n_hc; @@ -4142,6 +4416,15 @@ static int ds4_gpu_encode_mul_mm_id_mapped_tile( uint64_t nb33; } ds4_gpu_flash_attn_pad_args; +typedef struct { + uint32_t raw_cap; + uint32_t raw_start; + uint32_t n_raw; + uint32_t n_comp; + uint32_t pad_rows; + uint32_t shared_pad; +} ds4_gpu_flash_kv_stage_f16_args; + typedef struct { int32_t ne01; int32_t ne30; @@ -4218,6 +4501,26 @@ static int ds4_gpu_encode_mul_mm_id_mapped_tile( bool src2; } ds4_gpu_rope_tail_batch_args; +typedef struct { + uint64_t row_bytes; + uint64_t token_bytes; + int32_t head_dim; + int32_t n_dims; + int32_t n_ctx_orig; + int32_t inverse; + uint32_t pos0; + uint32_t pos_step; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; +} ds4_gpu_rope_affine_pair_args; + +_Static_assert(sizeof(ds4_gpu_rope_affine_pair_args) == 64, + "Metal affine RoPE argument ABI changed"); + static ds4_gpu_rope_tail_batch_args ds4_gpu_make_rope_tail_args( uint32_t n_tok, uint32_t n_head, @@ -4270,51 +4573,133 @@ static int ds4_gpu_encode_rope_tail_inplace( uint32_t head_dim, uint32_t pos0, uint32_t pos_step) { - int32_t pos_stack[256]; - int32_t *pos = pos_stack; - if (n_tok > (uint32_t)(sizeof(pos_stack) / sizeof(pos_stack[0]))) { - pos = malloc((size_t)n_tok * sizeof(*pos)); - if (!pos) { - fprintf(stderr, "ds4: failed to allocate Metal RoPE position buffer\n"); - return 0; - } + const uint32_t tail_threads = args->n_dims > 0 ? (uint32_t)args->n_dims : 0u; + const bool lane_compatible = + tail_threads <= head_dim && ((head_dim - tail_threads) & 31u) == 0u; + const bool force_affine_position = + getenv("DS4_METAL_ENABLE_AFFINE_ROPE_PAIR") != NULL; + const bool use_inplace_pair = + g_rope_tail_inplace_pair_pipeline != nil && + args->mode == 0 && !args->src2 && + lane_compatible && + (ds4_gpu_device_name_contains("M3") || + getenv("DS4_METAL_ENABLE_INPLACE_ROPE_PAIR") != NULL || + force_affine_position) && + getenv("DS4_METAL_DISABLE_M3_INPLACE_ROPE_PAIR") == NULL; + const bool use_shared_coeff = + use_inplace_pair && !force_affine_position && + g_rope_tail_inplace_pair_shared4_pipeline != nil && + /* The 256-thread grouped schedule helps long prefill, but reduces the + * per-head parallelism that short batches and decode rely on. */ + tail_threads == 64u && n_head >= 4u && n_tok >= 32u && + getenv("DS4_METAL_DISABLE_M3_SHARED_ROPE_COEFF") == NULL; + if (force_affine_position && + g_rope_tail_inplace_pair_affine_pipeline == nil && + getenv("DS4_METAL_DISABLE_M3_INPLACE_ROPE_PAIR") == NULL && + getenv("DS4_METAL_DISABLE_M3_AFFINE_ROPE_PAIR") == NULL) { + fprintf(stderr, + "ds4: forced affine-position RoPE pipeline is unavailable\n"); + return 0; } - for (uint32_t t = 0; t < n_tok; t++) pos[t] = (int32_t)(pos0 + t * pos_step); + /* Keep long prefill on the proven shared4 kernel. Reconstructing affine + * positions inside its coefficient cohort perturbs YaRN fast-math codegen; + * the compact affine specialization is exact for the decode pair schedule. */ + const bool use_affine_position = + use_inplace_pair && !use_shared_coeff && + g_rope_tail_inplace_pair_affine_pipeline != nil && + (n_tok == 1u || force_affine_position) && + (ds4_gpu_device_name_contains("M3") || force_affine_position) && + getenv("DS4_METAL_DISABLE_M3_AFFINE_ROPE_PAIR") == NULL; - const NSUInteger pos_bytes = (NSUInteger)n_tok * sizeof(*pos); + int32_t pos_stack[256]; + int32_t *pos = NULL; id posbuf = nil; - if (pos_bytes > 4096u) { - /* - * Metal inline setBytes data is meant for small constants. Long prefill - * RoPE calls need thousands of positions; passing that much inline can - * make the Apple driver abort the process instead of reporting a normal - * API error. - */ - posbuf = ds4_gpu_new_transient_buffer(pos_bytes, "ds4_rope_positions"); - if (!posbuf) { - if (pos != pos_stack) free(pos); - return 0; + const NSUInteger pos_bytes = (NSUInteger)n_tok * sizeof(int32_t); + if (!use_affine_position) { + pos = pos_stack; + if (n_tok > (uint32_t)(sizeof(pos_stack) / sizeof(pos_stack[0]))) { + pos = malloc((size_t)n_tok * sizeof(*pos)); + if (!pos) { + fprintf(stderr, "ds4: failed to allocate Metal RoPE position buffer\n"); + return 0; + } + } + for (uint32_t t = 0; t < n_tok; t++) { + pos[t] = (int32_t)(pos0 + t * pos_step); + } + + if (pos_bytes > 4096u) { + /* + * Metal inline setBytes data is meant for small constants. Long + * prefill RoPE calls need thousands of positions; passing that much + * inline can make the Apple driver abort the process. + */ + posbuf = ds4_gpu_new_transient_buffer( + pos_bytes, "ds4_rope_positions"); + if (!posbuf) { + if (pos != pos_stack) free(pos); + return 0; + } + memcpy([posbuf contents], pos, pos_bytes); } - memcpy([posbuf contents], pos, pos_bytes); } - const NSUInteger nth = (NSUInteger)(head_dim < 256u ? head_dim : 256u); + ds4_gpu_rope_affine_pair_args affine_args; + if (use_affine_position) { + const uint64_t row_bytes = (uint64_t)head_dim * sizeof(float); + affine_args = (ds4_gpu_rope_affine_pair_args) { + .row_bytes = row_bytes, + .token_bytes = (uint64_t)n_head * row_bytes, + .head_dim = (int32_t)head_dim, + .n_dims = args->n_dims, + .n_ctx_orig = args->n_ctx_orig, + .inverse = args->inverse, + .pos0 = pos0, + .pos_step = pos_step, + .freq_base = args->freq_base, + .freq_scale = args->freq_scale, + .ext_factor = args->ext_factor, + .attn_factor = args->attn_factor, + .beta_fast = args->beta_fast, + .beta_slow = args->beta_slow, + }; + } + const NSUInteger reference_nth = + (NSUInteger)(head_dim < 256u ? head_dim : 256u); + const NSUInteger pair_nth = + (NSUInteger)(tail_threads < 256u ? tail_threads : 256u); + const NSUInteger nth = use_shared_coeff ? 256u : + (use_inplace_pair ? pair_nth : reference_nth); + const NSUInteger head_groups = use_shared_coeff ? + (NSUInteger)((n_head + 3u) / 4u) : (NSUInteger)n_head; id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_rope_tail_batch_pipeline]; - [enc setBytes:args length:sizeof(*args) atIndex:0]; - [enc setBuffer:xbuf offset:xoff atIndex:1]; - if (posbuf) { - [enc setBuffer:posbuf offset:0 atIndex:2]; + id pipeline = use_affine_position ? + g_rope_tail_inplace_pair_affine_pipeline : + (use_shared_coeff ? + g_rope_tail_inplace_pair_shared4_pipeline : + (use_inplace_pair ? + g_rope_tail_inplace_pair_pipeline : g_rope_tail_batch_pipeline)); + [enc setComputePipelineState:pipeline]; + if (use_affine_position) { + [enc setBytes:&affine_args length:sizeof(affine_args) atIndex:0]; } else { - [enc setBytes:pos length:pos_bytes atIndex:2]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + } + [enc setBuffer:xbuf offset:xoff atIndex:1]; + if (!use_affine_position) { + if (posbuf) { + [enc setBuffer:posbuf offset:0 atIndex:2]; + } else { + [enc setBytes:pos length:pos_bytes atIndex:2]; + } + [enc setBuffer:xbuf offset:xoff atIndex:3]; } - [enc setBuffer:xbuf offset:xoff atIndex:3]; [enc setBuffer:xbuf offset:xoff atIndex:4]; - [enc dispatchThreadgroups:MTLSizeMake(n_head, n_tok, 1) + [enc dispatchThreadgroups:MTLSizeMake(head_groups, n_tok, 1) threadsPerThreadgroup:MTLSizeMake(nth ? nth : 1u, 1, 1)]; ds4_gpu_end_compute_encoder(cb, enc); - if (pos != pos_stack) free(pos); + if (pos && pos != pos_stack) free(pos); return 1; } @@ -4350,6 +4735,21 @@ static int ds4_gpu_encode_rope_tail_inplace( uint32_t width; } ds4_gpu_dsv4_ratio4_shift_args; +typedef struct { + uint32_t head_dim; + uint32_t n_comp; + uint32_t replay; + uint32_t n_threads; +} ds4_gpu_dsv4_compressor_pack_ratio4_args; + +typedef struct { + int64_t n_rows; + uint32_t head_dim; + uint32_t n_comp; + uint32_t replay; + uint32_t pad; +} ds4_gpu_dsv4_softmax_pool_ratio4_direct_args; + typedef struct { uint32_t width; uint32_t ratio; @@ -4373,6 +4773,13 @@ static int ds4_gpu_encode_rope_tail_inplace( uint64_t nb1; } ds4_gpu_dsv4_softmax_pool_args; +typedef struct { + uint32_t width; + uint32_t ratio; + uint32_t pos0; + uint32_t n_tokens; +} ds4_gpu_dsv4_compressor_score_ape_args; + typedef struct { int32_t ne00; int32_t ne01; @@ -4800,6 +5207,23 @@ int ds4_gpu_init(void) { return 0; } + fn = [library newFunctionWithName:@"kernel_cpy_contig_f32_f16_4"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_cpy_contig_f32_f16_4 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_cpy_contig_f32_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_cpy_contig_f32_f16_pipeline) { + fprintf(stderr, "ds4: Metal kernel_cpy_contig_f32_f16_4 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + fn = [library newFunctionWithName:@"kernel_cpy_f16_f32"]; if (!fn) { fprintf(stderr, "ds4: Metal kernel_cpy_f16_f32 function not found\n"); @@ -4817,6 +5241,55 @@ int ds4_gpu_init(void) { return 0; } + fn = [library newFunctionWithName:@"kernel_cpy_contig_f16_f32_4"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f32_4 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_cpy_contig_f16_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_cpy_contig_f16_f32_pipeline) { + fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f32_4 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_cpy_contig_f16_f16_bits_4"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f16_bits_4 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_cpy_contig_f16_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_cpy_contig_f16_f16_pipeline) { + fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f16_bits_4 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_dsv4_flash_kv_stage_f16"]; + if (fn) { + g_flash_kv_stage_f16_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_flash_kv_stage_f16_pipeline) { + fprintf(stderr, + "ds4: optional Metal gathered KV staging pipeline unavailable: %s\n", + [[error localizedDescription] UTF8String]); + } + } else { + fprintf(stderr, + "ds4: optional Metal gathered KV staging kernel unavailable\n"); + } + fn = [library newFunctionWithName:@"kernel_dsv4_fp8_kv_quantize_f32"]; if (!fn) { fprintf(stderr, "ds4: Metal kernel_dsv4_fp8_kv_quantize_f32 function not found\n"); @@ -5731,6 +6204,56 @@ int ds4_gpu_init(void) { return 0; } + error = nil; + fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32_inplace_pair"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_rope_tail_f32_inplace_pair function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_rope_tail_inplace_pair_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_rope_tail_inplace_pair_pipeline) { + fprintf(stderr, + "ds4: Metal kernel_dsv4_rope_tail_f32_inplace_pair pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32_inplace_pair_shared4"]; + if (!fn) { + fprintf(stderr, + "ds4: optional Metal shared-head RoPE kernel unavailable; using per-head path\n"); + } else { + g_rope_tail_inplace_pair_shared4_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_rope_tail_inplace_pair_shared4_pipeline) { + fprintf(stderr, + "ds4: optional Metal shared-head RoPE pipeline unavailable; using per-head path: %s\n", + [[error localizedDescription] UTF8String]); + } + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32_inplace_pair_affine"]; + if (fn) { + g_rope_tail_inplace_pair_affine_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_rope_tail_inplace_pair_affine_pipeline) { + fprintf(stderr, + "ds4: optional Metal affine-position RoPE pair pipeline unavailable: %s\n", + [[error localizedDescription] UTF8String]); + } + } else { + fprintf(stderr, + "ds4: optional Metal affine-position RoPE pair kernel unavailable\n"); + } + + error = nil; fn = [library newFunctionWithName:@"kernel_dsv4_softmax_pool"]; if (!fn) { fprintf(stderr, "ds4: Metal kernel_dsv4_softmax_pool function not found\n"); @@ -5947,6 +6470,36 @@ int ds4_gpu_init(void) { return 0; } + error = nil; + fn = [library newFunctionWithName:@"kernel_dsv4_hc_weighted_sum_norm4"]; + if (fn) { + g_hc_weighted_sum_norm_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_hc_weighted_sum_norm_pipeline) { + fprintf(stderr, + "ds4: optional Metal output HC sum/RMSNorm pipeline unavailable: %s\n", + [[error localizedDescription] UTF8String]); + } + } else { + fprintf(stderr, + "ds4: optional Metal output HC sum/RMSNorm kernel unavailable\n"); + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_dsv4_output_hc_weights4"]; + if (fn) { + g_output_hc_weights4_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_output_hc_weights4_pipeline) { + fprintf(stderr, + "ds4: optional Metal output HC weights4 pipeline unavailable: %s\n", + [[error localizedDescription] UTF8String]); + } + } else { + fprintf(stderr, + "ds4: optional Metal output HC weights4 kernel unavailable\n"); + } + MTLFunctionConstantValues *unary_sigmoid_constants = [[MTLFunctionConstantValues alloc] init]; int16_t unary_sigmoid_op = 102; bool unary_cnt = false; @@ -6163,6 +6716,12 @@ int ds4_gpu_init(void) { ds4_gpu_get_pipeline("kernel_dsv4_indexer_score_one_direct"); g_dsv4_compressor_store_one_pipeline = ds4_gpu_get_pipeline("kernel_dsv4_compressor_store_one"); + g_dsv4_compressor_pack_ratio4_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_compressor_pack_ratio4"); + g_dsv4_softmax_pool_ratio4_direct_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_softmax_pool_ratio4_direct"); + g_rms_norm_scale_pipeline = + ds4_gpu_get_pipeline("kernel_rms_norm_scale_f32_4"); g_dsv4_sort_i32_rows_asc_pipeline = ds4_gpu_get_pipeline("kernel_dsv4_sort_i32_rows_asc"); g_dsv4_indexed_attention_heads8_pipeline = @@ -6173,8 +6732,14 @@ int ds4_gpu_init(void) { ds4_gpu_get_pipeline("kernel_dsv4_softplus_sqrt_f32_4"); g_dsv4_router_finalize_one_pipeline = ds4_gpu_get_pipeline("kernel_dsv4_router_finalize_one"); + g_dsv4_router_finalize_one_simd_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_router_finalize_one_simd"); + g_dsv4_router_finalize_weights_one_simd_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_router_finalize_weights_one_simd"); g_dsv4_router_weights_one_pipeline = ds4_gpu_get_pipeline("kernel_dsv4_router_weights_one"); + g_dsv4_router_weights_batch_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_router_weights_batch"); g_dsv4_hc_expand4_pipeline = ds4_gpu_get_pipeline("kernel_dsv4_hc_expand4"); if (!g_dsv4_indexer_score_one_direct_pipeline || @@ -6709,7 +7274,11 @@ void ds4_gpu_cleanup(void) { g_concat_pipeline = nil; g_cpy_f32_f32_pipeline = nil; g_cpy_f32_f16_pipeline = nil; + g_cpy_contig_f32_f16_pipeline = nil; g_cpy_f16_f32_pipeline = nil; + g_cpy_contig_f16_f32_pipeline = nil; + g_cpy_contig_f16_f16_pipeline = nil; + g_flash_kv_stage_f16_pipeline = nil; g_swiglu_pipeline = nil; g_add_pipeline = nil; g_moe_sum6_pipeline = nil; @@ -6726,11 +7295,14 @@ void ds4_gpu_cleanup(void) { g_unary_fill_f16_pipeline = nil; g_rms_norm_pipeline = nil; g_rms_norm_plain_pipeline = nil; + g_rms_norm_scale_pipeline = nil; g_dsv4_qkv_rms_norm_pipeline = nil; g_hc_split_sinkhorn_pipeline = nil; g_hc_split_weighted_sum_pipeline = nil; g_hc_split_weighted_sum_norm_pipeline = nil; g_hc_weighted_sum_pipeline = nil; + g_hc_weighted_sum_norm_pipeline = nil; + g_output_hc_weights4_pipeline = nil; g_hc_expand_pipeline = nil; g_moe_mul_mv_id_iq2_xxs_pipeline = nil; g_moe_mul_mv_id_iq2_xxs_pair_pipeline = nil; @@ -6767,10 +7339,15 @@ void ds4_gpu_cleanup(void) { g_moe_table_q4_pair_up_encoder = nil; g_moe_table_q4_sum_down_encoder = nil; g_rope_tail_batch_pipeline = nil; + g_rope_tail_inplace_pair_pipeline = nil; + g_rope_tail_inplace_pair_shared4_pipeline = nil; + g_rope_tail_inplace_pair_affine_pipeline = nil; g_dsv4_fp8_kv_quantize_pipeline = nil; g_dsv4_indexer_qat_pipeline = nil; g_dsv4_kv_fp8_store_pipeline = nil; g_dsv4_ratio4_shift_pipeline = nil; + g_dsv4_compressor_pack_ratio4_pipeline = nil; + g_dsv4_softmax_pool_ratio4_direct_pipeline = nil; g_dsv4_softmax_pool_pipeline = nil; g_soft_max_f32_pipeline = nil; g_soft_max_f32_4_pipeline = nil; @@ -6787,12 +7364,17 @@ void ds4_gpu_cleanup(void) { g_dsv4_indexed_attention_heads8_rb16_pipeline = nil; g_dsv4_softplus_sqrt_pipeline = nil; g_dsv4_router_finalize_one_pipeline = nil; + g_dsv4_router_finalize_one_simd_pipeline = nil; + g_dsv4_router_finalize_weights_one_simd_pipeline = nil; g_dsv4_router_weights_one_pipeline = nil; + g_dsv4_router_weights_batch_pipeline = nil; g_dsv4_hc_expand4_pipeline = nil; g_flash_attn_mask_buffer = nil; + g_flash_attn_zero_mask_buffer = nil; g_flash_attn_pad_buffer = nil; g_flash_attn_tmp_buffer = nil; g_flash_attn_blk_buffer = nil; + ds4_gpu_clear_zero_prefix_prefill_mask_cache(); g_flash_attn_ring_buffer = nil; g_flash_attn_kv_buffer = nil; g_compressor_pool_kv_buffer = nil; @@ -6826,6 +7408,7 @@ void ds4_gpu_cleanup(void) { g_model_mapped_max_tensor_bytes = 0; ds4_gpu_tensor_tracking_reset(); g_flash_attn_mask_bytes = 0; + g_flash_attn_zero_mask_bytes = 0; g_flash_attn_pad_bytes = 0; g_flash_attn_tmp_bytes = 0; g_flash_attn_blk_bytes = 0; @@ -12430,9 +13013,9 @@ static int ds4_gpu_indexer_scores_batch_tensor( [enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(index_comp) atIndex:3]; [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:4]; if (use_nax) { - const NSUInteger q_shared = 16u * 32u; + const NSUInteger q_shared = 2u * 32u * 32u; const NSUInteger k_shared = 32u * 128u; - const NSUInteger dot_shared = 16u * 32u; + const NSUInteger dot_shared = 32u * 32u; [enc setThreadgroupMemoryLength:(q_shared + k_shared) * sizeof(uint16_t) + dot_shared * sizeof(float) atIndex:0]; [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + 31u) / 32u, @@ -12575,7 +13158,9 @@ int ds4_gpu_indexer_topk_tensor( .ne3 = 1, .top_k = block_top_k, }; - const NSUInteger smem = (((NSUInteger)nth * sizeof(int32_t)) + 15u) & ~(NSUInteger)15u; + // kernel_argsort_f32_i32_desc stages the block's scores behind the + // index array: nth int32 indices + nth float scores. + const NSUInteger smem = (((NSUInteger)nth * (sizeof(int32_t) + sizeof(float))) + 15u) & ~(NSUInteger)15u; NSUInteger cur_off = 0; NSUInteger next_off = (NSUInteger)scratch_row_bytes * n_tokens; @@ -12768,6 +13353,21 @@ static int ds4_gpu_matmul_q8_0_legacy_tensor( ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); if (out_dim > 65536u) mv_dispatch.nsg = 8; + const bool force_output_nr4 = + getenv("DS4_METAL_ENABLE_OUTPUT_Q8_NR4") != NULL; + const bool output_shape = + in_dim == 4096u && out_dim == 129280u; + const bool use_output_nr4 = + !g_quality_mode && (out_dim & 3u) == 0u && + (force_output_nr4 || + (output_shape && ds4_gpu_device_name_contains("M3"))) && + getenv("DS4_METAL_DISABLE_M3_OUTPUT_Q8_NR4") == NULL; + if (use_output_nr4) { + mv_dispatch.function_name = + "kernel_mul_mv_q8_0_f32_nr4"; + mv_dispatch.nr0 = 4; + mv_dispatch.smem = 32u * 4u * sizeof(float); + } mv_args.nr0 = mv_dispatch.nr0; id pipeline = ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg); @@ -12859,7 +13459,7 @@ static int ds4_gpu_matmul_q8_0_legacy_tensor( [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc setThreadgroupMemoryLength:64u * 32u * sizeof(uint16_t) atIndex:0]; + [enc setThreadgroupMemoryLength:2u * 64u * 32u * sizeof(uint16_t) atIndex:0]; [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)(n_tok / nax_tile_n), (NSUInteger)out_dim / 64u, 1) @@ -12984,28 +13584,108 @@ int ds4_gpu_matmul_q8_0_pair_tensor( uint64_t out1_dim, const ds4_gpu_tensor *x, uint64_t n_tok) { - (void)out0; (void)out1; (void)model_map; (void)model_size; - (void)weight0_offset; (void)weight1_offset; - (void)in_dim; (void)out0_dim; (void)out1_dim; (void)x; (void)n_tok; - return 0; -} + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out0 || !out1 || !model_map || !x || n_tok != 1 || + out0_dim == 0 || out1_dim == 0 || (in_dim & 31u) != 0 || + in_dim > UINT32_MAX || out0_dim > UINT32_MAX || out1_dim > UINT32_MAX) { + return 0; + } -int ds4_gpu_matmul_q8_0_f16_out_tensor( - ds4_gpu_tensor *out_h, - 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) { - (void)out_h; (void)model_map; (void)model_size; (void)weight_offset; - (void)in_dim; (void)out_dim; (void)x; (void)n_tok; - return 0; -} + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id out0buf = ds4_gpu_tensor_buffer(out0); + id out1buf = ds4_gpu_tensor_buffer(out1); + const uint64_t x_bytes = in_dim * sizeof(float); + const uint64_t out0_bytes = out0_dim * sizeof(float); + const uint64_t out1_bytes = out1_dim * sizeof(float); + if (!xbuf || !out0buf || !out1buf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(out0) < out0_bytes || + ds4_gpu_tensor_bytes(out1) < out1_bytes) { + fprintf(stderr, "ds4: Metal paired Q8_0 matvec received undersized activation buffers\n"); + return 0; + } -int ds4_gpu_shared_gate_up_swiglu_q8_0_tensor( - ds4_gpu_tensor *gate, + const uint64_t row_bytes = (in_dim / 32u) * 34u; + const uint64_t weight0_bytes = out0_dim * row_bytes; + const uint64_t weight1_bytes = out1_dim * row_bytes; + if (weight0_offset > model_size || weight0_bytes > model_size - weight0_offset || + weight1_offset > model_size || weight1_bytes > model_size - weight1_offset) { + fprintf(stderr, "ds4: Metal paired Q8_0 matvec range is outside the mapped model\n"); + return 0; + } + + uint64_t inner0 = 0; + uint64_t inner1 = 0; + id weight0buf = + ds4_gpu_wrap_model_range(model_map, model_size, + weight0_offset, weight0_bytes, &inner0); + id weight1buf = + ds4_gpu_wrap_model_range(model_map, model_size, + weight1_offset, weight1_bytes, &inner1); + if (!weight0buf || !weight1buf) return 0; + + ds4_gpu_mv_dispatch dispatch0 = ds4_gpu_make_q8_0_mv_dispatch(); + ds4_gpu_mv_dispatch dispatch1 = ds4_gpu_make_q8_0_mv_dispatch(); + if (out0_dim > 65536u) dispatch0.nsg = 8; + if (out1_dim > 65536u) dispatch1.nsg = 8; + /* A common threadgroup shape is required to retain each standalone + * reduction tree. Mixed 4/8-simdgroup extents use the existing fallback. */ + if (dispatch0.nsg != dispatch1.nsg) return 0; + + ds4_gpu_q8_0_matvec_args args0 = ds4_gpu_make_q8_0_mv_args(in_dim, out0_dim); + ds4_gpu_q8_0_matvec_args args1 = ds4_gpu_make_q8_0_mv_args(in_dim, out1_dim); + args0.nr0 = dispatch0.nr0; + args1.nr0 = dispatch1.nr0; + id pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_q8_0_f32_pair", dispatch0.nsg); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args0 length:sizeof(args0) atIndex:0]; + [enc setBytes:&args1 length:sizeof(args1) atIndex:1]; + [enc setBuffer:weight0buf offset:(NSUInteger)inner0 atIndex:2]; + [enc setBuffer:weight1buf offset:(NSUInteger)inner1 atIndex:3]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4]; + [enc setBuffer:out0buf offset:ds4_gpu_tensor_offset(out0) atIndex:5]; + [enc setBuffer:out1buf offset:ds4_gpu_tensor_offset(out1) atIndex:6]; + [enc setThreadgroupMemoryLength:2u * dispatch0.smem atIndex:0]; + const uint64_t max_out_dim = out0_dim > out1_dim ? out0_dim : out1_dim; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)max_out_dim + + (NSUInteger)dispatch0.nr0 - 1u) / + (NSUInteger)dispatch0.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)dispatch0.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "paired Q8_0 matvec")) return 0; + } + + return 1; +} + +int ds4_gpu_matmul_q8_0_f16_out_tensor( + ds4_gpu_tensor *out_h, + 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) { + (void)out_h; (void)model_map; (void)model_size; (void)weight_offset; + (void)in_dim; (void)out_dim; (void)x; (void)n_tok; + return 0; +} + +int ds4_gpu_shared_gate_up_swiglu_q8_0_tensor( + ds4_gpu_tensor *gate, ds4_gpu_tensor *up, ds4_gpu_tensor *mid, const void *model_map, @@ -13228,7 +13908,7 @@ int ds4_gpu_matmul_f16_tensor( [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; - [enc setThreadgroupMemoryLength:64u * 32u * sizeof(uint16_t) atIndex:0]; + [enc setThreadgroupMemoryLength:2u * 64u * 32u * sizeof(uint16_t) atIndex:0]; [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)(n_tok / nax_tile_n), (NSUInteger)out_dim / 64u, 1) @@ -13355,6 +14035,144 @@ int ds4_gpu_matmul_f16_pair_tensor( return 1; } +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) { + if (!g_initialized && !ds4_gpu_init()) return -1; + const bool force = + getenv("DS4_METAL_ENABLE_COMPRESSOR_PAIR_STATE_STORE") != NULL; + if ((g_quality_mode || + (!ds4_gpu_device_name_contains("M3") && !force)) || + getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_PAIR_STATE_STORE") != NULL || + getenv("DS4_METAL_DISABLE_COMPRESSOR_PAIR_PROJ") != NULL || + getenv("DS4_METAL_DISABLE_COMPRESSOR_STORE_ONE") != NULL) { + return 0; + } + if (!out_kv || !out_score || !state_kv || !state_score || !model_map || !x || + in_dim == 0 || width == 0 || ratio == 0 || + (ape_type != 0u && ape_type != 1u)) { + return -1; + } + if (in_dim != 4096u || + (width != 256u && width != 512u && width != 1024u) || + (ratio != 4u && ratio != 128u)) { + return 0; + } + + @autoreleasepool { + const uint32_t state_rows = ratio == 4u ? 2u * ratio : ratio; + const uint64_t row_bytes = in_dim * sizeof(uint16_t); + const uint64_t weight_bytes = row_bytes * width; + const uint64_t out_bytes = (uint64_t)width * sizeof(float); + const uint64_t state_bytes = + (uint64_t)state_rows * width * sizeof(float); + const uint64_t ape_elem = ape_type == 1u ? sizeof(uint16_t) : sizeof(float); + const uint64_t ape_bytes = (uint64_t)ratio * width * ape_elem; + if (weight_kv_offset > model_size || + weight_bytes > model_size - weight_kv_offset || + weight_score_offset > model_size || + weight_bytes > model_size - weight_score_offset || + ape_offset > model_size || ape_bytes > model_size - ape_offset) { + return -1; + } + + id xbuf = ds4_gpu_tensor_buffer(x); + id outkvbuf = ds4_gpu_tensor_buffer(out_kv); + id outscorebuf = ds4_gpu_tensor_buffer(out_score); + id statekvbuf = ds4_gpu_tensor_buffer(state_kv); + id statescbuf = ds4_gpu_tensor_buffer(state_score); + if (!xbuf || !outkvbuf || !outscorebuf || !statekvbuf || !statescbuf || + ds4_gpu_tensor_bytes(x) < in_dim * sizeof(float) || + ds4_gpu_tensor_bytes(out_kv) < out_bytes || + ds4_gpu_tensor_bytes(out_score) < out_bytes || + ds4_gpu_tensor_bytes(state_kv) < state_bytes || + ds4_gpu_tensor_bytes(state_score) < state_bytes) { + return -1; + } + + uint64_t weight_kv_inner = 0; + uint64_t weight_score_inner = 0; + uint64_t ape_inner = 0; + id weightkvbuf = ds4_gpu_wrap_model_range( + model_map, model_size, weight_kv_offset, weight_bytes, + &weight_kv_inner); + id weightscorebuf = ds4_gpu_wrap_model_range( + model_map, model_size, weight_score_offset, weight_bytes, + &weight_score_inner); + id apebuf = ds4_gpu_wrap_model_range( + model_map, model_size, ape_offset, ape_bytes, &ape_inner); + if (!weightkvbuf || !weightscorebuf || !apebuf) return -1; + + ds4_gpu_f16_matvec_args mv_args = + ds4_gpu_make_f16_mv_args(in_dim, width); + ds4_gpu_mv_dispatch mv_dispatch = + ds4_gpu_make_plain_mv_dispatch(in_dim, 0); + if (ds4_gpu_use_compressor_pair_nr4() && + (width == 512u || width == 1024u) && in_dim >= 4096u) { + mv_dispatch.nr0 = 4; + mv_dispatch.smem = 32u * 4u * sizeof(float); + } + mv_args.nr0 = mv_dispatch.nr0; + ds4_gpu_dsv4_compressor_store_one_args store_args = { + .width = width, + .ratio = ratio, + .pos = pos, + .ape_type = ape_type, + }; + id pipeline = ds4_gpu_get_mul_mv_pipeline( + "kernel_mul_mv_f16_f32_pair_compressor_store_4", + mv_dispatch.nsg); + if (!pipeline) return -1; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return -1; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBytes:&store_args length:sizeof(store_args) atIndex:1]; + [enc setBuffer:weightkvbuf offset:(NSUInteger)weight_kv_inner atIndex:2]; + [enc setBuffer:weightscorebuf offset:(NSUInteger)weight_score_inner atIndex:3]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4]; + [enc setBuffer:outkvbuf offset:ds4_gpu_tensor_offset(out_kv) atIndex:5]; + [enc setBuffer:outscorebuf offset:ds4_gpu_tensor_offset(out_score) atIndex:6]; + [enc setBuffer:apebuf offset:(NSUInteger)ape_inner atIndex:7]; + [enc setBuffer:statekvbuf offset:ds4_gpu_tensor_offset(state_kv) atIndex:8]; + [enc setBuffer:statescbuf offset:ds4_gpu_tensor_offset(state_score) atIndex:9]; + if (mv_dispatch.smem) { + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake( + ((NSUInteger)width + (NSUInteger)mv_dispatch.nr0 - 1u) / + (NSUInteger)mv_dispatch.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake( + 32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer( + cb, owned, "F16 paired matvec compressor state store")) { + return -1; + } + } + + return 1; +} + int ds4_gpu_matmul_f32_tensor( ds4_gpu_tensor *out, const void *model_map, @@ -13512,6 +14330,159 @@ int ds4_gpu_rms_norm_plain_rows_tensor( return 1; } +static int ds4_gpu_hc_rms_scale_project_mode( + uint32_t in_dim, + uint32_t out_dim, + uint32_t n_rows) { + const bool hard_shape = n_rows > 8u && + (in_dim == 16384u || in_dim == 28672u) && out_dim == 24u; + if (!hard_shape) return 0; + + const bool force = + getenv("DS4_METAL_ENABLE_HC_RMS_SCALE_PROJ") != NULL; + if (getenv("DS4_METAL_DISABLE_M3_HC_RMS_SCALE_PROJ") != NULL || + (!ds4_gpu_device_name_contains("M3") && !force)) { + return 0; + } + if (g_rms_norm_scale_pipeline == nil) { + return force ? -1 : 0; + } + return 1; +} + +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) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !scale_scratch || !model_map || !x || + in_dim == 0u || out_dim == 0u || n_rows == 0u || + n_rows > INT32_MAX) { + return 0; + } + + const int mode = + ds4_gpu_hc_rms_scale_project_mode(in_dim, out_dim, n_rows); + if (mode < 0) return 0; + if (mode == 0) { + return ds4_gpu_rms_norm_plain_rows_tensor( + scale_scratch, x, in_dim, n_rows, eps) != 0 && + ds4_gpu_matmul_f16_tensor( + out, model_map, model_size, weight_offset, + in_dim, out_dim, scale_scratch, n_rows) != 0; + } + + @autoreleasepool { + const uint64_t x_elems = (uint64_t)in_dim * n_rows; + const uint64_t out_elems = (uint64_t)out_dim * n_rows; + if (x_elems > UINT64_MAX / sizeof(float) || + out_elems > UINT64_MAX / sizeof(float)) { + return 0; + } + const uint64_t x_bytes = x_elems * sizeof(float); + const uint64_t out_bytes = out_elems * sizeof(float); + const uint64_t scale_bytes = (uint64_t)n_rows * sizeof(float); + + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + id scalebuf = ds4_gpu_tensor_buffer(scale_scratch); + if (!xbuf || !outbuf || !scalebuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes || + ds4_gpu_tensor_bytes(scale_scratch) < scale_bytes) { + fprintf(stderr, "ds4: Metal HC RMS scale projection received undersized activation buffers\n"); + return 0; + } + + const uint64_t weight_row_bytes = (uint64_t)in_dim * sizeof(uint16_t); + if ((uint64_t)out_dim > UINT64_MAX / weight_row_bytes) return 0; + const uint64_t weight_bytes = (uint64_t)out_dim * weight_row_bytes; + if (weight_offset > model_size || + weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal HC RMS scale projection weight range is outside the mapped model\n"); + return 0; + } + + const bool bc_inp = (in_dim % 32u) != 0u; + const bool bc_out = (out_dim % 64u) != 0u || + (n_rows % 32u) != 0u; + id mm_pipeline = + ds4_gpu_get_mul_mm_pipeline( + "kernel_mul_mm_f16_f32_scaled", bc_inp, bc_out); + if (!mm_pipeline) { + if (getenv("DS4_METAL_ENABLE_HC_RMS_SCALE_PROJ") != NULL) { + return 0; + } + return ds4_gpu_rms_norm_plain_rows_tensor( + scale_scratch, x, in_dim, n_rows, eps) != 0 && + ds4_gpu_matmul_f16_tensor( + out, model_map, model_size, weight_offset, + in_dim, out_dim, scale_scratch, n_rows) != 0; + } + + id scale_pipeline = ds4_gpu_hot_pipeline( + g_rms_norm_scale_pipeline, "kernel_rms_norm_scale_f32_4"); + if (!scale_pipeline) return 0; + + uint64_t weight_inner = 0; + id weightbuf = ds4_gpu_wrap_model_range( + model_map, model_size, weight_offset, weight_bytes, &weight_inner); + if (!weightbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_rms_norm_args norm_args = + ds4_gpu_make_rms_norm_args(in_dim, n_rows, eps); + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:scale_pipeline]; + [enc setBytes:&norm_args length:sizeof(norm_args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; + [enc setBuffer:scalebuf + offset:ds4_gpu_tensor_offset(scale_scratch) + atIndex:2]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake( + ds4_gpu_rms_norm_threads(in_dim), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + ds4_gpu_mul_mm_args mm_args = ds4_gpu_make_mm_args( + in_dim, out_dim, n_rows, weight_row_bytes); + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:mm_pipeline]; + [enc setBytes:&mm_args length:sizeof(mm_args) atIndex:0]; + [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setBuffer:scalebuf + offset:ds4_gpu_tensor_offset(scale_scratch) + atIndex:4]; + [enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake( + ((NSUInteger)n_rows + 31u) / 32u, + ((NSUInteger)out_dim + 63u) / 64u, + 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer( + cb, owned, "HC RMS scale F16 projection")) { + return 0; + } + } + + return 1; +} + int ds4_gpu_rms_norm_weight_tensor( ds4_gpu_tensor *out, const ds4_gpu_tensor *x, @@ -13960,34 +14931,20 @@ static int ds4_gpu_encode_f16_round_copy_for_raw_store( return 0; } - ds4_gpu_cpy_args f32_to_f16 = - ds4_gpu_make_cpy_1d_args(n, sizeof(float), sizeof(uint16_t)); - ds4_gpu_cpy_args f16_to_f32 = - ds4_gpu_make_cpy_1d_args(n, sizeof(uint16_t), sizeof(float)); - const NSUInteger nth_f32_f16 = ds4_gpu_cpy_threads(n, g_cpy_f32_f16_pipeline); - const NSUInteger nth_f16_f32 = ds4_gpu_cpy_threads(n, g_cpy_f16_f32_pipeline); - const NSUInteger groups_f32_f16 = ((NSUInteger)n + nth_f32_f16 - 1u) / nth_f32_f16; - const NSUInteger groups_f16_f32 = ((NSUInteger)n + nth_f16_f32 - 1u) / nth_f16_f32; - - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_cpy_f32_f16_pipeline]; - [enc setBytes:&f32_to_f16 length:sizeof(f32_to_f16) atIndex:0]; - [enc setBuffer:srcbuf offset:ds4_gpu_tensor_offset(src) atIndex:1]; - [enc setBuffer:g_f16_round_scratch_buffer offset:0 atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(groups_f32_f16, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth_f32_f16, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_cpy_f16_f32_pipeline]; - [enc setBytes:&f16_to_f32 length:sizeof(f16_to_f32) atIndex:0]; - [enc setBuffer:g_f16_round_scratch_buffer offset:0 atIndex:1]; - [enc setBuffer:g_raw_store_round_buffer offset:0 atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(groups_f16_f32, 1, 1) - threadsPerThreadgroup:MTLSizeMake(nth_f16_f32, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); - - return 1; + if (!ds4_gpu_encode_cpy_f32_f16_1d(cb, + srcbuf, + ds4_gpu_tensor_offset(src), + g_f16_round_scratch_buffer, + 0, + n)) { + return 0; + } + return ds4_gpu_encode_cpy_f16_f32_1d(cb, + g_f16_round_scratch_buffer, + 0, + g_raw_store_round_buffer, + 0, + n); } static int ds4_gpu_encode_set_rows_f32_i32( @@ -14288,6 +15245,42 @@ static int ds4_gpu_encode_compressor_score_with_ape( return 0; } const uint32_t total_elems = (uint32_t)total_elems64; + + const bool force_fused = + getenv("DS4_METAL_ENABLE_COMPRESSOR_APE_ADD") != NULL; + const bool use_fused = + (ds4_gpu_device_name_contains("M3") || force_fused) && + getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_APE_ADD") == NULL; + if (use_fused) { + id pipeline = ds4_gpu_get_pipeline( + ape_type == 1u ? "kernel_dsv4_compressor_score_ape_f16" + : "kernel_dsv4_compressor_score_ape_f32"); + if (pipeline) { + ds4_gpu_dsv4_compressor_score_ape_args args = { + .width = width, + .ratio = ratio, + .pos0 = pos0, + .n_tokens = n_tokens, + }; + NSUInteger nth = pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth == 0) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:score_src offset:score_src_offset atIndex:1]; + [enc setBuffer:apebuf offset:ape_offset atIndex:2]; + [enc setBuffer:score_dst offset:score_dst_offset atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake( + ((NSUInteger)total_elems + nth - 1u) / nth, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; + } + if (force_fused) return 0; + } + const NSUInteger scratch_bytes = (NSUInteger)total_elems * sizeof(float); if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_ape_buffer, &g_compressor_store_ape_bytes, @@ -14973,6 +15966,163 @@ static int ds4_gpu_encode_concat_f32_dim1( return 1; } +static int ds4_gpu_compressor_pack_ratio4_fusion_mode(uint32_t head_dim) { + const bool force = + getenv("DS4_METAL_ENABLE_COMPRESSOR_RATIO4_PACK_FUSION") != NULL; + const bool default_shape = head_dim == 128u || head_dim == 512u; + if (getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_RATIO4_PACK_FUSION") != NULL || + (!(ds4_gpu_device_name_contains("M3") && default_shape) && !force)) { + return 0; + } + if (g_dsv4_compressor_pack_ratio4_pipeline == nil) { + return force ? -1 : 0; + } + return 1; +} + +static bool ds4_gpu_buffer_ranges_overlap( + id a, + NSUInteger a_offset, + uint64_t a_bytes, + id b, + NSUInteger b_offset, + uint64_t b_bytes) { + if (a != b || a_bytes == 0u || b_bytes == 0u) return false; + if (a_offset <= b_offset) { + return a_bytes > (uint64_t)(b_offset - a_offset); + } + return b_bytes > (uint64_t)(a_offset - b_offset); +} + +static int ds4_gpu_compressor_ratio4_direct_pool_mode( + uint32_t head_dim, + uint32_t n_comp) { + // One compressed row intentionally uses the legacy GGML reduction graph. + if (n_comp <= 1u) return 0; + + const bool force = + getenv("DS4_METAL_ENABLE_COMPRESSOR_RATIO4_DIRECT_POOL") != NULL; + const bool default_shape = head_dim == 128u || head_dim == 512u; + if (getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_RATIO4_DIRECT_POOL") != NULL || + (!(ds4_gpu_device_name_contains("M3") && default_shape) && !force)) { + return 0; + } + if (g_dsv4_softmax_pool_ratio4_direct_pipeline == nil) { + return force ? -1 : 0; + } + return 1; +} + +static int ds4_gpu_encode_compressor_ratio4_direct_pool( + id cb, + ds4_gpu_tensor *out, + id kvbuf, + NSUInteger kv_offset, + id scorebuf, + NSUInteger score_offset, + id statekvbuf, + NSUInteger state_kv_offset, + id statescbuf, + NSUInteger state_score_offset, + uint32_t head_dim, + uint32_t n_comp, + bool replay) { + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t n = (uint64_t)head_dim * n_comp; + if (!cb || !outbuf || !kvbuf || !scorebuf || !statekvbuf || !statescbuf || + head_dim == 0u || n_comp <= 1u || n > UINT32_MAX || + ds4_gpu_tensor_bytes(out) < n * sizeof(float)) { + return 0; + } + + id pipeline = ds4_gpu_hot_pipeline( + g_dsv4_softmax_pool_ratio4_direct_pipeline, + "kernel_dsv4_softmax_pool_ratio4_direct"); + if (!pipeline) return 0; + + ds4_gpu_dsv4_softmax_pool_ratio4_direct_args args = { + .n_rows = 8, + .head_dim = head_dim, + .n_comp = n_comp, + .replay = replay ? 1u : 0u, + .pad = 0u, + }; + + NSUInteger nth = 256u; + if (nth > pipeline.maxTotalThreadsPerThreadgroup) { + nth = pipeline.maxTotalThreadsPerThreadgroup; + } + if (nth == 0u) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:kvbuf offset:kv_offset atIndex:1]; + [enc setBuffer:scorebuf offset:score_offset atIndex:2]; + [enc setBuffer:statekvbuf offset:state_kv_offset atIndex:3]; + [enc setBuffer:statescbuf offset:state_score_offset atIndex:4]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:5]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n + nth - 1u) / nth, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_compressor_pack_ratio4( + id cb, + id kvbuf, + NSUInteger kv_offset, + id scorebuf, + NSUInteger score_offset, + id statekvbuf, + NSUInteger state_kv_offset, + id statescbuf, + NSUInteger state_score_offset, + uint32_t head_dim, + uint32_t n_comp, + bool replay) { + if (!cb || !kvbuf || !scorebuf || !statekvbuf || !statescbuf || + !g_compressor_pool_kv_buffer || !g_compressor_pool_score_buffer || + head_dim == 0 || n_comp == 0 || head_dim > UINT32_MAX / 2u) { + return 0; + } + + const uint64_t total_elems64 = (uint64_t)n_comp * 8u * head_dim; + if (total_elems64 > UINT32_MAX || n_comp > UINT32_MAX / 8u) return 0; + id pipeline = ds4_gpu_hot_pipeline( + g_dsv4_compressor_pack_ratio4_pipeline, + "kernel_dsv4_compressor_pack_ratio4"); + if (!pipeline) return 0; + + NSUInteger nth = head_dim; + if (nth > 256u) nth = 256u; + if (nth > pipeline.maxTotalThreadsPerThreadgroup) { + nth = pipeline.maxTotalThreadsPerThreadgroup; + } + if (nth == 0) return 0; + + ds4_gpu_dsv4_compressor_pack_ratio4_args args = { + .head_dim = head_dim, + .n_comp = n_comp, + .replay = replay ? 1u : 0u, + .n_threads = (uint32_t)nth, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:kvbuf offset:kv_offset atIndex:1]; + [enc setBuffer:scorebuf offset:score_offset atIndex:2]; + [enc setBuffer:statekvbuf offset:state_kv_offset atIndex:3]; + [enc setBuffer:statescbuf offset:state_score_offset atIndex:4]; + [enc setBuffer:g_compressor_pool_kv_buffer offset:0 atIndex:5]; + [enc setBuffer:g_compressor_pool_score_buffer offset:0 atIndex:6]; + [enc dispatchThreadgroups:MTLSizeMake(n_comp, 8u, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + static int ds4_gpu_encode_compressor_pool( id cb, ds4_gpu_tensor *out, @@ -15292,18 +16442,71 @@ int ds4_gpu_compressor_prefill_tensor( } if (ok && ratio == 4u) { + const int direct_pool_mode = + ds4_gpu_compressor_ratio4_direct_pool_mode(head_dim, n_comp); + if (ok && direct_pool_mode < 0) ok = 0; + const uint64_t direct_output_bytes = + (uint64_t)n_comp * head_dim * sizeof(float); + const uint64_t direct_input_bytes = + (uint64_t)n_comp * 4u * width * sizeof(float); + const bool direct_pool_overlap = direct_pool_mode > 0 && + ds4_gpu_buffer_ranges_overlap( + compbuf, + ds4_gpu_tensor_offset(comp_cache), + direct_output_bytes, + kvbuf, + ds4_gpu_tensor_offset(kv), + direct_input_bytes); + const bool use_direct_pool = + direct_pool_mode > 0 && !direct_pool_overlap; const NSUInteger pack_bytes = (NSUInteger)n_comp * 8u * head_dim * sizeof(float); - if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_kv_buffer, - &g_compressor_pool_kv_bytes, - pack_bytes, - "ds4_compressor_pool_kv") || - !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_buffer, - &g_compressor_pool_score_bytes, - pack_bytes, - "ds4_compressor_pool_score")) { + if (ok && !use_direct_pool && + (!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_kv_buffer, + &g_compressor_pool_kv_bytes, + pack_bytes, + "ds4_compressor_pool_kv") || + !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_buffer, + &g_compressor_pool_score_bytes, + pack_bytes, + "ds4_compressor_pool_score"))) { ok = 0; } - if (ok) { + const int pack_fusion_mode = use_direct_pool ? 0 : + ds4_gpu_compressor_pack_ratio4_fusion_mode(head_dim); + if (ok && pack_fusion_mode < 0) ok = 0; + const bool use_pack_fusion = pack_fusion_mode > 0; + if (ok && use_direct_pool) { + ok = ds4_gpu_encode_compressor_ratio4_direct_pool( + cb, + comp_cache, + kvbuf, + ds4_gpu_tensor_offset(kv), + g_compressor_store_score_buffer, + 0, + statekvbuf, + ds4_gpu_tensor_offset(state_kv), + statescbuf, + ds4_gpu_tensor_offset(state_score), + head_dim, + n_comp, + false); + } + if (ok && use_pack_fusion) { + ok = ds4_gpu_encode_compressor_pack_ratio4( + cb, + kvbuf, + ds4_gpu_tensor_offset(kv), + g_compressor_store_score_buffer, + 0, + statekvbuf, + ds4_gpu_tensor_offset(state_kv), + statescbuf, + ds4_gpu_tensor_offset(state_score), + head_dim, + n_comp, + false); + } + if (ok && !use_direct_pool && !use_pack_fusion) { ok = ds4_gpu_encode_fill_f32_rows(cb, g_compressor_pool_kv_buffer, 0, @@ -15317,7 +16520,7 @@ int ds4_gpu_compressor_prefill_tensor( 8u * n_comp, ds4_gpu_negative_infinity()); } - if (ok) { + if (ok && !use_direct_pool && !use_pack_fusion) { const uint64_t src_row_stride = (uint64_t)width * sizeof(float); const uint64_t src_plane_stride = (uint64_t)ratio * src_row_stride; const uint64_t dst_row_stride = (uint64_t)head_dim * sizeof(float); @@ -15348,7 +16551,7 @@ int ds4_gpu_compressor_prefill_tensor( dst_row_stride, dst_plane_stride); } - if (ok && n_comp > 1u) { + if (ok && !use_direct_pool && !use_pack_fusion && n_comp > 1u) { const uint64_t src_row_stride = (uint64_t)width * sizeof(float); const uint64_t src_plane_stride = (uint64_t)ratio * src_row_stride; const uint64_t dst_row_stride = (uint64_t)head_dim * sizeof(float); @@ -15378,7 +16581,7 @@ int ds4_gpu_compressor_prefill_tensor( dst_row_stride, dst_plane_stride); } - if (ok) { + if (ok && !use_direct_pool) { ok = ds4_gpu_encode_dsv4_softmax_pool(cb, comp_cache, g_compressor_pool_kv_buffer, @@ -15534,18 +16737,52 @@ int ds4_gpu_compressor_prefill_ratio4_replay_tensor( const NSUInteger score_bytes = (NSUInteger)n_tokens * width * sizeof(float); const NSUInteger pack_bytes = (NSUInteger)n_comp * 8u * head_dim * sizeof(float); + const int direct_pool_mode = + ds4_gpu_compressor_ratio4_direct_pool_mode(head_dim, n_comp); + if (ok && direct_pool_mode < 0) ok = 0; + const uint64_t direct_output_bytes = + (uint64_t)n_comp * head_dim * sizeof(float); + const uint64_t direct_input_bytes = + (uint64_t)n_tokens * width * sizeof(float); + const uint64_t direct_state_bytes = + (uint64_t)4u * width * sizeof(float); + const bool direct_pool_overlap = direct_pool_mode > 0 && + (ds4_gpu_buffer_ranges_overlap( + compbuf, + ds4_gpu_tensor_offset(comp_cache), + direct_output_bytes, + kvbuf, + ds4_gpu_tensor_offset(kv), + direct_input_bytes) || + ds4_gpu_buffer_ranges_overlap( + compbuf, + ds4_gpu_tensor_offset(comp_cache), + direct_output_bytes, + statekvbuf, + ds4_gpu_tensor_offset(state_kv), + direct_state_bytes) || + ds4_gpu_buffer_ranges_overlap( + compbuf, + ds4_gpu_tensor_offset(comp_cache), + direct_output_bytes, + statescbuf, + ds4_gpu_tensor_offset(state_score), + direct_state_bytes)); + const bool use_direct_pool = + direct_pool_mode > 0 && !direct_pool_overlap; if (ok && (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_score_buffer, &g_compressor_store_score_bytes, score_bytes, "ds4_compressor_store_score") || - !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_kv_buffer, - &g_compressor_pool_kv_bytes, - pack_bytes, - "ds4_compressor_pool_kv") || - !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_buffer, - &g_compressor_pool_score_bytes, - pack_bytes, - "ds4_compressor_pool_score"))) { + (!use_direct_pool && + (!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_kv_buffer, + &g_compressor_pool_kv_bytes, + pack_bytes, + "ds4_compressor_pool_kv") || + !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_buffer, + &g_compressor_pool_score_bytes, + pack_bytes, + "ds4_compressor_pool_score"))))) { ok = 0; } @@ -15564,7 +16801,51 @@ int ds4_gpu_compressor_prefill_ratio4_replay_tensor( n_tokens); } - if (ok) { + const uint64_t src_row_stride = (uint64_t)width * sizeof(float); + const uint64_t src_plane_stride = (uint64_t)ratio * src_row_stride; + const uint64_t dst_row_stride = (uint64_t)head_dim * sizeof(float); + const uint64_t dst_plane_stride = 8ull * dst_row_stride; + const NSUInteger state_off = ds4_gpu_tensor_offset(state_kv); + const NSUInteger state_score_off = ds4_gpu_tensor_offset(state_score); + const int pack_fusion_mode = use_direct_pool ? 0 : + ds4_gpu_compressor_pack_ratio4_fusion_mode(head_dim); + if (ok && pack_fusion_mode < 0) ok = 0; + const bool use_pack_fusion = pack_fusion_mode > 0; + + if (ok && use_direct_pool) { + ok = ds4_gpu_encode_compressor_ratio4_direct_pool( + cb, + comp_cache, + kvbuf, + ds4_gpu_tensor_offset(kv), + g_compressor_store_score_buffer, + 0, + statekvbuf, + state_off, + statescbuf, + state_score_off, + head_dim, + n_comp, + true); + } + + if (ok && use_pack_fusion) { + ok = ds4_gpu_encode_compressor_pack_ratio4( + cb, + kvbuf, + ds4_gpu_tensor_offset(kv), + g_compressor_store_score_buffer, + 0, + statekvbuf, + state_off, + statescbuf, + state_score_off, + head_dim, + n_comp, + true); + } + + if (ok && !use_direct_pool && !use_pack_fusion) { ok = ds4_gpu_encode_fill_f32_rows(cb, g_compressor_pool_kv_buffer, 0, @@ -15579,14 +16860,7 @@ int ds4_gpu_compressor_prefill_ratio4_replay_tensor( ds4_gpu_negative_infinity()); } - const uint64_t src_row_stride = (uint64_t)width * sizeof(float); - const uint64_t src_plane_stride = (uint64_t)ratio * src_row_stride; - const uint64_t dst_row_stride = (uint64_t)head_dim * sizeof(float); - const uint64_t dst_plane_stride = 8ull * dst_row_stride; - const NSUInteger state_off = ds4_gpu_tensor_offset(state_kv); - const NSUInteger state_score_off = ds4_gpu_tensor_offset(state_score); - - if (ok) { + if (ok && !use_direct_pool && !use_pack_fusion) { /* * The aligned nonzero ratio-4 path replays the current ubatch * compressor, but seeds the first compressed row with the previous @@ -15618,7 +16892,7 @@ int ds4_gpu_compressor_prefill_ratio4_replay_tensor( dst_row_stride, dst_plane_stride); } - if (ok) { + if (ok && !use_direct_pool && !use_pack_fusion) { ok = ds4_gpu_encode_cpy_f32_f32_3d(cb, kvbuf, ds4_gpu_tensor_offset(kv) + @@ -15645,7 +16919,7 @@ int ds4_gpu_compressor_prefill_ratio4_replay_tensor( dst_row_stride, dst_plane_stride); } - if (ok && n_comp > 1u) { + if (ok && !use_direct_pool && !use_pack_fusion && n_comp > 1u) { ok = ds4_gpu_encode_cpy_f32_f32_3d(cb, kvbuf, ds4_gpu_tensor_offset(kv), @@ -15671,7 +16945,7 @@ int ds4_gpu_compressor_prefill_ratio4_replay_tensor( dst_row_stride, dst_plane_stride); } - if (ok) { + if (ok && !use_direct_pool) { ok = ds4_gpu_encode_dsv4_softmax_pool(cb, comp_cache, g_compressor_pool_kv_buffer, @@ -15885,7 +17159,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) { if (!g_initialized && !ds4_gpu_init()) return 0; if (!kv_cur || !sc_cur || !state_kv || !state_score || !comp_cache || !model_map || head_dim == 0 || ratio == 0 || @@ -15926,34 +17201,36 @@ int ds4_gpu_compressor_update_tensor( return 0; } - const bool use_store_one = - getenv("DS4_METAL_DISABLE_COMPRESSOR_STORE_ONE") == NULL; - const int store_ok = use_store_one - ? ds4_gpu_compressor_store_one_tensor(kv_cur, - sc_cur, - state_kv, - state_score, - model_map, - model_size, - ape_offset, - ape_type, - width, - ratio, - pos) - : 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); - if (!store_ok) { - return 0; + if (!state_already_stored) { + const bool use_store_one = + getenv("DS4_METAL_DISABLE_COMPRESSOR_STORE_ONE") == NULL; + const int store_ok = use_store_one + ? ds4_gpu_compressor_store_one_tensor(kv_cur, + sc_cur, + state_kv, + state_score, + model_map, + model_size, + ape_offset, + ape_type, + width, + ratio, + pos) + : 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); + if (!store_ok) { + return 0; + } } if (!emit) return 1; @@ -16645,14 +17922,26 @@ static int ds4_gpu_encode_cpy_f32_f16_1d( uint32_t n) { if (!cb || !src || !dst || n == 0) return 0; - ds4_gpu_cpy_args args = - ds4_gpu_make_cpy_1d_args(n, sizeof(float), sizeof(uint16_t)); - const NSUInteger nth = ds4_gpu_cpy_threads(n, g_cpy_f32_f16_pipeline); - const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; + const int use_contiguous = + ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F32_F16_COPY") <= 0; + id pipeline = use_contiguous + ? g_cpy_contig_f32_f16_pipeline + : g_cpy_f32_f16_pipeline; + const NSUInteger work_items = use_contiguous + ? ((NSUInteger)n + 3u) / 4u + : (NSUInteger)n; + const NSUInteger nth = ds4_gpu_cpy_threads((uint32_t)work_items, pipeline); + const NSUInteger groups = (work_items + nth - 1u) / nth; id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_cpy_f32_f16_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setComputePipelineState:pipeline]; + if (use_contiguous) { + [enc setBytes:&n length:sizeof(n) atIndex:0]; + } else { + ds4_gpu_cpy_args args = + ds4_gpu_make_cpy_1d_args(n, sizeof(float), sizeof(uint16_t)); + [enc setBytes:&args length:sizeof(args) atIndex:0]; + } [enc setBuffer:src offset:src_off atIndex:1]; [enc setBuffer:dst offset:dst_off atIndex:2]; [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) @@ -16717,14 +18006,26 @@ static int ds4_gpu_encode_cpy_f16_f32_1d( uint32_t n) { if (!cb || !src || !dst || n == 0) return 0; - ds4_gpu_cpy_args args = - ds4_gpu_make_cpy_1d_args(n, sizeof(uint16_t), sizeof(float)); - const NSUInteger nth = ds4_gpu_cpy_threads(n, g_cpy_f16_f32_pipeline); - const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; + const int use_contiguous = + ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F32_F16_COPY") <= 0; + id pipeline = use_contiguous + ? g_cpy_contig_f16_f32_pipeline + : g_cpy_f16_f32_pipeline; + const NSUInteger work_items = use_contiguous + ? ((NSUInteger)n + 3u) / 4u + : (NSUInteger)n; + const NSUInteger nth = ds4_gpu_cpy_threads((uint32_t)work_items, pipeline); + const NSUInteger groups = (work_items + nth - 1u) / nth; id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_cpy_f16_f32_pipeline]; - [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setComputePipelineState:pipeline]; + if (use_contiguous) { + [enc setBytes:&n length:sizeof(n) atIndex:0]; + } else { + ds4_gpu_cpy_args args = + ds4_gpu_make_cpy_1d_args(n, sizeof(uint16_t), sizeof(float)); + [enc setBytes:&args length:sizeof(args) atIndex:0]; + } [enc setBuffer:src offset:src_off atIndex:1]; [enc setBuffer:dst offset:dst_off atIndex:2]; [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) @@ -16747,16 +18048,282 @@ static int ds4_gpu_encode_copy_to_f16_1d( if (!src_is_f16) { return ds4_gpu_encode_cpy_f32_f16_1d(cb, src, src_off, dst, dst_off, n); } - - if (g_batch_cb && cb == g_batch_cb) ds4_gpu_close_batch_encoder(); - id blit = [cb blitCommandEncoder]; - if (!blit) return 0; - [blit copyFromBuffer:src - sourceOffset:src_off - toBuffer:dst - destinationOffset:dst_off - size:(NSUInteger)n * sizeof(uint16_t)]; - [blit endEncoding]; + + if (ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F16_F16_COPY") <= 0) { + const NSUInteger work_items = ((NSUInteger)n + 3u) / 4u; + const NSUInteger nth = ds4_gpu_cpy_threads( + (uint32_t)work_items, + g_cpy_contig_f16_f16_pipeline); + const NSUInteger groups = (work_items + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_cpy_contig_f16_f16_pipeline]; + [enc setBytes:&n length:sizeof(n) atIndex:0]; + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; + } + + if (g_batch_cb && cb == g_batch_cb) ds4_gpu_close_batch_encoder(); + id blit = [cb blitCommandEncoder]; + if (!blit) return 0; + [blit copyFromBuffer:src + sourceOffset:src_off + toBuffer:dst + destinationOffset:dst_off + size:(NSUInteger)n * sizeof(uint16_t)]; + [blit endEncoding]; + return 1; +} + +static int ds4_gpu_encode_copy_raw_ring_to_f16( + id cb, + id raw, + NSUInteger raw_offset, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_raw, + uint32_t head_dim, + id dst, + NSUInteger dst_offset) { + if (!cb || !raw || !dst || raw_cap == 0 || raw_start >= raw_cap || + n_raw == 0 || n_raw > raw_cap || head_dim == 0) { + return 0; + } + + const NSUInteger raw_row_bytes = (NSUInteger)head_dim * sizeof(float); + const NSUInteger dst_row_bytes = (NSUInteger)head_dim * sizeof(uint16_t); + const uint32_t tail_rows = raw_cap - raw_start < n_raw + ? raw_cap - raw_start + : n_raw; + const uint32_t head_rows = n_raw - tail_rows; + const uint64_t tail_count = (uint64_t)tail_rows * head_dim; + const uint64_t head_count = (uint64_t)head_rows * head_dim; + const uint64_t raw_inner = (uint64_t)raw_start * raw_row_bytes; + const uint64_t dst_inner = (uint64_t)tail_rows * dst_row_bytes; + if (tail_count > UINT32_MAX || head_count > UINT32_MAX || + raw_inner > NSUIntegerMax - raw_offset || + dst_inner > NSUIntegerMax - dst_offset) { + return 0; + } + + if (tail_rows && + !ds4_gpu_encode_cpy_f32_f16_1d( + cb, + raw, + raw_offset + (NSUInteger)raw_inner, + dst, + dst_offset, + (uint32_t)tail_count)) { + return 0; + } + if (head_rows && + !ds4_gpu_encode_cpy_f32_f16_1d( + cb, + raw, + raw_offset, + dst, + dst_offset + (NSUInteger)dst_inner, + (uint32_t)head_count)) { + return 0; + } + return 1; +} + +static int ds4_gpu_encode_flash_kv_stage_f16( + id cb, + id raw, + NSUInteger raw_offset, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_raw, + id comp, + NSUInteger comp_offset, + bool comp_is_f16, + uint32_t n_comp, + uint32_t head_dim, + id dst, + NSUInteger dst_offset, + id mask, + NSUInteger mask_offset, + id pad, + NSUInteger pad_offset, + bool fuse_pad, + bool shared_pad, + bool *did_fuse_pad) { + if (did_fuse_pad) *did_fuse_pad = false; + if (!cb || !raw || !comp || !dst || raw_cap == 0 || + raw_start >= raw_cap || n_raw == 0 || n_raw > raw_cap || + n_comp == 0 || head_dim == 0) { + return 0; + } + + const bool force = + getenv("DS4_METAL_ENABLE_GATHERED_KV_STAGE") != NULL; + const bool disabled = + getenv("DS4_METAL_DISABLE_M3_GATHERED_KV_STAGE") != NULL; + const bool require = + getenv("DS4_METAL_REQUIRE_GATHERED_KV_STAGE") != NULL; + const bool supported_shape = + comp_is_f16 && head_dim == 512u && raw_cap <= UINT32_MAX / 128u; + const uint64_t row_vecs64 = 128u; + const uint64_t comp_count64 = (uint64_t)n_comp * head_dim; + const uint64_t dst_comp_inner64 = + (uint64_t)n_raw * head_dim * sizeof(uint16_t); + if (comp_count64 > UINT32_MAX || + dst_comp_inner64 > NSUIntegerMax - dst_offset) { + return 0; + } + const uint64_t total_vecs64 = + ((uint64_t)n_raw + n_comp) * row_vecs64; + const bool valid_grid = + total_vecs64 != 0 && total_vecs64 <= UINT32_MAX; + const bool eligible = + supported_shape && valid_grid && !g_quality_mode && !disabled && + g_flash_kv_stage_f16_pipeline != nil && + (ds4_gpu_device_name_contains("M3") || force); + const bool component_disabled = eligible && + (ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F32_F16_COPY") > 0 || + ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F16_F16_COPY") > 0); + const bool use_fusion = eligible && !component_disabled; + const bool use_pad_fusion = + use_fusion && fuse_pad && mask != nil && pad != nil && + getenv("DS4_METAL_DISABLE_M3_GATHERED_KV_PAD_FUSION") == NULL; + if (require && supported_shape && !use_fusion) { + fprintf(stderr, + "ds4: required Metal gathered KV staging kernel was not selected\n"); + return 0; + } + + if (use_fusion) { + ds4_gpu_flash_kv_stage_f16_args args = { + .raw_cap = raw_cap, + .raw_start = raw_start, + .n_raw = n_raw, + .n_comp = n_comp, + .pad_rows = use_pad_fusion ? 32u : 0u, + .shared_pad = use_pad_fusion && shared_pad ? 1u : 0u, + }; + const NSUInteger total_vecs = (NSUInteger)total_vecs64 + + (use_pad_fusion ? 32u * 128u + 32u : 0u); + NSUInteger nth = + g_flash_kv_stage_f16_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth > total_vecs) nth = total_vecs; + if (nth == 0) return 0; + const NSUInteger groups = (total_vecs + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_flash_kv_stage_f16_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:raw offset:raw_offset atIndex:1]; + [enc setBuffer:comp offset:comp_offset atIndex:2]; + [enc setBuffer:dst offset:dst_offset atIndex:3]; + [enc setBuffer:(use_pad_fusion ? mask : dst) + offset:(use_pad_fusion ? mask_offset : dst_offset) + atIndex:4]; + [enc setBuffer:(use_pad_fusion ? pad : dst) + offset:(use_pad_fusion ? pad_offset : dst_offset) + atIndex:5]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + if (did_fuse_pad) *did_fuse_pad = use_pad_fusion; + return 1; + } + + if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb, + raw, + raw_offset, + raw_cap, + raw_start, + n_raw, + head_dim, + dst, + dst_offset)) { + return 0; + } + return ds4_gpu_encode_copy_to_f16_1d( + cb, + comp, + comp_offset, + comp_is_f16, + dst, + dst_offset + (NSUInteger)dst_comp_inner64, + (uint32_t)comp_count64); +} + +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) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!dst || !raw || !comp || raw_cap == 0 || raw_start >= raw_cap || + n_raw == 0 || n_raw > raw_cap || n_comp == 0 || + comp_is_f16 == 0 || head_dim != 512u) { + return 0; + } + + @autoreleasepool { + const uint64_t raw_bytes = + (uint64_t)raw_cap * head_dim * sizeof(float); + const uint64_t comp_bytes = + (uint64_t)n_comp * head_dim * + (comp_is_f16 ? sizeof(uint16_t) : sizeof(float)); + const uint64_t dst_bytes = + ((uint64_t)n_raw + n_comp) * head_dim * sizeof(uint16_t); + id rawbuf = ds4_gpu_tensor_buffer(raw); + id compbuf = ds4_gpu_tensor_buffer(comp); + id dstbuf = ds4_gpu_tensor_buffer(dst); + if (!rawbuf || !compbuf || !dstbuf || + ds4_gpu_tensor_bytes(raw) < raw_bytes || + ds4_gpu_tensor_bytes(comp) < comp_bytes || + ds4_gpu_tensor_bytes(dst) < dst_bytes) { + fprintf(stderr, + "ds4: Metal gathered KV staging received undersized buffers\n"); + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + if (!ds4_gpu_encode_flash_kv_stage_f16( + cb, + rawbuf, + ds4_gpu_tensor_offset(raw), + raw_cap, + raw_start, + n_raw, + compbuf, + ds4_gpu_tensor_offset(comp), + comp_is_f16 != 0, + n_comp, + head_dim, + dstbuf, + ds4_gpu_tensor_offset(dst), + nil, + 0, + nil, + 0, + false, + false, + NULL)) { + return 0; + } + if (!ds4_gpu_finish_command_buffer( + cb, owned, "gathered KV staging")) { + return 0; + } + } return 1; } @@ -16860,6 +18427,7 @@ static int ds4_gpu_encode_flash_attention_raw_heads( id vec_pipeline = ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512", true, true, false, false, (n_raw % ncpsg) != 0, + false, (int32_t)head_dim, (int32_t)head_dim, (int32_t)nsg, @@ -16868,49 +18436,15 @@ static int ds4_gpu_encode_flash_attention_raw_heads( ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, (int32_t)nwg); if (!vec_pipeline || !reduce_pipeline) return 0; - id kvbuf = rawbuf; - NSUInteger kvoff = ds4_gpu_tensor_offset(raw_kv); - if (raw_start != 0) { - const NSUInteger ring_bytes = (NSUInteger)n_raw * row_bytes; - const uint32_t tail_avail = raw_cap - raw_start; - const uint32_t tail_rows = tail_avail < n_raw ? tail_avail : n_raw; - const uint32_t head_rows = n_raw - tail_rows; - const uint32_t tail_elems = tail_rows * head_dim; - const uint32_t head_elems = head_rows * head_dim; - if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_ring_buffer, - &g_flash_attn_ring_bytes, - ring_bytes, - "ds4_flash_attn_ring")) { - return 0; - } - - if ((tail_rows && - !ds4_gpu_encode_cpy_f32_f32_1d(cb, - rawbuf, - ds4_gpu_tensor_offset(raw_kv) + (NSUInteger)raw_start * row_bytes, - g_flash_attn_ring_buffer, - 0, - tail_elems)) || - (head_rows && - !ds4_gpu_encode_cpy_f32_f32_1d(cb, + if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb, rawbuf, ds4_gpu_tensor_offset(raw_kv), - g_flash_attn_ring_buffer, - (NSUInteger)tail_rows * row_bytes, - head_elems))) { - return 0; - } - - kvbuf = g_flash_attn_ring_buffer; - kvoff = 0; - } - - if (!ds4_gpu_encode_cpy_f32_f16_1d(cb, - kvbuf, - kvoff, - g_flash_attn_kv_buffer, - 0, - n_raw * head_dim)) { + raw_cap, + raw_start, + n_raw, + head_dim, + g_flash_attn_kv_buffer, + 0)) { return 0; } @@ -17162,7 +18696,37 @@ static int ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec_long const NSUInteger nblk1 = ((NSUInteger)n_tokens + nqptg - 1u) / nqptg; const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u); - id mask_buffer = ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); + const uint32_t mask_cache_kind = ratio == 4u + ? DS4_GPU_PREFILL_MASK_CACHE_RATIO4 + : (ratio == 128u ? DS4_GPU_PREFILL_MASK_CACHE_RATIO128 : 0u); + bool mask_cache_created = false; + ds4_gpu_zero_prefix_prefill_mask_cache_entry *mask_cache = + use_comp_mask == 0u && mask_cache_kind != 0u + ? ds4_gpu_get_zero_prefix_prefill_mask_cache(mask_cache_kind, + n_tokens, + n_comp, + n_keys, + window, + ratio, + nqptg, + ncpsg, + has_kvpad, + bc_mask, + mask_bytes, + blk_bytes, + &mask_cache_created) + : NULL; + id mask_buffer = mask_cache + ? mask_cache->mask + : ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); + if (mask_cache && mask_cache_created) { + ds4_gpu_fill_static_mixed_prefill_mask((uint16_t *)[mask_buffer contents], + n_tokens, + n_comp, + window, + ratio); + mask_cache->valid = true; + } if (!mask_buffer || !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, &g_flash_attn_kv_bytes, @@ -17172,12 +18736,16 @@ static int ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec_long &g_flash_attn_pad_bytes, pad_bytes, "ds4_flash_attn_pad") || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer, - &g_flash_attn_blk_bytes, - blk_bytes, - "ds4_flash_attn_blk")) { + (!mask_cache && + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer, + &g_flash_attn_blk_bytes, + blk_bytes, + "ds4_flash_attn_blk"))) { return 0; } + id blk_buffer = mask_cache + ? mask_cache->blk + : g_flash_attn_blk_buffer; const bool flash_stage_profile = getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL && g_batch_cb != nil; @@ -17226,11 +18794,13 @@ static int ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec_long DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_comp"); } - ds4_gpu_fill_static_mixed_prefill_mask((uint16_t *)[mask_buffer contents], - n_tokens, - n_comp, - window, - ratio); + if (!mask_cache) { + ds4_gpu_fill_static_mixed_prefill_mask((uint16_t *)[mask_buffer contents], + n_tokens, + n_comp, + window, + ratio); + } DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_fill"); if (use_comp_mask && n_comp != 0) { if (!ds4_gpu_encode_cpy_f32_f16_2d(cb, @@ -17306,14 +18876,17 @@ static int ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec_long }; id enc = nil; - enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:blk_pipeline]; - [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0]; - [enc setBuffer:mask_buffer offset:0 atIndex:1]; - [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1) - threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); + if (!mask_cache || !mask_cache->blk_ready) { + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:blk_pipeline]; + [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0]; + [enc setBuffer:mask_buffer offset:0 atIndex:1]; + [enc setBuffer:blk_buffer offset:0 atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + if (mask_cache) mask_cache->blk_ready = true; + } DS4_METAL_PROFILE_FLASH_ATTN_STAGE("block_map"); ds4_gpu_flash_attn_vec_args args = { @@ -17365,7 +18938,7 @@ static int ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec_long [enc setBuffer:mask_buffer offset:0 atIndex:4]; [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; - [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:7]; + [enc setBuffer:blk_buffer offset:0 atIndex:7]; [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8]; [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; [enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1) @@ -17530,6 +19103,7 @@ static int ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_vec( id vec_pipeline = ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512", true, true, false, false, has_kvpad, + false, (int32_t)head_dim, (int32_t)head_dim, (int32_t)nsg, @@ -17742,8 +19316,29 @@ static int ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec( const NSUInteger nblk1 = ((NSUInteger)n_tokens + nqptg - 1u) / nqptg; const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u); - id mask_buffer = - ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); + bool mask_cache_created = false; + ds4_gpu_zero_prefix_prefill_mask_cache_entry *mask_cache = + ds4_gpu_get_zero_prefix_prefill_mask_cache( + DS4_GPU_PREFILL_MASK_CACHE_RAW, + n_tokens, + 0u, + n_tokens, + window, + 0u, + nqptg, + ncpsg, + has_kvpad, + bc_mask, + mask_bytes, + blk_bytes, + &mask_cache_created); + id mask_buffer = mask_cache + ? mask_cache->mask + : ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); + if (mask_cache && mask_cache_created) { + ds4_gpu_fill_raw_prefill_mask((uint16_t *)[mask_buffer contents], n_tokens, window); + mask_cache->valid = true; + } if (!mask_buffer || !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, &g_flash_attn_kv_bytes, @@ -17753,12 +19348,16 @@ static int ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec( &g_flash_attn_pad_bytes, pad_bytes, "ds4_flash_attn_pad") || - !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer, - &g_flash_attn_blk_bytes, - blk_bytes, - "ds4_flash_attn_blk")) { + (!mask_cache && + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer, + &g_flash_attn_blk_bytes, + blk_bytes, + "ds4_flash_attn_blk"))) { return 0; } + id blk_buffer = mask_cache + ? mask_cache->blk + : g_flash_attn_blk_buffer; const bool flash_stage_profile = getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL && g_batch_cb != nil; @@ -17784,7 +19383,9 @@ static int ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec( } \ } while (0) - ds4_gpu_fill_raw_prefill_mask((uint16_t *)[mask_buffer contents], n_tokens, window); + if (!mask_cache) { + ds4_gpu_fill_raw_prefill_mask((uint16_t *)[mask_buffer contents], n_tokens, window); + } DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_fill"); id pad_pipeline = nil; @@ -17855,14 +19456,18 @@ static int ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec( .nb33 = mask_bytes, }; - id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:blk_pipeline]; - [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0]; - [enc setBuffer:mask_buffer offset:0 atIndex:1]; - [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:2]; - [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1) - threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; - ds4_gpu_end_compute_encoder(cb, enc); + id enc = nil; + if (!mask_cache || !mask_cache->blk_ready) { + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:blk_pipeline]; + [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0]; + [enc setBuffer:mask_buffer offset:0 atIndex:1]; + [enc setBuffer:blk_buffer offset:0 atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + if (mask_cache) mask_cache->blk_ready = true; + } DS4_METAL_PROFILE_FLASH_ATTN_STAGE("block_map"); ds4_gpu_flash_attn_vec_args args = { @@ -17914,7 +19519,7 @@ static int ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec( [enc setBuffer:mask_buffer offset:0 atIndex:4]; [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; - [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:7]; + [enc setBuffer:blk_buffer offset:0 atIndex:7]; [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8]; [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; [enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1) @@ -18035,6 +19640,7 @@ static int ds4_gpu_encode_flash_attention_prefill_raw_heads( id vec_pipeline = ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512", true, true, false, false, true, + false, (int32_t)head_dim, (int32_t)head_dim, (int32_t)nsg, @@ -18216,10 +19822,18 @@ static int ds4_gpu_encode_flash_attention_gathered_heads( const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) + nrows * (2u * (NSUInteger)nwg) * sizeof(float); - if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_mask_buffer, - &g_flash_attn_mask_bytes, - mask_bytes, - "ds4_flash_attn_mask") || + const bool use_persistent_zero_mask = + use_mask == 0u && + (ds4_gpu_device_name_contains("M3") || + getenv("DS4_METAL_ENABLE_PERSISTENT_ZERO_ATTN_MASK") != NULL) && + getenv("DS4_METAL_DISABLE_M3_PERSISTENT_ZERO_ATTN_MASK") == NULL; + + if (!(use_persistent_zero_mask + ? ds4_gpu_ensure_zero_attention_mask(mask_bytes) + : ds4_gpu_ensure_scratch_buffer(&g_flash_attn_mask_buffer, + &g_flash_attn_mask_bytes, + mask_bytes, + "ds4_flash_attn_mask")) || !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, &g_flash_attn_kv_bytes, kv_bytes, @@ -18234,15 +19848,20 @@ static int ds4_gpu_encode_flash_attention_gathered_heads( "ds4_flash_attn_tmp")) { return 0; } + id flash_mask_buffer = use_persistent_zero_mask + ? g_flash_attn_zero_mask_buffer + : g_flash_attn_mask_buffer; - id pad_pipeline = nil; - if ((n_keys % ncpsg) != 0) { - pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); - if (!pad_pipeline) return 0; - } + const bool has_kvpad = (n_keys % ncpsg) != 0; + const bool use_shared_kvpad = + has_kvpad && + (ds4_gpu_device_name_contains("M3") || + getenv("DS4_METAL_ENABLE_SHARED_KV_PAD") != NULL) && + getenv("DS4_METAL_DISABLE_M3_SHARED_KV_PAD") == NULL; id vec_pipeline = ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512", - true, true, false, false, (n_keys % ncpsg) != 0, + true, true, false, false, has_kvpad, + use_shared_kvpad, (int32_t)head_dim, (int32_t)head_dim, (int32_t)nsg, @@ -18251,74 +19870,52 @@ static int ds4_gpu_encode_flash_attention_gathered_heads( ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, (int32_t)nwg); if (!vec_pipeline || !reduce_pipeline) return 0; - id raw_linear_buf = rawbuf; - NSUInteger raw_linear_offset = ds4_gpu_tensor_offset(raw_kv); - if (raw_start != 0) { - const NSUInteger ring_bytes = (NSUInteger)n_raw * row_bytes; - const uint32_t tail_rows = raw_cap - raw_start < n_raw ? raw_cap - raw_start : n_raw; - const uint32_t head_rows = n_raw - tail_rows; - if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_ring_buffer, - &g_flash_attn_ring_bytes, - ring_bytes, - "ds4_flash_attn_ring")) { - return 0; - } - - if ((tail_rows && - !ds4_gpu_encode_cpy_f32_f32_1d(cb, - rawbuf, - ds4_gpu_tensor_offset(raw_kv) + (NSUInteger)raw_start * row_bytes, - g_flash_attn_ring_buffer, - 0, - tail_rows * head_dim)) || - (head_rows && - !ds4_gpu_encode_cpy_f32_f32_1d(cb, - rawbuf, - ds4_gpu_tensor_offset(raw_kv), - g_flash_attn_ring_buffer, - (NSUInteger)tail_rows * row_bytes, - head_rows * head_dim))) { - return 0; - } - - raw_linear_buf = g_flash_attn_ring_buffer; - raw_linear_offset = 0; - } - - if (!ds4_gpu_encode_cpy_f32_f16_1d(cb, - raw_linear_buf, - raw_linear_offset, - g_flash_attn_kv_buffer, - 0, - n_raw * head_dim)) { - return 0; - } - if (n_comp) { - if (!ds4_gpu_encode_copy_to_f16_1d(cb, - compbuf, - ds4_gpu_tensor_offset(comp_kv), - comp_kv_f16 != 0, - g_flash_attn_kv_buffer, - (NSUInteger)n_raw * row_bytes_f16, - n_comp * head_dim)) { - return 0; - } - } - - if (!ds4_gpu_encode_fill_f16_1d(cb, g_flash_attn_mask_buffer, 0, n_keys, 0.0f)) { + if (!use_persistent_zero_mask && + !ds4_gpu_encode_fill_f16_1d(cb, flash_mask_buffer, 0, n_keys, 0.0f)) { return 0; } if (use_mask && n_comp && !ds4_gpu_encode_cpy_f32_f16_1d(cb, maskbuf, ds4_gpu_tensor_offset(comp_mask), - g_flash_attn_mask_buffer, + flash_mask_buffer, (NSUInteger)n_raw * sizeof(uint16_t), n_comp)) { return 0; } - if ((n_keys % ncpsg) != 0) { + bool pad_fused = false; + if (!ds4_gpu_encode_flash_kv_stage_f16( + cb, + rawbuf, + ds4_gpu_tensor_offset(raw_kv), + raw_cap, + raw_start, + n_raw, + compbuf, + ds4_gpu_tensor_offset(comp_kv), + comp_kv_f16 != 0, + n_comp, + head_dim, + g_flash_attn_kv_buffer, + 0, + flash_mask_buffer, + 0, + g_flash_attn_pad_buffer, + 0, + has_kvpad, + use_shared_kvpad, + &pad_fused)) { + return 0; + } + + id pad_pipeline = nil; + if (has_kvpad && !pad_fused) { + pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); + if (!pad_pipeline) return 0; + } + + if (has_kvpad && !pad_fused) { ds4_gpu_flash_attn_pad_args pad_args = { .ne11 = (int32_t)n_keys, .ne_12_2 = 1, @@ -18342,7 +19939,7 @@ static int ds4_gpu_encode_flash_attention_gathered_heads( [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1]; [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; - [enc setBuffer:g_flash_attn_mask_buffer offset:0 atIndex:3]; + [enc setBuffer:flash_mask_buffer offset:0 atIndex:3]; [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; @@ -18395,7 +19992,7 @@ static int ds4_gpu_encode_flash_attention_gathered_heads( [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; - [enc setBuffer:g_flash_attn_mask_buffer offset:0 atIndex:4]; + [enc setBuffer:flash_mask_buffer offset:0 atIndex:4]; [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:7]; @@ -18486,45 +20083,15 @@ static int ds4_gpu_encode_flash_attention_decode_raw_batch_heads( return 0; } - id kvbuf = rawbuf; - NSUInteger kvoff = ds4_gpu_tensor_offset(raw_kv); - if (raw_start != 0) { - const NSUInteger ring_bytes = (NSUInteger)n_raw * row_bytes; - const uint32_t tail_avail = raw_cap - raw_start; - const uint32_t tail_rows = tail_avail < n_raw ? tail_avail : n_raw; - const uint32_t head_rows = n_raw - tail_rows; - if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_ring_buffer, - &g_flash_attn_ring_bytes, - ring_bytes, - "ds4_flash_attn_ring")) { - return 0; - } - if ((tail_rows && - !ds4_gpu_encode_cpy_f32_f32_1d(cb, - rawbuf, - ds4_gpu_tensor_offset(raw_kv) + (NSUInteger)raw_start * row_bytes, - g_flash_attn_ring_buffer, - 0, - tail_rows * head_dim)) || - (head_rows && - !ds4_gpu_encode_cpy_f32_f32_1d(cb, + if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb, rawbuf, ds4_gpu_tensor_offset(raw_kv), - g_flash_attn_ring_buffer, - (NSUInteger)tail_rows * row_bytes, - head_rows * head_dim))) { - return 0; - } - kvbuf = g_flash_attn_ring_buffer; - kvoff = 0; - } - - if (!ds4_gpu_encode_cpy_f32_f16_1d(cb, - kvbuf, - kvoff, - g_flash_attn_kv_buffer, - 0, - n_raw * head_dim)) { + raw_cap, + raw_start, + n_raw, + head_dim, + g_flash_attn_kv_buffer, + 0)) { return 0; } @@ -18757,45 +20324,15 @@ static int ds4_gpu_encode_flash_attention_decode_mixed_batch_heads( return 0; } - id kvbuf = rawbuf; - NSUInteger kvoff = ds4_gpu_tensor_offset(raw_kv); - if (raw_start != 0) { - const NSUInteger ring_bytes = (NSUInteger)n_raw * row_bytes; - const uint32_t tail_avail = raw_cap - raw_start; - const uint32_t tail_rows = tail_avail < n_raw ? tail_avail : n_raw; - const uint32_t head_rows = n_raw - tail_rows; - if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_ring_buffer, - &g_flash_attn_ring_bytes, - ring_bytes, - "ds4_flash_attn_ring")) { - return 0; - } - if ((tail_rows && - !ds4_gpu_encode_cpy_f32_f32_1d(cb, - rawbuf, - ds4_gpu_tensor_offset(raw_kv) + (NSUInteger)raw_start * row_bytes, - g_flash_attn_ring_buffer, - 0, - tail_rows * head_dim)) || - (head_rows && - !ds4_gpu_encode_cpy_f32_f32_1d(cb, + if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb, rawbuf, ds4_gpu_tensor_offset(raw_kv), - g_flash_attn_ring_buffer, - (NSUInteger)tail_rows * row_bytes, - head_rows * head_dim))) { - return 0; - } - kvbuf = g_flash_attn_ring_buffer; - kvoff = 0; - } - - if (!ds4_gpu_encode_cpy_f32_f16_1d(cb, - kvbuf, - kvoff, - g_flash_attn_kv_buffer, - 0, - n_raw * head_dim) || + raw_cap, + raw_start, + n_raw, + head_dim, + g_flash_attn_kv_buffer, + 0) || !ds4_gpu_encode_copy_to_f16_1d(cb, compbuf, ds4_gpu_tensor_offset(comp_kv), @@ -21492,7 +23029,7 @@ static int ds4_gpu_encode_attn_out_low_q8_mpp( [enc setBuffer:src0 offset:src0_off atIndex:1]; [enc setBuffer:src1 offset:src1_off atIndex:2]; [enc setBuffer:dst offset:dst_off atIndex:3]; - [enc setThreadgroupMemoryLength:4096u atIndex:0]; + [enc setThreadgroupMemoryLength:8192u atIndex:0]; [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)mm_args->ne21 + (NSUInteger)tile_n - 1u) / (NSUInteger)tile_n, ((NSUInteger)mm_args->ne0 + 63u) / 64u, (NSUInteger)mm_args->ne02) @@ -21907,16 +23444,45 @@ static int ds4_gpu_encode_router_select( if (flash_router_fast_path && !g_quality_mode && n_tokens == 1 && getenv("DS4_METAL_DISABLE_ROUTER_SELECT_FUSION") == NULL) { + const bool force_simd_weights_fusion = + getenv("DS4_METAL_ENABLE_ROUTER_SIMD_WEIGHTS_FUSION") != NULL; + const bool use_simd_finalize = + !hash_mode && + g_dsv4_router_finalize_one_simd_pipeline != nil && + g_dsv4_router_finalize_one_simd_pipeline.threadExecutionWidth == 32u && + g_dsv4_router_finalize_one_simd_pipeline.maxTotalThreadsPerThreadgroup >= 256u && + (ds4_gpu_device_name_contains("M3") || + getenv("DS4_METAL_ENABLE_ROUTER_SIMD_FINALIZE") != NULL || + force_simd_weights_fusion) && + getenv("DS4_METAL_DISABLE_M3_ROUTER_SIMD_FINALIZE") == NULL; + const bool use_simd_weights_fusion = + use_simd_finalize && + g_dsv4_router_finalize_weights_one_simd_pipeline != nil && + g_dsv4_router_finalize_weights_one_simd_pipeline.threadExecutionWidth == 32u && + g_dsv4_router_finalize_weights_one_simd_pipeline.maxTotalThreadsPerThreadgroup >= 256u && + (ds4_gpu_device_name_contains("M3") || force_simd_weights_fusion) && + getenv("DS4_METAL_DISABLE_M3_ROUTER_SIMD_WEIGHTS_FUSION") == NULL; id softplus_sqrt_pipeline = ds4_gpu_hot_pipeline(g_dsv4_softplus_sqrt_pipeline, "kernel_dsv4_softplus_sqrt_f32_4"); id router_finalize_pipeline = - ds4_gpu_hot_pipeline(g_dsv4_router_finalize_one_pipeline, - "kernel_dsv4_router_finalize_one"); - id router_weights_pipeline = - ds4_gpu_hot_pipeline(g_dsv4_router_weights_one_pipeline, - "kernel_dsv4_router_weights_one"); - if (!softplus_sqrt_pipeline || !router_finalize_pipeline || !router_weights_pipeline) return 0; + ds4_gpu_hot_pipeline( + use_simd_weights_fusion + ? g_dsv4_router_finalize_weights_one_simd_pipeline + : use_simd_finalize + ? g_dsv4_router_finalize_one_simd_pipeline + : g_dsv4_router_finalize_one_pipeline, + use_simd_weights_fusion + ? "kernel_dsv4_router_finalize_weights_one_simd" + : use_simd_finalize + ? "kernel_dsv4_router_finalize_one_simd" + : "kernel_dsv4_router_finalize_one"); + id router_weights_pipeline = use_simd_weights_fusion + ? nil + : ds4_gpu_hot_pipeline(g_dsv4_router_weights_one_pipeline, + "kernel_dsv4_router_weights_one"); + if (!softplus_sqrt_pipeline || !router_finalize_pipeline || + (!use_simd_weights_fusion && !router_weights_pipeline)) return 0; ok = ds4_gpu_encode_unary_f32_rows(cb, softplus_sqrt_pipeline, @@ -21968,11 +23534,19 @@ static int ds4_gpu_encode_router_select( [enc setBytes:&zero_i32 length:sizeof(zero_i32) atIndex:4]; } [enc setBuffer:selectedbuf offset:selected_off atIndex:5]; - [enc setThreadgroupMemoryLength:256u * sizeof(float) + 256u * sizeof(int32_t) atIndex:0]; + if (use_simd_weights_fusion) { + [enc setBuffer:weightsbuf offset:weights_off atIndex:6]; + } + const NSUInteger router_finalize_scratch_bytes = use_simd_finalize + ? 2u * (256u * sizeof(float) + 256u * sizeof(int32_t)) + : 256u * sizeof(float) + 256u * sizeof(int32_t); + [enc setThreadgroupMemoryLength:router_finalize_scratch_bytes atIndex:0]; [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; ds4_gpu_end_compute_encoder(cb, enc); + if (use_simd_weights_fusion) return 1; + enc = ds4_gpu_compute_encoder(cb); [enc setComputePipelineState:router_weights_pipeline]; [enc setBuffer:probsbuf offset:probs_off atIndex:0]; @@ -21984,14 +23558,6 @@ static int ds4_gpu_encode_router_select( return 1; } - const NSUInteger sum_bytes = (NSUInteger)n_tokens * sizeof(float); - if (!ds4_gpu_ensure_scratch_buffer(&g_router_weight_sum_buffer, - &g_router_weight_sum_bytes, - sum_bytes, - "ds4_router_weight_sum")) { - return 0; - } - if (flash_router_fast_path && !g_quality_mode && n_tokens == 1) { id softplus_sqrt_pipeline = ds4_gpu_hot_pipeline(g_dsv4_softplus_sqrt_pipeline, @@ -22083,6 +23649,28 @@ static int ds4_gpu_encode_router_select( } if (!ok) return 0; + const bool use_batch_weights_fusion = + flash_router_fast_path && !g_quality_mode && n_tokens > 1u && + g_dsv4_router_weights_batch_pipeline != nil && + (ds4_gpu_device_name_contains("M3") || + getenv("DS4_METAL_ENABLE_ROUTER_WEIGHTS_BATCH_FUSION") != NULL) && + getenv("DS4_METAL_DISABLE_M3_ROUTER_WEIGHTS_BATCH_FUSION") == NULL && + getenv("DS4_METAL_DISABLE_ROUTER_SELECT_FUSION") == NULL; + if (use_batch_weights_fusion) { + const float scale = expert_weight_scale; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_router_weights_batch_pipeline]; + [enc setBytes:&scale length:sizeof(scale) atIndex:0]; + [enc setBuffer:probsbuf offset:probs_off atIndex:1]; + [enc setBuffer:selectedbuf offset:selected_off atIndex:2]; + [enc setBuffer:weightsbuf offset:weights_off atIndex:3]; + [enc setThreadgroupMemoryLength:40u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake(6, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; + } + if (flash_router_fast_path && !g_quality_mode && n_tokens == 1) { id router_weights_pipeline = ds4_gpu_hot_pipeline(g_dsv4_router_weights_one_pipeline, @@ -22099,6 +23687,14 @@ static int ds4_gpu_encode_router_select( return 1; } + const NSUInteger sum_bytes = (NSUInteger)n_tokens * sizeof(float); + if (!ds4_gpu_ensure_scratch_buffer(&g_router_weight_sum_buffer, + &g_router_weight_sum_bytes, + sum_bytes, + "ds4_router_weight_sum")) { + return 0; + } + ok = ds4_gpu_encode_get_rows_f32_router_weights(cb, probsbuf, probs_off, @@ -24934,13 +26530,21 @@ int ds4_gpu_routed_moe_batch_tensor( !g_quality_mode && !use_q4_batch_expert_table && !use_iq2_batch_selected_addr; + /* + * Fused gate+up grouped matmul with the SwiGLU epilogue. The IQ2 + * variant stays opt-in behind its env flag; the Q4_K variant is the + * default path — same MMA accumulation order and epilogue math as the + * separate GEMMs + swiglu pass, so the mid tensor is bit-identical. + */ const bool use_mm_id_pair_swiglu = use_mm_id && request_mid_f16 && - gate_type == DS4_METAL_TENSOR_IQ2_XXS && - down_type == DS4_METAL_TENSOR_Q2_K && n_expert == 6 && - getenv("DS4_METAL_ENABLE_MOE_MM_ID_PAIR_SWIGLU") != NULL && + ((gate_type == DS4_METAL_TENSOR_IQ2_XXS && + down_type == DS4_METAL_TENSOR_Q2_K && + getenv("DS4_METAL_ENABLE_MOE_MM_ID_PAIR_SWIGLU") != NULL) || + (gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K)) && getenv("DS4_METAL_DISABLE_MOE_MM_ID_PAIR_SWIGLU") == NULL && getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL && getenv("DS4_METAL_GRAPH_DUMP_PREFIX") == NULL; @@ -24965,7 +26569,9 @@ int ds4_gpu_routed_moe_batch_tensor( ds4_gpu_routed_mm_pipeline(down_type); if (use_mm_id_pair_swiglu) { pair_swiglu_mm_pipeline = - ds4_gpu_get_pipeline("kernel_mul_mm_id_iq2_xxs_pair_swiglu_f16"); + ds4_gpu_get_pipeline(gate_type == DS4_METAL_TENSOR_Q4_K ? + "kernel_mul_mm_id_q4_K_pair_swiglu_f16" : + "kernel_mul_mm_id_iq2_xxs_pair_swiglu_f16"); } if (!map_pipeline || !gate_mm_pipeline || !up_mm_pipeline || !down_mm_pipeline || (use_mm_id_pair_swiglu && !pair_swiglu_mm_pipeline)) { @@ -25822,6 +27428,123 @@ int ds4_gpu_hc_weighted_sum_tensor( "HC weighted sum"); } +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) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !norm_out || !residual_hc || !weights || !model_map) return 0; + + const bool force = + getenv("DS4_METAL_ENABLE_OUTPUT_HC_SUM_NORM_FUSION") != NULL; + const bool disabled = + getenv("DS4_METAL_DISABLE_M3_OUTPUT_HC_SUM_NORM_FUSION") != NULL; + const bool require = + getenv("DS4_METAL_REQUIRE_OUTPUT_HC_SUM_NORM_FUSION") != NULL; + const bool supported_shape = + n_hc == 4u && (n_embd == 4096u || n_embd == 7168u); + const bool auto_shape = + n_embd == 4096u && ds4_gpu_device_name_contains("M3"); + const bool use_fusion = + supported_shape && !g_quality_mode && !disabled && + g_hc_weighted_sum_norm_pipeline != nil && + (auto_shape || force); + if (require && supported_shape && !use_fusion) { + fprintf(stderr, + "ds4: required Metal output HC sum/RMSNorm fusion was not selected\n"); + return 0; + } + if (!use_fusion) return 0; + + @autoreleasepool { + const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float); + const uint64_t residual_bytes = (uint64_t)n_hc * out_row_bytes; + const uint64_t weight_bytes = (uint64_t)n_hc * sizeof(float); + id xbuf = ds4_gpu_tensor_buffer(residual_hc); + id wbuf = ds4_gpu_tensor_buffer(weights); + id outbuf = ds4_gpu_tensor_buffer(out); + id normbuf = ds4_gpu_tensor_buffer(norm_out); + if (!xbuf || !wbuf || !outbuf || !normbuf || + ds4_gpu_tensor_bytes(residual_hc) < residual_bytes || + ds4_gpu_tensor_bytes(weights) < weight_bytes || + ds4_gpu_tensor_bytes(out) != out_row_bytes || + ds4_gpu_tensor_bytes(norm_out) < out_row_bytes) { + fprintf(stderr, + "ds4: Metal output HC sum/RMSNorm fusion received invalid activation buffers\n"); + return 0; + } + if (norm_weight_offset > model_size || + out_row_bytes > model_size - norm_weight_offset) { + fprintf(stderr, + "ds4: Metal output HC sum/RMSNorm weight range is outside the mapped model\n"); + return 0; + } + + uint64_t norm_inner = 0; + id normwbuf = ds4_gpu_wrap_model_range( + model_map, model_size, norm_weight_offset, + out_row_bytes, &norm_inner); + if (!normwbuf) return 0; + + ds4_gpu_hc_weighted_sum_norm_args args = { + .n_embd = (int64_t)n_embd, + .n_hc = (int64_t)n_hc, + .n_tokens = 1, + .nb_x0 = sizeof(float), + .nb_x1 = out_row_bytes, + .nb_x2 = residual_bytes, + .nb_w0 = sizeof(float), + .nb_w1 = weight_bytes, + .nb0 = sizeof(float), + .nb1 = out_row_bytes, + .nb_norm1 = out_row_bytes, + .norm_eps = norm_eps, + }; + const NSUInteger nth = ds4_gpu_rms_norm_threads(n_embd); + const NSUInteger shared_bytes = + ((NSUInteger)n_embd + 32u) * sizeof(float); + if (nth > g_hc_weighted_sum_norm_pipeline.maxTotalThreadsPerThreadgroup || + shared_bytes > [g_device maxThreadgroupMemoryLength]) { + if (require) { + fprintf(stderr, + "ds4: required Metal output HC sum/RMSNorm fusion exceeds device limits\n"); + } + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_hc_weighted_sum_norm_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:1]; + [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setBuffer:normwbuf offset:(NSUInteger)norm_inner atIndex:4]; + [enc setBuffer:normbuf offset:ds4_gpu_tensor_offset(norm_out) atIndex:5]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer( + cb, owned, "output HC sum/RMSNorm fused")) { + return 0; + } + } + + return 1; +} + int ds4_gpu_hc_weighted_sum_split_tensor( ds4_gpu_tensor *out, const ds4_gpu_tensor *residual_hc, @@ -25957,9 +27680,9 @@ int ds4_gpu_hc_split_weighted_sum_tensor( return 1; } -/* Decode-only HC-pre plus the immediately following weighted RMSNorm. This is - * specialized for DS4's HC=4 decode shape; larger batched prefill keeps using - * the existing two-stage path. */ +/* HC-pre plus the immediately following weighted RMSNorm, specialized for + * DS4's HC=4 shape. Both decode and batched prefill use this implementation; + * the kernel preserves their established single-row and batched scale formulas. */ int ds4_gpu_hc_split_weighted_sum_norm_tensor( ds4_gpu_tensor *out, ds4_gpu_tensor *norm_out, @@ -26154,6 +27877,51 @@ int ds4_gpu_output_hc_weights_tensor( &base_inner); if (!scalebuf || !basebuf) return 0; + const bool force_weights4 = + getenv("DS4_METAL_ENABLE_OUTPUT_HC_WEIGHTS4") != NULL; + const bool disable_weights4 = + getenv("DS4_METAL_DISABLE_M3_OUTPUT_HC_WEIGHTS4") != NULL; + const bool require_weights4 = + getenv("DS4_METAL_REQUIRE_OUTPUT_HC_WEIGHTS4") != NULL; + const bool weights4_shape = n_hc == 4u && n_tokens64 == 1u; + const bool use_weights4 = + weights4_shape && !g_quality_mode && !disable_weights4 && + g_output_hc_weights4_pipeline != nil && + (ds4_gpu_device_name_contains("M3") || force_weights4) && + g_output_hc_weights4_pipeline.maxTotalThreadsPerThreadgroup >= 2u; + if (require_weights4 && weights4_shape && !use_weights4) { + fprintf(stderr, + "ds4: required Metal output HC weights4 kernel was not selected\n"); + return 0; + } + + if (use_weights4) { + ds4_gpu_output_hc_weights4_args args = { + .post_scale = 1.0f, + .eps = eps, + }; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_output_hc_weights4_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:prebuf offset:ds4_gpu_tensor_offset(pre) atIndex:1]; + [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2]; + [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(2, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer( + cb, owned, "output HC weights4")) { + return 0; + } + return 1; + } + const uint32_t n_tokens = (uint32_t)n_tokens64; ds4_gpu_bin_args mul_args = ds4_gpu_make_bin_rows_args(n_hc, n_tokens, 1); ds4_gpu_bin_args add_args = ds4_gpu_make_bin_rows_args(n_hc, n_tokens, n_hc); diff --git a/metal/argsort.metal b/metal/argsort.metal index 5081022bc..77c473f96 100644 --- a/metal/argsort.metal +++ b/metal/argsort.metal @@ -65,6 +65,15 @@ kernel void kernel_argsort_f32_i32( // initialize indices shmem_i32[col] = i00 + col; + // Stage this block's score slice in threadgroup memory (indices stay in + // [i00, i00+ntg.x), so shmem_f32[idx - i00] replaces the device gather). + // The host allocates ntg.x extra floats after the index array. Values and + // the comparison network are unchanged, so the permutation is identical. + threadgroup float * shmem_f32 = (threadgroup float *) (shmem_i32 + ntg.x); + if (i00 + col < args.ne00) { + shmem_f32[col] = src0_row[i00 + col]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); for (int k = 2; k <= ntg.x; k *= 2) { @@ -74,16 +83,16 @@ kernel void kernel_argsort_f32_i32( if ((col & k) == 0) { if (shmem_i32[col] >= args.ne00 || (shmem_i32[ixj] < args.ne00 && (order == DS4_SORT_ORDER_ASC ? - src0_row[shmem_i32[col]] > src0_row[shmem_i32[ixj]] : - src0_row[shmem_i32[col]] < src0_row[shmem_i32[ixj]])) + shmem_f32[shmem_i32[col] - i00] > shmem_f32[shmem_i32[ixj] - i00] : + shmem_f32[shmem_i32[col] - i00] < shmem_f32[shmem_i32[ixj] - i00])) ) { SWAP(shmem_i32[col], shmem_i32[ixj]); } } else { if (shmem_i32[ixj] >= args.ne00 || (shmem_i32[col] < args.ne00 && (order == DS4_SORT_ORDER_ASC ? - src0_row[shmem_i32[col]] < src0_row[shmem_i32[ixj]] : - src0_row[shmem_i32[col]] > src0_row[shmem_i32[ixj]])) + shmem_f32[shmem_i32[col] - i00] < shmem_f32[shmem_i32[ixj] - i00] : + shmem_f32[shmem_i32[col] - i00] > shmem_f32[shmem_i32[ixj] - i00])) ) { SWAP(shmem_i32[col], shmem_i32[ixj]); } diff --git a/metal/cpy.metal b/metal/cpy.metal index 3aa00ac19..f606d185b 100644 --- a/metal/cpy.metal +++ b/metal/cpy.metal @@ -55,3 +55,192 @@ typedef decltype(kernel_cpy_t_t) kernel_cpy_t; template [[host_name("kernel_cpy_f32_f32")]] kernel kernel_cpy_t kernel_cpy_t_t; template [[host_name("kernel_cpy_f32_f16")]] kernel kernel_cpy_t kernel_cpy_t_t; template [[host_name("kernel_cpy_f16_f32")]] kernel kernel_cpy_t kernel_cpy_t_t; + +// Contiguous 1D conversions avoid the generic tensor-index reconstruction +// above. Packed vector types retain scalar alignment, so tensor views whose +// offsets are float/half aligned do not need additional 16/8-byte alignment. +// The final vector is converted element-by-element when n is not divisible by +// four, preserving the generic kernel's bounds and conversion semantics. +kernel void kernel_cpy_contig_f32_f16_4( + constant uint & n, + device const packed_float4 * src, + device packed_half4 * dst, + uint gid [[thread_position_in_grid]]) { + const uint i = gid * 4u; + if (i >= n) { + return; + } + + const uint remaining = n - i; + if (remaining >= 4u) { + const float4 value = float4(src[gid]); + dst[gid] = packed_half4(half4(value)); + return; + } + + device const float * src_scalar = (device const float *)src; + device half * dst_scalar = (device half *)dst; + for (uint lane = 0; lane < remaining; ++lane) { + dst_scalar[i + lane] = half(src_scalar[i + lane]); + } +} + +kernel void kernel_cpy_contig_f16_f32_4( + constant uint & n, + device const packed_half4 * src, + device packed_float4 * dst, + uint gid [[thread_position_in_grid]]) { + const uint i = gid * 4u; + if (i >= n) { + return; + } + + const uint remaining = n - i; + if (remaining >= 4u) { + const half4 value = half4(src[gid]); + dst[gid] = packed_float4(float4(value)); + return; + } + + device const half * src_scalar = (device const half *)src; + device float * dst_scalar = (device float *)dst; + for (uint lane = 0; lane < remaining; ++lane) { + dst_scalar[i + lane] = float(src_scalar[i + lane]); + } +} + +// Bitwise F16 transport for cache staging. Use ushort rather than half so NaN +// payloads and every other binary16 encoding pass through unchanged. +kernel void kernel_cpy_contig_f16_f16_bits_4( + constant uint & n, + device const packed_ushort4 * src, + device packed_ushort4 * dst, + uint gid [[thread_position_in_grid]]) { + const uint i = gid * 4u; + if (i >= n) { + return; + } + + const uint remaining = n - i; + if (remaining >= 4u) { + dst[gid] = src[gid]; + return; + } + + device const ushort * src_scalar = (device const ushort *)src; + device ushort * dst_scalar = (device ushort *)dst; + for (uint lane = 0; lane < remaining; ++lane) { + dst_scalar[i + lane] = src_scalar[i + lane]; + } +} + +struct ds4_metal_args_flash_kv_stage_f16 { + uint raw_cap; + uint raw_start; + uint n_raw; + uint n_comp; + uint pad_rows; + uint shared_pad; +}; + +// Decode-time gathered attention consumes a logical raw-cache ring followed +// by an already-F16 compressed cache. Pack both regions into the contiguous +// F16 FlashAttention scratch in one dispatch. The raw conversion expression +// and compressed ushort4 transport exactly match the standalone copy kernels. +kernel void kernel_dsv4_flash_kv_stage_f16( + constant ds4_metal_args_flash_kv_stage_f16 & args, + device const char * raw_src, + device const char * comp_src, + device char * dst, + device const char * mask_src, + device char * pad_dst, + uint gid [[thread_position_in_grid]]) { + constexpr uint row_vecs = 128; + const uint raw_vecs = args.n_raw * row_vecs; + const uint n_keys = args.n_raw + args.n_comp; + const uint total_vecs = n_keys * row_vecs; + + if (gid < raw_vecs) { + const uint logical_row = gid >> 7; + const uint col = gid & 127u; + uint physical_row = args.raw_start + logical_row; + if (physical_row >= args.raw_cap) { + physical_row -= args.raw_cap; + } + device const packed_float4 *raw = + (device const packed_float4 *)raw_src; + device packed_half4 *dst_half = (device packed_half4 *)dst; + const float4 value = + float4(raw[physical_row * row_vecs + col]); + dst_half[gid] = packed_half4(half4(value)); + return; + } + + if (gid < total_vecs) { + device const packed_ushort4 *comp = + (device const packed_ushort4 *)comp_src; + device packed_ushort4 *dst_bits = (device packed_ushort4 *)dst; + dst_bits[gid] = comp[gid - raw_vecs]; + return; + } + + // The vector FlashAttention kernel redirects its final partial block to + // a compact K/V/mask buffer. When requested, append those writes to this + // dispatch so gathered decode does not need a standalone pad dispatch. + const uint pad_rows = args.pad_rows; + const uint pad_vecs = pad_rows * row_vecs; + uint pad_gid = gid - total_vecs; + if (pad_gid < pad_vecs) { + const uint row = pad_gid >> 7; + const uint col = pad_gid & 127u; + const uint valid_rows = n_keys % pad_rows; + device packed_half4 *pad_half = (device packed_half4 *)pad_dst; + device packed_ushort4 *pad_bits = (device packed_ushort4 *)pad_dst; + if (row >= valid_rows) { + const packed_half4 zero = packed_half4(half4(0.0h)); + pad_half[pad_gid] = zero; + if (!args.shared_pad) { + pad_half[pad_vecs + pad_gid] = zero; + } + return; + } + + const uint logical_row = n_keys - valid_rows + row; + if (logical_row < args.n_raw) { + uint physical_row = args.raw_start + logical_row; + if (physical_row >= args.raw_cap) { + physical_row -= args.raw_cap; + } + device const packed_float4 *raw = + (device const packed_float4 *)raw_src; + const float4 value = + float4(raw[physical_row * row_vecs + col]); + const packed_half4 value_half = packed_half4(half4(value)); + pad_half[pad_gid] = value_half; + if (!args.shared_pad) { + pad_half[pad_vecs + pad_gid] = value_half; + } + } else { + device const packed_ushort4 *comp = + (device const packed_ushort4 *)comp_src; + const packed_ushort4 value_bits = + comp[(logical_row - args.n_raw) * row_vecs + col]; + pad_bits[pad_gid] = value_bits; + if (!args.shared_pad) { + pad_bits[pad_vecs + pad_gid] = value_bits; + } + } + return; + } + + pad_gid -= pad_vecs; + if (pad_gid < pad_rows) { + const uint valid_rows = n_keys % pad_rows; + device const ushort *mask_bits = (device const ushort *)mask_src; + device ushort *pad_mask_bits = + (device ushort *)pad_dst + 2u * pad_vecs * 4u; + pad_mask_bits[pad_gid] = pad_gid < valid_rows + ? mask_bits[n_keys - valid_rows + pad_gid] + : 0xfbffu; + } +} diff --git a/metal/dense.metal b/metal/dense.metal index ad1dd8790..f3913484f 100644 --- a/metal/dense.metal +++ b/metal/dense.metal @@ -25,6 +25,13 @@ struct ds4_metal_args_mul_mv { short r3; }; +struct ds4_metal_args_compressor_pair_store { + uint32_t width; + uint32_t ratio; + uint32_t pos; + uint32_t ape_type; +}; + struct ds4_metal_args_mul_mm { int32_t ne00; int32_t ne02; @@ -190,6 +197,143 @@ kernel void kernel_mul_mv_q8_0_f32( kernel_mul_mv_q8_0_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); } +// The single-token vocabulary projection has enough output rows that grouping +// four independent rows per threadgroup halves scheduling and activation +// rereads. Each row retains the exact Q8 block walk and reduction tree above. +[[host_name("kernel_mul_mv_q8_0_f32_nr4")]] +kernel void kernel_mul_mv_q8_0_f32_nr4( + constant ds4_metal_args_mul_mv & args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + kernel_mul_mv_q8_0_f32_impl<4, constant ds4_metal_args_mul_mv &>( + args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); +} + +// Decode Q-A/KV pair. Both projections consume the same activation row but +// have independent weight ranges and output extents. Keep the standalone Q8_0 +// lane/block traversal and two-stage reduction verbatim for each bank; only +// the activation load and threadgroup scheduling are shared. +[[host_name("kernel_mul_mv_q8_0_f32_pair")]] +kernel void kernel_mul_mv_q8_0_f32_pair( + constant ds4_metal_args_mul_mv & args0, + constant ds4_metal_args_mul_mv & args1, + device const char * src0_a, + device const char * src0_b, + device const char * src1, + device char * dst_a, + device char * dst_b, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; + constexpr short NW = N_SIMDWIDTH; + constexpr short NQ = 8; + constexpr short NR0 = N_R0_Q8_0; + + const int nb = args0.ne00 / QK8_0; + const int r0 = tgpig.x * NR0; + const bool active_a = r0 < args0.ne01; + const bool active_b = r0 < args1.ne01; + + device const float *y = (device const float *)src1; + device const block_q8_0 *ax_a[NR0]; + device const block_q8_0 *ax_b[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + const int out_row = r0 + row; + ax_a[row] = active_a && out_row < args0.ne01 + ? (device const block_q8_0 *)(src0_a + (uint64_t)out_row * args0.nb01) + : (device const block_q8_0 *)src0_a; + ax_b[row] = active_b && out_row < args1.ne01 + ? (device const block_q8_0 *)(src0_b + (uint64_t)out_row * args1.nb01) + : (device const block_q8_0 *)src0_b; + } + + float suma[NR0] = { 0.f }; + float sumb[NR0] = { 0.f }; + + const short ix = tiisg / (NW / NQ); + const short il = tiisg % (NW / NQ); + const int ib0 = sgitg * NQ + ix; + float yl[NQ]; + device const float *yb = y + ib0 * QK8_0 + il * NQ; + + for (int ib = ib0; ib < nb; ib += NSG * NQ) { + FOR_UNROLL (short i = 0; i < NQ; ++i) { + yl[i] = yb[i]; + } + + FOR_UNROLL (short row = 0; row < NR0; ++row) { + const int out_row = r0 + row; + if (active_a && out_row < args0.ne01) { + device const int8_t *qs = ax_a[row][ib].qs + il * NQ; + float sumq = 0.f; + FOR_UNROLL (short i = 0; i < NQ; ++i) { + sumq += qs[i] * yl[i]; + } + suma[row] += sumq * ax_a[row][ib].d; + } + if (active_b && out_row < args1.ne01) { + device const int8_t *qs = ax_b[row][ib].qs + il * NQ; + float sumq = 0.f; + FOR_UNROLL (short i = 0; i < NQ; ++i) { + sumq += qs[i] * yl[i]; + } + sumb[row] += sumq * ax_b[row][ib].d; + } + } + + yb += NSG * NQ * QK8_0; + } + + threadgroup float *shared = (threadgroup float *)shmem; + threadgroup float *sha[NR0]; + threadgroup float *shb[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + sha[row] = shared + NW * row; + shb[row] = shared + NW * (NR0 + row); + if (sgitg == 0) { + sha[row][tiisg] = 0.0f; + if (active_b) shb[row][tiisg] = 0.0f; + } + suma[row] = simd_sum(suma[row]); + if (active_b) sumb[row] = simd_sum(sumb[row]); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + FOR_UNROLL (short row = 0; row < NR0; ++row) { + if (tiisg == 0) { + sha[row][sgitg] = suma[row]; + if (active_b) shb[row][sgitg] = sumb[row]; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + device float *out_a = (device float *)dst_a; + device float *out_b = (device float *)dst_b; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + const float total_a = simd_sum(sha[row][tiisg]); + if (tiisg == 0 && sgitg == 0) { + const int out_row = r0 + row; + if (active_a && out_row < args0.ne01) out_a[out_row] = total_a; + } + if (active_b) { + const float total_b = simd_sum(shb[row][tiisg]); + if (tiisg == 0 && sgitg == 0) { + const int out_row = r0 + row; + if (out_row < args1.ne01) out_b[out_row] = total_b; + } + } + } +} + // Decode shared-expert gate/up projections followed by SwiGLU: // // mid = silu(min(gate, limit)) * clamp(up, -limit, limit) @@ -689,6 +833,59 @@ kernel void kernel_mul_mv_f16_f32_pair_4( args, src0_a, src0_b, src1, dst_a, dst_b, shmem, tgpig, tiisg, sgitg); } +// Decode compressor projection plus recurrent-state append. The paired +// matvec remains unchanged and still materializes both F32 outputs. After a +// device-memory barrier, the first NR0 threads reload those exact stored bits +// and perform the same state write and score+APE addition as +// kernel_dsv4_compressor_store_one. +kernel void kernel_mul_mv_f16_f32_pair_compressor_store_4( + constant ds4_metal_args_mul_mv & args, + constant ds4_metal_args_compressor_pair_store & store, + device const char * src0_a, + device const char * src0_b, + device const char * src1, + device char * dst_a, + device char * dst_b, + device const char * ape, + device float * state_kv, + device float * state_score, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + kernel_mul_mv_f16_f32_pair_4_disp( + args, src0_a, src0_b, src1, dst_a, dst_b, + shmem, tgpig, tiisg, sgitg); + + threadgroup_barrier(mem_flags::mem_device); + + if (tiitg >= args.nr0 || store.width == 0u || store.ratio == 0u) { + return; + } + const uint col = tgpig.x * (uint)args.nr0 + tiitg; + if (col >= store.width) return; + + const uint pos_mod = store.pos % store.ratio; + const uint dst_row = store.ratio == 4u ? store.ratio + pos_mod : pos_mod; + const uint dst = dst_row * store.width + col; + const uint ape_i = pos_mod * store.width + col; + + device volatile const float * projected_kv = + (device volatile const float *)dst_a; + device volatile const float * projected_score = + (device volatile const float *)dst_b; + float ape_v; + if (store.ape_type == 1u) { + ape_v = (float)(((device const half *)ape)[ape_i]); + } else { + ape_v = ((device const float *)ape)[ape_i]; + } + + state_kv[dst] = projected_kv[col]; + state_score[dst] = projected_score[col] + ape_v; +} + template void kernel_mul_mv_t_t_short_impl( args_t args, @@ -776,6 +973,27 @@ void dequantize_q8_0(device const block_q8_0 *xb, short il, thread type4x4 & reg reg = (type4x4) reg_f; } +/* + * Bit-identical twin of dequantize_q8_0 for the MPP staging loop: same + * half(float(qs[i]) * d) per element, but the 16 consecutive int8 lanes are + * fetched as eight aligned 16-bit loads instead of sixteen byte loads. + * xb->qs is always 2-byte aligned (block_q8_0.d is a half). + */ +void dequantize_q8_0_pairs(device const block_q8_0 *xb, short il, thread half4x4 & reg) { + device const ushort *qs16 = (device const ushort *)(xb->qs + 16*il); + const float d = xb->d; + + float4x4 reg_f; + + FOR_UNROLL (short i = 0; i < 8; i++) { + const ushort u = qs16[i]; + reg_f[i/2][(i%2)*2 + 0] = ((float)(int8_t)(u & 0xFF)) * d; + reg_f[i/2][(i%2)*2 + 1] = ((float)(int8_t)(u >> 8)) * d; + } + + reg = (half4x4) reg_f; +} + template void dequantize_q8_0_t4(device const block_q8_0 *xb, short il, thread type4 & reg) { device const int8_t * qs = ((device const int8_t *)xb->qs); @@ -1044,6 +1262,11 @@ template [[host_name("kernel_mul_mm_f16_f32_mpp")]] kernel mul_mm_mpp_t kernel_ // aligned F16/Q8_0 prompt matmuls. The host selects the widest token tile that // evenly divides the batch, with 128-token tiles retained after the 64-token // retest was neutral or slower. +// +// The host dispatch guarantees M % NR0 == 0 and K % NK == 0, so the dequant +// stage does no bounds work. The weight tile is double-buffered: the next +// k-step's dequant overlaps the current cooperative matmul, so the k-loop +// needs one threadgroup barrier per step instead of two. template< short NR1, typename SA, typename SA_4x4, typename block_q, short nl, @@ -1077,7 +1300,8 @@ kernel void kernel_mul_mm_mpp_direct_rhs( const uint64_t offset0 = (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; threadgroup SA *sa = (threadgroup SA *)shmem; - auto tA = tensor(sa, dextents(NK, NR0)); + auto tA0 = tensor(sa, dextents(NK, NR0)); + auto tA1 = tensor(sa + NR0*NK, dextents(NK, NR0)); device T1 *ptrB = (device T1 *)(srcB + args.nb12*i12 + args.nb13*i13); const int strideB = args.nb11/sizeof(T1); @@ -1088,7 +1312,7 @@ kernel void kernel_mul_mm_mpp_direct_rhs( matmul2d_descriptor::mode::multiply_accumulate), execution_simdgroups<4>> mm; - auto cT = mm.template get_destination_cooperative_tensor(); + auto cT = mm.template get_destination_cooperative_tensor(); #pragma unroll for (uint16_t i = 0; i < cT.get_capacity(); ++i) { @@ -1097,43 +1321,48 @@ kernel void kernel_mul_mm_mpp_direct_rhs( } } - for (int loop_k = 0; loop_k < K; loop_k += NK) { - for (int work = tiitg; work < NR0*NL; work += NUM_THREADS) { - const int row = work/NL; - const int k_chunk = work%NL; + // NR0*NL/NUM_THREADS 16-value weight chunks per thread (1 at NR0=64). + auto stage_tile = [&](const int loop_k, threadgroup SA *buf) { + FOR_UNROLL (int work = tiitg; work < NR0*NL; work += NUM_THREADS) { + const int row = work / NL; + const int k_chunk = work % NL; const int k_pos = loop_k + k_chunk*16; const short k_base = k_chunk*16; - - if (r0 + row < M) { - if (is_same::value && FC_mul_mm_bc_inp) { - device const T0 *row_ptr = (device const T0 *)(srcA + args.nb01*(r0 + row) + offset0); - FOR_UNROLL (short i = 0; i < 16; i++) { - sa[row*NK + k_base + i] = (k_pos + i < K) ? (SA)row_ptr[k_pos + i] : (SA)0; - } - } else { - const int block_idx = k_pos/(16*nl); - const short il = (k_pos/16)%nl; - device const block_q *row_ptr = (device const block_q *)(srcA + args.nb01*(r0 + row) + offset0); - - SA_4x4 temp_a; - dequantize_func(row_ptr + block_idx, il, temp_a); - FOR_UNROLL (short i = 0; i < 16; i++) { - sa[row*NK + k_base + i] = (k_pos + i < K) ? temp_a[i/4][i%4] : (SA)0; - } - } - } else { + if (is_same::value && FC_mul_mm_bc_inp) { + device const T0 *row_ptr_f = + (device const T0 *)(srcA + args.nb01*(r0 + row) + offset0); FOR_UNROLL (short i = 0; i < 16; i++) { - sa[row*NK + k_base + i] = (SA)0; + buf[row*NK + k_base + i] = (SA)row_ptr_f[k_pos + i]; } + } else { + device const block_q *row_ptr = + (device const block_q *)(srcA + args.nb01*(r0 + row) + offset0); + SA_4x4 temp_a; + dequantize_func(row_ptr + k_pos/(16*nl), (k_pos/16)%nl, temp_a); + typedef vec SA4; + threadgroup SA4 *dst4 = (threadgroup SA4 *)(buf + row*NK + k_base); + dst4[0] = temp_a[0]; + dst4[1] = temp_a[1]; + dst4[2] = temp_a[2]; + dst4[3] = temp_a[3]; } } + }; - threadgroup_barrier(mem_flags::mem_threadgroup); + stage_tile(0, sa); + threadgroup_barrier(mem_flags::mem_threadgroup); - auto mA = tA.slice(0, 0); + uint buf_sel = 0; + for (int loop_k = 0; loop_k < K; loop_k += NK) { + auto mA = (buf_sel ? tA1 : tA0).slice(0, 0); auto mB = tB.slice(loop_k, r1); mm.run(mB, mA, cT); + const int next_k = loop_k + NK; + if (next_k < K) { + buf_sel ^= 1u; + stage_tile(next_k, buf_sel ? sa + NR0*NK : sa); + } threadgroup_barrier(mem_flags::mem_threadgroup); } @@ -1148,9 +1377,9 @@ typedef decltype(kernel_mul_mm_mpp_direct_rhs<32, half, half4x4, float4x4, 1, de template [[host_name("kernel_mul_mm_f16_f32_mpp_direct_rhs")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<32, half, half4x4, half4x4, 1, dequantize_f16, half, half4x4, float>; template [[host_name("kernel_mul_mm_f16_f32_mpp_direct_rhs_n64")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<64, half, half4x4, half4x4, 1, dequantize_f16, half, half4x4, float>; template [[host_name("kernel_mul_mm_f16_f32_mpp_direct_rhs_n128")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<128, half, half4x4, half4x4, 1, dequantize_f16, half, half4x4, float>; -template [[host_name("kernel_mul_mm_q8_0_f32_nax_direct_rhs")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<32, half, half4x4, block_q8_0, 2, dequantize_q8_0, float, float4x4, float>; -template [[host_name("kernel_mul_mm_q8_0_f32_nax_direct_rhs_n64")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<64, half, half4x4, block_q8_0, 2, dequantize_q8_0, float, float4x4, float>; -template [[host_name("kernel_mul_mm_q8_0_f32_nax_direct_rhs_n128")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<128, half, half4x4, block_q8_0, 2, dequantize_q8_0, float, float4x4, float>; +template [[host_name("kernel_mul_mm_q8_0_f32_nax_direct_rhs")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<32, half, half4x4, block_q8_0, 2, dequantize_q8_0_pairs, float, float4x4, float>; +template [[host_name("kernel_mul_mm_q8_0_f32_nax_direct_rhs_n64")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<64, half, half4x4, block_q8_0, 2, dequantize_q8_0_pairs, float, float4x4, float>; +template [[host_name("kernel_mul_mm_q8_0_f32_nax_direct_rhs_n128")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<128, half, half4x4, block_q8_0, 2, dequantize_q8_0_pairs, float, float4x4, float>; #endif // Tiled matrix-matrix kernel used for prompt batches larger than 8. DS4 uses @@ -1357,6 +1586,217 @@ kernel void kernel_mul_mm( } } +// Legacy F16-weight/F32-RHS prefill matmul with a per-row RMS scale applied at +// the existing F32-to-F16 RHS staging boundary. The tile layout, half inputs, +// float accumulators, and output path intentionally mirror kernel_mul_mm so the +// only arithmetic change versus materializing RMSNorm first is where x*scale is +// rounded from F32 to F16. +kernel void kernel_mul_mm_f16_f32_scaled( + constant ds4_metal_args_mul_mm & args, + device const char * src0, + device const char * src1, + device char * dst, + device const float * scales, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + + threadgroup half * sa = (threadgroup half *)(shmem); + threadgroup half * sb = (threadgroup half *)(shmem + 4096); + + constexpr int NR0 = 64; + constexpr int NR1 = 32; + + constexpr int NK = 32; + constexpr int NL0 = NK/16; + constexpr int NL1 = NK/8; + + const int im = tgpig.z; + const int r0 = tgpig.y*NR0; + const int r1 = tgpig.x*NR1; + + // if this block is of 64x32 shape or smaller + const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0; + const short nr1 = (args.ne1 - r1 < NR1) ? (args.ne1 - r1) : NR1; + + // a thread shouldn't load data outside of the matrix + const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1; // 0 .. 63 + const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1; // 0 .. 31 + + const short il0 = (tiitg % NL0); + + short il = il0; + + const int i12 = im%args.ne12; + const int i13 = im/args.ne12; + + const uint64_t offset0 = (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; + const short offset1 = il0; + + device const half4x4 * x = (device const half4x4 *)(src0 + args.nb01*(r0 + lr0) + offset0) + offset1; + + const short iy = 8*(tiitg % NL1); + + device const float * y = (device const float *)(src1 + + args.nb13*i13 + + args.nb12*i12 + + args.nb11*(r1 + lr1) + + args.nb10*iy); + const float row_scale = scales[r1 + lr1]; + + simdgroup_half8x8 ma[4]; + simdgroup_half8x8 mb[2]; + + simdgroup_float8x8 mc[8]; + + for (short i = 0; i < 8; i++){ + mc[i] = make_filled_simdgroup_matrix(0.f); + } + + for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) { + // load data and store to threadgroup memory + if (FC_mul_mm_bc_inp) { + threadgroup_barrier(mem_flags::mem_threadgroup); + + // no need for dequantization + for (short i = 0; i < 16; i++) { + const short sx = 2*il0 + i/8; + const short sy = (tiitg/NL0)/8; + + const short lx = (tiitg/NL0)%8; + const short ly = i%8; + + const short ib = 8*sx + sy; + + *(sa + 64*ib + 8*ly + lx) = loop_k + 16*il + i < args.ne00 ? *((device half *) x + i) : 0; + } + } else { + half4x4 temp_a; + dequantize_f16(x, il, temp_a); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + FOR_UNROLL (short i = 0; i < 16; i++) { + const short sx = 2*il0 + i/8; + const short sy = (tiitg/NL0)/8; + + const short lx = (tiitg/NL0)%8; + const short ly = i%8; + + const short ib = 8*sx + sy; + + // Pointer-form store matches the legacy F16 tile layout. + *(sa + 64*ib + 8*ly + lx) = temp_a[i/4][i%4]; + } + } + + if (FC_mul_mm_bc_inp) { + for (short i = 0; i < 8; ++i) { + const short sx = (tiitg%NL1); + const short sy = (tiitg/NL1)/8; + + const short lx = i; + const short ly = (tiitg/NL1)%8; + + const short ib = 4*sx + sy; + const float scaled = loop_k + iy + i < args.ne00 ? y[i] * row_scale : 0.0f; + + *(sb + 64*ib + 8*ly + lx) = (half)scaled; + } + } else { + const short sx = (tiitg%NL1); + const short sy = (tiitg/NL1)/8; + + const short ly = (tiitg/NL1)%8; + + const short ib = 4*sx + sy; + + const float2x4 raw = *((device const float2x4 *) y); + float2x4 scaled; + scaled[0] = raw[0] * row_scale; + scaled[1] = raw[1] * row_scale; + *(threadgroup half2x4 *)(sb + 64*ib + 8*ly) = (half2x4)scaled; + } + + il = (il + 2 < 1) ? il + 2 : il % 2; + x = (il < 2) ? x + 2 : x; + + y += NK; + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // load matrices from threadgroup memory and conduct outer products + threadgroup const half * lsma = (sa + 4*64*(sgitg%2)); + threadgroup const half * lsmb = (sb + 2*64*(sgitg/2)); + + FOR_UNROLL (short ik = 0; ik < NK/8; ik++) { + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 4; i++) { + simdgroup_load(ma[i], lsma + 64*i, 8, 0, false); + } + + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 2; i++) { + simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false); + } + + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 8; i++){ + simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]); + } + + lsma += 8*64; + lsmb += 4*64; + } + } + + if (!FC_mul_mm_bc_out || (r0 + NR0 <= args.ne0 && r1 + NR1 <= args.ne1)) { + // if no bounds checks on the output are needed, we can directly write to device memory + device float * C = (device float *) dst + + (r0 + 32*(sgitg & 1)) + \ + (r1 + 16*(sgitg >> 1)) * args.ne0 + im*args.ne1*args.ne0; + + for (short i = 0; i < 8; i++) { + simdgroup_store(mc[i], C + 8*(i%4) + 8*args.ne0*(i/4), args.ne0, 0, false); + } + } else { + // block is smaller than 64x32, we should avoid writing data outside of the matrix + threadgroup_barrier(mem_flags::mem_threadgroup); + + threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0; + + for (short i = 0; i < 8; i++) { + simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (sgitg == 0) { + for (int j = tiitg; j < nr1; j += NR1) { + device float * D = (device float *) dst + r0 + (r1 + j)*args.ne0 + im*args.ne1*args.ne0; + device float4 * D4 = (device float4 *) D; + + threadgroup float * C = temp_str + (j*NR0); + threadgroup float4 * C4 = (threadgroup float4 *) C; + + int i = 0; + for (; i < nr0/4; i++) { + *(D4 + i) = *(C4 + i); + } + + i *= 4; + for (; i < nr0; i++) { + *(D + i) = *(C + i); + } + } + } + } +} + kernel void kernel_mul_mm_f16_f32_pair( constant ds4_metal_args_mul_mm & args, device const char * src0_a, diff --git a/metal/dsv4_hc.metal b/metal/dsv4_hc.metal index fb67011eb..af5c0e98c 100644 --- a/metal/dsv4_hc.metal +++ b/metal/dsv4_hc.metal @@ -21,6 +21,26 @@ struct ds4_metal_args_dsv4_hc_weighted_sum { uint64_t nb1; }; +struct ds4_metal_args_dsv4_hc_weighted_sum_norm { + int64_t n_embd; + int64_t n_hc; + int64_t n_tokens; + uint64_t nb_x0; + uint64_t nb_x1; + uint64_t nb_x2; + uint64_t nb_w0; + uint64_t nb_w1; + uint64_t nb0; + uint64_t nb1; + uint64_t nb_norm1; + float norm_eps; +}; + +struct ds4_metal_args_dsv4_output_hc_weights4 { + float post_scale; + float eps; +}; + struct ds4_metal_args_dsv4_hc_split_weighted_sum { int64_t n_embd; int32_t n_hc; @@ -506,10 +526,12 @@ kernel void kernel_dsv4_hc_split_weighted_sum_norm4( device const float4 *x1 = (device const float4 *)(x + 1 * args.nb_x1 + (uint64_t)row * args.nb_x2); device const float4 *x2 = (device const float4 *)(x + 2 * args.nb_x1 + (uint64_t)row * args.nb_x2); device const float4 *x3 = (device const float4 *)(x + 3 * args.nb_x1 + (uint64_t)row * args.nb_x2); - const float4 v = x0[i] * pre_shmem[0] + - x1[i] * pre_shmem[1] + - x2[i] * pre_shmem[2] + - x3[i] * pre_shmem[3]; + // Preserve the standalone HC collapse's explicit accumulation order. + float4 v = 0.0f; + v += x0[i] * pre_shmem[0]; + v += x1[i] * pre_shmem[1]; + v += x2[i] * pre_shmem[2]; + v += x3[i] * pre_shmem[3]; row_shmem[i] = v; sumf += dot(v, v); } @@ -523,7 +545,11 @@ kernel void kernel_dsv4_hc_split_weighted_sum_norm4( sumf = sum_shmem[tiisg]; sumf = simd_sum(sumf); - const float norm_scale = rsqrt(sumf / float(n_embd) + args.norm_eps); + // Batched prefill must match kernel_rms_norm_fuse_impl so enabling this + // fusion does not change its scale by an ULP. Keep the established + // single-row decode result unchanged; that path historically used rsqrt. + const float norm_arg = sumf / float(n_embd) + args.norm_eps; + const float norm_scale = args.n_rows > 1 ? 1.0f / sqrt(norm_arg) : rsqrt(norm_arg); device float4 *dst4 = (device float4 *)(dst + (uint64_t)row * args.nb1); device const float4 *w4 = (device const float4 *)norm_weight; @@ -883,3 +909,109 @@ kernel void kernel_dsv4_hc_weighted_sum( *((device float *) (dst + d*args.nb0 + t*args.nb1)) = acc; } + +// The one-row output head immediately applies a learned RMSNorm after reducing +// its four HC streams. Preserve the standalone scalar HC accumulation, write +// the collapsed row for diagnostics, then reload its F32 values from +// threadgroup memory using the standalone RMSNorm's float4 reduction mapping. +kernel void kernel_dsv4_hc_weighted_sum_norm4( + constant ds4_metal_args_dsv4_hc_weighted_sum_norm & args, + device const char * x, + device const char * weights, + device char * dst, + device const char * norm_weight, + device char * norm_dst, + threadgroup float * shared [[threadgroup(0)]], + ushort tid [[thread_position_in_threadgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort ntg [[threads_per_threadgroup]]) { + if (args.n_tokens != 1 || args.n_hc != 4 || + args.n_embd <= 0 || (args.n_embd & 3) != 0) { + return; + } + + const uint n_embd = uint(args.n_embd); + const uint n4 = n_embd >> 2; + threadgroup float *row_shmem = shared; + threadgroup float *sum_shmem = shared + n_embd; + + if (sgitg == 0) { + sum_shmem[tiisg] = 0.0f; + } + + for (uint d = tid; d < n_embd; d += ntg) { + float acc = 0.0f; + for (int64_t h = 0; h < args.n_hc; ++h) { + const float xv = *((device const float *)( + x + (uint64_t)d*args.nb_x0 + (uint64_t)h*args.nb_x1)); + const float wv = *((device const float *)( + weights + (uint64_t)h*args.nb_w0)); + acc += xv * wv; + } + row_shmem[d] = acc; + *((device float *)(dst + (uint64_t)d*args.nb0)) = acc; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + threadgroup const float4 *row4 = + (threadgroup const float4 *)row_shmem; + float sumf = 0.0f; + for (uint i = tid; i < n4; i += ntg) { + sumf += dot(row4[i], row4[i]); + } + sumf = simd_sum(sumf); + + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tiisg == 0) { + sum_shmem[sgitg] = sumf; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + sumf = sum_shmem[tiisg]; + sumf = simd_sum(sumf); + + const float mean = sumf/args.n_embd; + const float scale = 1.0f/sqrt(mean + args.norm_eps); + device const float4 *w4 = (device const float4 *)norm_weight; + device float4 *norm4 = (device float4 *)norm_dst; + for (uint i = tid; i < n4; i += ntg) { + norm4[i] = (row4[i]*scale)*w4[i]; + } +} + +// The one-row HC=4 output head historically materializes four device-F32 +// stages across separate launches. Collapse those launches into one tiny +// two-thread group while preserving the scalar/vector lane mapping and every +// global rounding boundary. +kernel void kernel_dsv4_output_hc_weights4( + constant ds4_metal_args_dsv4_output_hc_weights4 & args, + device const float * pre, + device const float * hc_scale, + device const float * hc_base, + device float * dst, + ushort tid [[thread_position_in_threadgroup]]) { + device volatile float *stage = (device volatile float *)dst; + + for (uint i = tid; i < 4; i += 2) { + stage[i] = pre[i] * hc_scale[0]; + } + threadgroup_barrier(mem_flags::mem_device); + + for (uint i = tid; i < 4; i += 2) { + stage[i] = stage[i] + hc_base[i]; + } + threadgroup_barrier(mem_flags::mem_device); + + if (tid == 0) { + const float4 x = *((device volatile float4 *)stage); + *((device volatile float4 *)stage) = 1 / (1 + exp(-x)); + } + threadgroup_barrier(mem_flags::mem_device); + + if (tid == 0) { + const float4 x = *((device volatile float4 *)stage); + *((device volatile float4 *)stage) = + args.post_scale * x + args.eps; + } +} diff --git a/metal/dsv4_kv.metal b/metal/dsv4_kv.metal index f91bdbf46..4f44a43ea 100644 --- a/metal/dsv4_kv.metal +++ b/metal/dsv4_kv.metal @@ -41,6 +41,13 @@ struct ds4_metal_args_dsv4_ratio4_shift { uint32_t width; }; +struct ds4_metal_args_dsv4_compressor_pack_ratio4 { + uint32_t head_dim; + uint32_t n_comp; + uint32_t replay; + uint32_t n_threads; +}; + struct ds4_metal_args_dsv4_compressor_store_one { uint32_t width; uint32_t ratio; @@ -264,6 +271,54 @@ kernel void kernel_dsv4_kv_fp8_store_f32( } } +// Builds the two ratio-4 softmax-pool packs together. Each output plane holds +// four previous-half rows followed by four current-half rows. Normal prefill +// seeds plane zero with 0/-inf; replay seeds it from the previous compressor +// state. Later planes take their previous half from the prior input group. +kernel void kernel_dsv4_compressor_pack_ratio4( + constant ds4_metal_args_dsv4_compressor_pack_ratio4 & args, + device const uint * kv, + device const uint * score, + device const uint * state_kv, + device const uint * state_score, + device uint * packed_kv, + device uint * packed_score, + uint2 group [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + if (group.x >= args.n_comp || group.y >= 8u || + args.head_dim == 0u || args.n_threads == 0u) { + return; + } + + const uint plane = group.x; + const uint row = group.y; + const uint64_t input_row_stride = 2ull * args.head_dim; + const uint64_t dst_row = ((uint64_t)plane * 8u + row) * args.head_dim; + + for (uint col = tid; col < args.head_dim; col += args.n_threads) { + const uint64_t dst = dst_row + col; + if (row >= 4u) { + const uint token = plane * 4u + (row - 4u); + const uint64_t src = (uint64_t)token * input_row_stride + + args.head_dim + col; + packed_kv[dst] = kv[src]; + packed_score[dst] = score[src]; + } else if (plane != 0u) { + const uint token = (plane - 1u) * 4u + row; + const uint64_t src = (uint64_t)token * input_row_stride + col; + packed_kv[dst] = kv[src]; + packed_score[dst] = score[src]; + } else if (args.replay != 0u) { + const uint64_t src = (uint64_t)row * input_row_stride + col; + packed_kv[dst] = state_kv[src]; + packed_score[dst] = state_score[src]; + } else { + packed_kv[dst] = 0u; + packed_score[dst] = 0xff800000u; + } + } +} + // Ratio-4 compression keeps two 4-row halves of recurrent state. After an // emitted compressed row, the second half becomes the next window's previous // half. The old encoder expressed this as four generic copies; this DS4-specific diff --git a/metal/dsv4_misc.metal b/metal/dsv4_misc.metal index 7e8cbd8a7..1195bee4c 100644 --- a/metal/dsv4_misc.metal +++ b/metal/dsv4_misc.metal @@ -43,6 +43,21 @@ struct ds4_metal_args_dsv4_softmax_pool { uint64_t nb1; }; +struct ds4_metal_args_dsv4_softmax_pool_ratio4_direct { + int64_t n_rows; + uint32_t head_dim; + uint32_t n_comp; + uint32_t replay; + uint32_t pad; +}; + +struct ds4_metal_args_dsv4_compressor_score_ape { + uint32_t width; + uint32_t ratio; + uint32_t pos0; + uint32_t n_tokens; +}; + struct ds4_metal_args_dsv4_indexed_attention { uint32_t n_tokens; uint32_t n_head; @@ -218,6 +233,54 @@ kernel void kernel_dsv4_router_weights_one( w[tid] = p[s[tid]] / sum * 1.5f; } +// Batched Flash-router weight finalization after selection is already known. +// Six active lanes deliberately match kernel_sum_rows_f32_f32's reduction +// topology. The denominator and divided weights cross threadgroup storage +// boundaries so division cannot be reassociated with the final scale. +kernel void kernel_dsv4_router_weights_batch( + constant float &scale, + device const float *probs, + device const int32_t *selected, + device float *weights, + threadgroup volatile float *scratch [[threadgroup(0)]], + uint row [[threadgroup_position_in_grid]], + ushort tid [[thread_position_in_threadgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]]) { + if (tid >= 6) return; + + threadgroup volatile float *sum_scratch = scratch; + threadgroup volatile float *denom_scratch = scratch + 32; + threadgroup volatile float *div_scratch = scratch + 33; + const uint out_index = row * 6u + (uint)tid; + const int32_t expert = selected[out_index]; + const float p = probs[row * 256u + (uint)expert]; + + // Keep this sequence identical to kernel_sum_rows_f32_f32 for width 6. + if (sgitg == 0) { + sum_scratch[tiisg] = 0.0f; + } + float sumf = 0.0f; + sumf += p; + sumf = simd_sum(sumf); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tiisg == 0) { + sum_scratch[sgitg] = sumf; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + sumf = sum_scratch[tiisg]; + sumf = simd_sum(sumf); + + if (tid == 0) { + denom_scratch[0] = clamp(sumf, 6.103515625e-5f, INFINITY); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + div_scratch[tid] = p / denom_scratch[0]; + threadgroup_barrier(mem_flags::mem_threadgroup); + weights[out_index] = div_scratch[tid] * scale; +} + // Decode router selection for one token after the existing // sqrt(softplus(logit)) probability kernel has run. Bias affects only top-k // selection. Route-weight normalization deliberately stays in the old one-token @@ -279,6 +342,176 @@ kernel void kernel_dsv4_router_finalize_one( threadgroup_barrier(mem_flags::mem_threadgroup); } +// M3 decode specialization for the non-hash one-token router. Scores and ids +// stay in registers. Intra-SIMD bitonic stages use shuffle-xor; the six stages +// that cross 32-lane SIMD groups exchange through alternating threadgroup +// banks. The next bank's publish barrier proves every prior-bank read finished; +// by the time a bank is reused two cross stages later, no reader can remain. +kernel void kernel_dsv4_router_finalize_one_simd( + constant ds4_metal_args_dsv4_router_select_one & args, + device const float *probs, + device const float *bias, + device const int32_t *hash, + device const int32_t *tokens, + device int32_t *selected, + threadgroup float *scratch [[threadgroup(0)]], + uint tid [[thread_position_in_threadgroup]]) { + if (tid >= 256 || args.hash_mode) return; + + (void)hash; + (void)tokens; + threadgroup float *score0_tg = scratch; + threadgroup int32_t *idx0_tg = + (threadgroup int32_t *)(scratch + 256); + threadgroup float *score1_tg = scratch + 512; + threadgroup int32_t *idx1_tg = + (threadgroup int32_t *)(scratch + 768); + const float p = probs[tid]; + float score = args.has_bias ? p + bias[tid] : p; + int32_t idx = (int32_t)tid; + uint cross_stage = 0; + + for (uint k = 2; k <= 256; k <<= 1) { + for (uint j = k >> 1; j > 0; j >>= 1) { + float peer_score; + int32_t peer_idx; + bool take_peer; + const bool lower = (tid & j) == 0; + const bool descending = (tid & k) == 0; + + if (j < 32) { + peer_score = simd_shuffle_xor(score, (ushort)j); + peer_idx = simd_shuffle_xor(idx, (ushort)j); + take_peer = descending + ? (lower ? score < peer_score : score > peer_score) + : (lower ? score > peer_score : score < peer_score); + if (take_peer) { + score = peer_score; + idx = peer_idx; + } + } else { + threadgroup float *score_tg = + (cross_stage & 1u) != 0u ? score1_tg : score0_tg; + threadgroup int32_t *idx_tg = + (cross_stage & 1u) != 0u ? idx1_tg : idx0_tg; + score_tg[tid] = score; + idx_tg[tid] = idx; + threadgroup_barrier(mem_flags::mem_threadgroup); + + const uint other = tid ^ j; + peer_score = score_tg[other]; + peer_idx = idx_tg[other]; + take_peer = descending + ? (lower ? score < peer_score : score > peer_score) + : (lower ? score > peer_score : score < peer_score); + if (take_peer) { + score = peer_score; + idx = peer_idx; + } + cross_stage++; + } + } + } + + if (tid < 6) { + selected[tid] = idx; + } +} + +// M3 decode specialization that extends the register/TG SIMD selection above +// through the existing six-value serial weight normalization. The selected ids +// cross the same device-memory boundary as the standalone weight kernel; +// volatile TG stores pin its left-fold and scaled-reciprocal rounding points. +kernel void kernel_dsv4_router_finalize_weights_one_simd( + constant ds4_metal_args_dsv4_router_select_one & args, + device const float *probs, + device const float *bias, + device const int32_t *hash, + device const int32_t *tokens, + device int32_t *selected, + device float *weights, + threadgroup float *scratch [[threadgroup(0)]], + uint tid [[thread_position_in_threadgroup]]) { + if (tid >= 256 || args.hash_mode) return; + + (void)hash; + (void)tokens; + threadgroup float *score0_tg = scratch; + threadgroup int32_t *idx0_tg = + (threadgroup int32_t *)(scratch + 256); + threadgroup float *score1_tg = scratch + 512; + threadgroup int32_t *idx1_tg = + (threadgroup int32_t *)(scratch + 768); + const float p = probs[tid]; + float score = args.has_bias ? p + bias[tid] : p; + int32_t idx = (int32_t)tid; + uint cross_stage = 0; + + for (uint k = 2; k <= 256; k <<= 1) { + for (uint j = k >> 1; j > 0; j >>= 1) { + float peer_score; + int32_t peer_idx; + bool take_peer; + const bool lower = (tid & j) == 0; + const bool descending = (tid & k) == 0; + + if (j < 32) { + peer_score = simd_shuffle_xor(score, (ushort)j); + peer_idx = simd_shuffle_xor(idx, (ushort)j); + take_peer = descending + ? (lower ? score < peer_score : score > peer_score) + : (lower ? score > peer_score : score < peer_score); + if (take_peer) { + score = peer_score; + idx = peer_idx; + } + } else { + threadgroup float *score_tg = + (cross_stage & 1u) != 0u ? score1_tg : score0_tg; + threadgroup int32_t *idx_tg = + (cross_stage & 1u) != 0u ? idx1_tg : idx0_tg; + score_tg[tid] = score; + idx_tg[tid] = idx; + threadgroup_barrier(mem_flags::mem_threadgroup); + + const uint other = tid ^ j; + peer_score = score_tg[other]; + peer_idx = idx_tg[other]; + take_peer = descending + ? (lower ? score < peer_score : score > peer_score) + : (lower ? score > peer_score : score < peer_score); + if (take_peer) { + score = peer_score; + idx = peer_idx; + } + cross_stage++; + } + } + } + + if (tid < 6) { + selected[tid] = idx; + } + threadgroup_barrier(mem_flags::mem_device); + + threadgroup volatile float *norm_scratch = + (threadgroup volatile float *)scratch; + if (tid == 0) { + device const int32_t *s = selected; + norm_scratch[0] = 0.0f; + for (uint i = 0; i < 6; i++) { + norm_scratch[0] = norm_scratch[0] + probs[s[i]]; + } + norm_scratch[0] = max(norm_scratch[0], 6.103515625e-5f); + norm_scratch[1] = 1.5f / norm_scratch[0]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 6) { + device const int32_t *s = selected; + weights[tid] = probs[s[tid]] * norm_scratch[1]; + } +} + // Fills the dense compressed-attention mask with -inf. The selected top-k rows // are enabled by kernel_dsv4_topk_mask_scatter in a second ordered dispatch. kernel void kernel_dsv4_topk_mask( @@ -1118,6 +1351,15 @@ kernel void kernel_dsv4_indexer_scores_tiled( // compressed-row dot tile. The kernel intentionally leaves top-k selection and // indexed attention semantics unchanged; all 512 selected rows remain available // to the later attention kernel. +// +// Each matmul processes a pair of heads (TQ = 2 x TM q rows): the per-element +// dot is still a 128-deep reduction in 32-wide k-steps, so scores are +// bit-identical to single-head tiles while the run count halves. The q tile +// is double-buffered, so the next k-step's stage overlaps the current +// cooperative matmul and each pair needs 5 barriers instead of 10. q and k +// staging use one float4/half4 per lane (each thread covers one row of 8/32 +// consecutive elements), which is the same half(float) conversion per element +// as the scalar form. kernel void kernel_dsv4_indexer_scores_nax( constant ds4_metal_args_dsv4_indexer_scores_fused & args, device const char *q, @@ -1128,6 +1370,7 @@ kernel void kernel_dsv4_indexer_scores_nax( uint2 tgpig [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { constexpr int TM = 16; + constexpr int TQ = 32; constexpr int TN = 32; constexpr int NK = 32; constexpr int D = 128; @@ -1139,9 +1382,9 @@ kernel void kernel_dsv4_indexer_scores_nax( const uint c0 = tgpig.x * TN; const uint t0 = tgpig.y * TM; - threadgroup half *qtg = shared; // [16][32] - threadgroup half *ktg = qtg + TM*NK; // [32][128] - threadgroup float *dot = (threadgroup float *)(ktg + TN*D); // [16][32], column-major + threadgroup half *qtg = shared; // 2 x [TQ][NK] + threadgroup half *ktg = qtg + 2*TQ*NK; // [32][128] + threadgroup float *dot = (threadgroup float *)(ktg + TN*D); // [TQ][TN], column-major const uint last_token = min(t0 + (uint)TM, args.n_tokens); const uint max_visible = last_token > t0 ? @@ -1162,17 +1405,21 @@ kernel void kernel_dsv4_indexer_scores_nax( return; } - for (uint work = tid; work < TN*D; work += NUM_THREADS) { - const uint cc = work / D; - const uint d = work - cc*D; + { + // One compressed row per 4 threads, 32 consecutive floats per thread. + const uint cc = tid / 4; const uint comp = c0 + cc; - half v = half(0.0f); + device const float *krow = nullptr; if (comp < args.n_comp) { - device const float *krow = (device const float *)(index_comp + + krow = (device const float *)(index_comp + (uint64_t)comp * args.index_row_stride); - v = half(krow[d]); } - ktg[cc*D + d] = v; + const uint d0 = (tid % 4) * 32; + FOR_UNROLL (uint j = 0; j < 8; j++) { + const float4 kv = krow ? *(device const float4 *)(krow + d0 + 4*j) + : float4(0.0f); + *(threadgroup half4 *)(ktg + cc*D + d0 + 4*j) = half4(kv); + } } threadgroup_barrier(mem_flags::mem_threadgroup); @@ -1182,17 +1429,45 @@ kernel void kernel_dsv4_indexer_scores_nax( acc[j] = 0.0f; } - auto tq = tensor(qtg, dextents(NK, TM)); + auto tq0 = tensor(qtg, dextents(NK, TQ)); + auto tq1 = tensor(qtg + TQ*NK, dextents(NK, TQ)); auto tk = tensor(ktg, dextents(D, TN)); - auto td = tensor(dot, dextents(TM, TN), array({1, TM})); + auto td = tensor(dot, dextents(TQ, TN), array({1, TQ})); matmul2d< - matmul2d_descriptor(TN, TM, NK, false, true, false, + matmul2d_descriptor(TN, TQ, NK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate), execution_simdgroups<4>> mm; - for (uint head = 0; head < args.n_head; head++) { - auto ct = mm.template get_destination_cooperative_tensor(); + // One q row per 4 threads, 8 consecutive floats per thread. Row r covers + // head (r / TM) of the pair and token row (r % TM). + const uint q_r = tid / 4; + const uint q_k4 = (tid % 4) * 8; + const uint q_hl = q_r / TM; + const uint q_tr = q_r % TM; + const uint q_token = t0 + q_tr; + device const char *q_row_base = nullptr; + if (q_token < args.n_tokens) { + q_row_base = q + (uint64_t)q_token * args.q_token_stride; + } + + auto stage_q = [&](const uint head0, const uint loop_k, threadgroup half *buf) { + const uint head = head0 + q_hl; + half4 v0 = half4(0.0f); + half4 v1 = half4(0.0f); + if (q_row_base && head < args.n_head) { + device const float4 *src4 = (device const float4 *) + (q_row_base + (uint64_t)head * args.q_head_stride + + (uint64_t)(loop_k + q_k4) * sizeof(float)); + v0 = half4(src4[0]); + v1 = half4(src4[1]); + } + *(threadgroup half4 *)(buf + q_r*NK + q_k4) = v0; + *(threadgroup half4 *)(buf + q_r*NK + q_k4 + 4) = v1; + }; + + for (uint head0 = 0; head0 < args.n_head; head0 += 2) { + auto ct = mm.template get_destination_cooperative_tensor(); #pragma unroll for (uint16_t i = 0; i < ct.get_capacity(); i++) { if (ct.is_valid_element(i)) { @@ -1200,27 +1475,19 @@ kernel void kernel_dsv4_indexer_scores_nax( } } - for (uint loop_k = 0; loop_k < D; loop_k += NK) { - for (uint work = tid; work < TM*NK; work += NUM_THREADS) { - const uint r = work / NK; - const uint k = work - r*NK; - const uint token = t0 + r; - half v = half(0.0f); - if (token < args.n_tokens) { - device const float *qrow = (device const float *)(q + - (uint64_t)token * args.q_token_stride + - (uint64_t)head * args.q_head_stride); - v = half(qrow[loop_k + k]); - } - qtg[r*NK + k] = v; - } - threadgroup_barrier(mem_flags::mem_threadgroup); + stage_q(head0, 0, qtg); + threadgroup_barrier(mem_flags::mem_threadgroup); - auto mq = tq.slice(0, 0); - auto mk = tk.slice(loop_k, 0); + uint qsel = 0; + FOR_UNROLL (uint i = 0; i < 4; i++) { + auto mk = tk.slice(i*NK, 0); + auto mq = (qsel ? tq1 : tq0).slice(0, 0); mm.run(mk, mq, ct); - - threadgroup_barrier(mem_flags::mem_threadgroup); + if (i < 3) { + qsel ^= 1u; + stage_q(head0, (i + 1)*NK, qsel ? qtg + TQ*NK : qtg); + threadgroup_barrier(mem_flags::mem_threadgroup); + } } ct.store(td); @@ -1236,12 +1503,16 @@ kernel void kernel_dsv4_indexer_scores_nax( if (token < args.n_tokens) { device const float *w = (device const float *)(weights + (uint64_t)token * args.weights_token_stride); - acc[j] += max(dot[cc*TM + r], 0.0f) * (w[head] * args.scale); + acc[j] += max(dot[cc*TQ + r], 0.0f) * (w[head0] * args.scale); + if (head0 + 1 < args.n_head) { + acc[j] += max(dot[cc*TQ + TM + r], 0.0f) * (w[head0 + 1] * args.scale); + } } } } - - threadgroup_barrier(mem_flags::mem_threadgroup); + // No barrier here: the next pair's q stage and these dot reads touch + // different buffers, and the next q-stage barrier separates the next + // ct.store from these reads. } #pragma unroll @@ -1290,6 +1561,39 @@ kernel void kernel_dsv4_indexer_weighted_sum( *((device float *) (dst + ic*args.nb0 + it*args.nb1)) = acc; } +// Adds the periodic compressor APE directly to projected scores. The legacy +// path materializes one repeated APE segment per period and then performs this +// same single F32 add; these kernels remove only that intermediate copy graph. +kernel void kernel_dsv4_compressor_score_ape_f32( + constant ds4_metal_args_dsv4_compressor_score_ape & args, + device const float *score, + device const float *ape, + device float *dst, + uint gid [[thread_position_in_grid]]) { + const uint64_t total = (uint64_t)args.n_tokens * args.width; + if ((uint64_t)gid >= total) return; + + const uint token = gid / args.width; + const uint col = gid - token*args.width; + const uint ape_row = (uint)(((uint64_t)args.pos0 + token) % args.ratio); + dst[gid] = score[gid] + ape[(uint64_t)ape_row*args.width + col]; +} + +kernel void kernel_dsv4_compressor_score_ape_f16( + constant ds4_metal_args_dsv4_compressor_score_ape & args, + device const float *score, + device const half *ape, + device float *dst, + uint gid [[thread_position_in_grid]]) { + const uint64_t total = (uint64_t)args.n_tokens * args.width; + if ((uint64_t)gid >= total) return; + + const uint token = gid / args.width; + const uint col = gid - token*args.width; + const uint ape_row = (uint)(((uint64_t)args.pos0 + token) % args.ratio); + dst[gid] = score[gid] + float(ape[(uint64_t)ape_row*args.width + col]); +} + // Fused softmax-weighted pooling of compressed KV rows. It is used when several // compressor rows are present; the one-row case deliberately follows the // unfused softmax/mul/sum graph in Objective-C to keep identical reductions. @@ -1325,3 +1629,86 @@ kernel void kernel_dsv4_softmax_pool( *((device float *) (dst + id*args.nb0 + ic*args.nb1)) = acc/sum; } + +// Ratio-4 compressor pooling without materializing the [n_comp, 8, head_dim] +// KV and score packs. The row mapping and both reduction loops deliberately +// match kernel_dsv4_softmax_pool so the arithmetic order is unchanged. +kernel void kernel_dsv4_softmax_pool_ratio4_direct( + constant ds4_metal_args_dsv4_softmax_pool_ratio4_direct & args, + device const float * kv, + device const float * score, + device const float * state_kv, + device const float * state_score, + device float * dst, + uint gid [[thread_position_in_grid]]) { + const uint64_t n = (uint64_t)args.head_dim * args.n_comp; + if ((uint64_t)gid >= n || args.head_dim == 0u) { + return; + } + + const uint64_t id = gid % args.head_dim; + const uint64_t ic = gid / args.head_dim; + const uint64_t input_row_stride = 2ull * args.head_dim; + + float max_s = -INFINITY; + float sum = 0.0f; + float acc = 0.0f; + if (ic != 0u) { + const int64_t token_base = (int64_t)ic * 4 - 4; + for (int64_t ir = 0; ir < args.n_rows; ++ir) { + const uint64_t token = (uint64_t)(token_base + ir); + const uint64_t src = token * input_row_stride + + ((uint64_t)ir >> 2u) * args.head_dim + id; + const float s = score[src]; + max_s = max(max_s, s); + } + + for (int64_t ir = 0; ir < args.n_rows; ++ir) { + const uint64_t token = (uint64_t)(token_base + ir); + const uint64_t src = token * input_row_stride + + ((uint64_t)ir >> 2u) * args.head_dim + id; + const float s = score[src]; + const float w = exp(s - max_s); + const float v = kv[src]; + sum += w; + acc += v*w; + } + } else { + for (int64_t ir = 0; ir < args.n_rows; ++ir) { + float s; + if (ir >= 4) { + const uint64_t src = (uint64_t)(ir - 4) * input_row_stride + + args.head_dim + id; + s = score[src]; + } else if (args.replay != 0u) { + s = state_score[(uint64_t)ir * input_row_stride + id]; + } else { + s = -INFINITY; + } + max_s = max(max_s, s); + } + + for (int64_t ir = 0; ir < args.n_rows; ++ir) { + float s; + float v; + if (ir >= 4) { + const uint64_t src = (uint64_t)(ir - 4) * input_row_stride + + args.head_dim + id; + s = score[src]; + v = kv[src]; + } else if (args.replay != 0u) { + const uint64_t src = (uint64_t)ir * input_row_stride + id; + s = state_score[src]; + v = state_kv[src]; + } else { + s = -INFINITY; + v = 0.0f; + } + const float w = exp(s - max_s); + sum += w; + acc += v*w; + } + } + + dst[ic * args.head_dim + id] = acc/sum; +} diff --git a/metal/dsv4_rope.metal b/metal/dsv4_rope.metal index b32075612..94a83c7c0 100644 --- a/metal/dsv4_rope.metal +++ b/metal/dsv4_rope.metal @@ -24,6 +24,23 @@ struct ds4_metal_args_dsv4_rope_tail { bool src2; }; +struct ds4_metal_args_dsv4_rope_affine_pair { + uint64_t row_bytes; + uint64_t token_bytes; + int32_t head_dim; + int32_t n_dims; + int32_t n_ctx_orig; + int32_t inverse; + uint32_t pos0; + uint32_t pos_step; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; +}; + static float rope_yarn_ramp(const float low, const float high, const int i0) { const float y = (i0 / 2 - low) / max(0.001f, high - low); return 1.0f - min(1.0f, max(0.0f, y)); @@ -163,3 +180,206 @@ kernel void kernel_dsv4_rope_tail_f32( } } } + +// DS4 only calls the Metal RoPE helper in-place and uses the adjacent-pair +// layout (mode 0). The generic kernel above still copies the no-position +// prefix, even though source and destination alias, and dispatches enough +// lanes for the full head. This specialization maps lanes directly to the +// rotated pairs and deliberately never reads or writes the unchanged prefix. +// Keep the arithmetic below in the same order as the mode-0 branch above so +// the optimized and reference paths remain bit-identical. +kernel void kernel_dsv4_rope_tail_f32_inplace_pair( + constant ds4_metal_args_dsv4_rope_tail & args, + device const char * src0, + device const char * src1, + device const char * src2, + device char * dst, + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + if (args.mode != 0) { + return; + } + + const int i1 = tgpig[0]; + const int i2 = tgpig[1]; + const int i3 = tgpig[2]; + const int n_nope = args.ne00 - args.n_dims; + if (n_nope < 0) { + return; + } + + device const int32_t * pos = (device const int32_t *) src1; + + float corr_dims[2]; + rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims); + + const float theta_base = (float) pos[i2]; + const float inv_ndims = -1.f/args.n_dims; + device const char * src_base = src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01; + device char * dst_base = dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1; + + // Keep each pair on the same SIMD lane as the generic in-place dispatch. + // n_nope is 32-aligned for the supported DS4 shapes, so logical tail index + // r ran on lane r%32 in the reference kernel. Compacting pairs would move + // them to lane (r/2)%32 and can perturb fast-math code generation. + for (int r = tid; r < args.n_dims; r += ntg.x) { + if ((r & 1) != 0) { + continue; + } + const int ic = r/2; +#ifdef DS4_METAL_ROPE_EXP2_LOG2 + const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base)); +#else + const float theta = theta_base * pow(args.freq_base, inv_ndims*r); +#endif + const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f; + + float cos_theta; + float sin_theta; + rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, r, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta); + if (args.inverse) { + sin_theta = -sin_theta; + } + + const int j0 = n_nope + r; + const int j1 = j0 + 1; + const float x0 = *((device const float *) (src_base + j0*args.nb00)); + const float x1 = *((device const float *) (src_base + j1*args.nb00)); + + *((device float *) (dst_base + j0*args.nb0)) = x0*cos_theta - x1*sin_theta; + *((device float *) (dst_base + j1*args.nb0)) = x0*sin_theta + x1*cos_theta; + } +} + +// The Q/K RoPE calls use the same position and scaling parameters for every +// head. Group four heads into one threadgroup so one 64-thread cohort computes +// the 32 adjacent-pair coefficients and the other cohorts reuse them. Keeping +// r on the same r%32 SIMD lane as the per-head specialization preserves the +// fast-math instruction mapping; only the redundant coefficient work changes. +kernel void kernel_dsv4_rope_tail_f32_inplace_pair_shared4( + constant ds4_metal_args_dsv4_rope_tail & args, + device const char * src0, + device const char * src1, + device const char * src2, + device char * dst, + uint tid [[thread_index_in_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + if (args.mode != 0 || args.n_dims != 64) { + return; + } + + const int n_nope = args.ne00 - args.n_dims; + if (n_nope < 0) { + return; + } + + const uint cohort = tid >> 6; + const int r = (int)(tid & 63u); + threadgroup float cos_shared[32]; + threadgroup float sin_shared[32]; + + device const int32_t * pos = (device const int32_t *) src1; + float corr_dims[2]; + rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims); + + const int i2 = tgpig[1]; + const float theta_base = (float) pos[i2]; + const float inv_ndims = -1.f/args.n_dims; + + if (cohort == 0 && (r & 1) == 0) { + const int ic = r/2; +#ifdef DS4_METAL_ROPE_EXP2_LOG2 + const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base)); +#else + const float theta = theta_base * pow(args.freq_base, inv_ndims*r); +#endif + const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f; + + float cos_theta; + float sin_theta; + rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, r, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta); + if (args.inverse) { + sin_theta = -sin_theta; + } + cos_shared[ic] = cos_theta; + sin_shared[ic] = sin_theta; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + const int i1 = (int)tgpig[0]*4 + (int)cohort; + if (i1 >= args.ne01 || (r & 1) != 0) { + return; + } + + const int i3 = tgpig[2]; + device const char * src_base = src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01; + device char * dst_base = dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1; + const int j0 = n_nope + r; + const int j1 = j0 + 1; + const float x0 = *((device const float *) (src_base + j0*args.nb00)); + const float x1 = *((device const float *) (src_base + j1*args.nb00)); + const float cos_theta = cos_shared[r/2]; + const float sin_theta = sin_shared[r/2]; + + *((device float *) (dst_base + j0*args.nb0)) = x0*cos_theta - x1*sin_theta; + *((device float *) (dst_base + j1*args.nb0)) = x0*sin_theta + x1*cos_theta; +} + +// DS4 positions are always affine within one RoPE dispatch. This variant +// reconstructs the same wrapped int32 position in-kernel, avoiding the host +// position array and its buffer binding while preserving the pair lane mapping +// and all floating-point operations of the specialization above. +kernel void kernel_dsv4_rope_tail_f32_inplace_pair_affine( + constant ds4_metal_args_dsv4_rope_affine_pair & args [[buffer(0)]], + device const char * src0 [[buffer(1)]], + device char * dst [[buffer(4)]], + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const int i1 = tgpig[0]; + const int i2 = tgpig[1]; + const int n_nope = args.head_dim - args.n_dims; + if (n_nope < 0) { + return; + } + + float corr_dims[2]; + rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims); + + const uint raw_pos = args.pos0 + (uint)i2 * args.pos_step; + const float theta_base = (float)as_type(raw_pos); + const float inv_ndims = -1.f/args.n_dims; + device const char * src_base = + src0 + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes; + device char * dst_base = + dst + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes; + + for (int r = tid; r < args.n_dims; r += ntg.x) { + if ((r & 1) != 0) { + continue; + } +#ifdef DS4_METAL_ROPE_EXP2_LOG2 + const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base)); +#else + const float theta = theta_base * pow(args.freq_base, inv_ndims*r); +#endif + const float freq_factor = 1.0f; + + float cos_theta; + float sin_theta; + rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, r, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta); + if (args.inverse) { + sin_theta = -sin_theta; + } + + const int j0 = n_nope + r; + const int j1 = j0 + 1; + const float x0 = *((device const float *) (src_base + j0*sizeof(float))); + const float x1 = *((device const float *) (src_base + j1*sizeof(float))); + + *((device float *) (dst_base + j0*sizeof(float))) = x0*cos_theta - x1*sin_theta; + *((device float *) (dst_base + j1*sizeof(float))) = x0*sin_theta + x1*cos_theta; + } +} diff --git a/metal/flash_attn.metal b/metal/flash_attn.metal index d069d43cc..8c8cd5f7a 100644 --- a/metal/flash_attn.metal +++ b/metal/flash_attn.metal @@ -934,6 +934,7 @@ constant bool FC_flash_attn_ext_vec_has_sinks [[function_constant(FC_FLASH_ATTN_ constant bool FC_flash_attn_ext_vec_has_bias [[function_constant(FC_FLASH_ATTN_EXT_VEC + 2)]]; constant bool FC_flash_attn_ext_vec_has_scap [[function_constant(FC_FLASH_ATTN_EXT_VEC + 3)]]; constant bool FC_flash_attn_ext_vec_has_kvpad [[function_constant(FC_FLASH_ATTN_EXT_VEC + 4)]]; +constant bool FC_flash_attn_ext_vec_shared_kvpad [[function_constant(FC_FLASH_ATTN_EXT_VEC + 5)]]; constant int32_t FC_flash_attn_ext_vec_ns10 [[function_constant(FC_FLASH_ATTN_EXT_VEC + 20)]]; constant int32_t FC_flash_attn_ext_vec_ns20 [[function_constant(FC_FLASH_ATTN_EXT_VEC + 21)]]; constant int32_t FC_flash_attn_ext_vec_nsg [[function_constant(FC_FLASH_ATTN_EXT_VEC + 22)]]; @@ -1069,8 +1070,14 @@ kernel void kernel_flash_attn_ext_vec( if (FC_flash_attn_ext_vec_has_kvpad && ic + C > args.ne11) { k = pad; - v = k + args.nb11*C*args.ne_12_2*args.ne_12_3; - mask = v + args.nb21*C*args.ne_12_2*args.ne_12_3; + if (FC_flash_attn_ext_vec_shared_kvpad) { + v = k; + mask = k + args.nb11*C*args.ne_12_2*args.ne_12_3 + + args.nb21*C*args.ne_12_2*args.ne_12_3; + } else { + v = k + args.nb11*C*args.ne_12_2*args.ne_12_3; + mask = v + args.nb21*C*args.ne_12_2*args.ne_12_3; + } const short ikv2 = iq2/(args.ne02/args.ne_12_2); const short ikv3 = iq3/(args.ne03/args.ne_12_3); diff --git a/metal/moe.metal b/metal/moe.metal index d47fab203..13eca92c2 100644 --- a/metal/moe.metal +++ b/metal/moe.metal @@ -4256,7 +4256,13 @@ kernel void kernel_mul_mm_id( } } -kernel void kernel_mul_mm_id_iq2_xxs_pair_swiglu_f16( +// Fused routed gate+up grouped matmul with the SwiGLU epilogue, generic over +// the expert quant block. Both A streams share one staged B tile per k-step; +// each output keeps the exact MMA accumulation order of the separate GEMMs, +// and the epilogue matches kernel_dsv4_moe_swiglu_weight_f16, so the fused +// result is bit-identical to the unfused path. +template +kernel void kernel_mul_mm_id_pair_swiglu_f16_impl( constant ds4_metal_args_mul_mm_id & args, constant ds4_metal_dsv4_moe_swiglu_weight_args & act, device const char * src0_gate, @@ -4271,8 +4277,9 @@ kernel void kernel_mul_mm_id_iq2_xxs_pair_swiglu_f16( ushort tiitg[[thread_index_in_threadgroup]], ushort tiisg[[thread_index_in_simdgroup]], ushort sgitg[[simdgroup_index_in_threadgroup]]) { - threadgroup half *sa = (threadgroup half *)(shmem); - threadgroup half *sb = (threadgroup half *)(shmem + 4096); + threadgroup half *sa_gate = (threadgroup half *)(shmem); + threadgroup half *sa_up = (threadgroup half *)(shmem + 4096); + threadgroup half *sb = (threadgroup half *)(shmem + 8192); constexpr int NR0 = 64; constexpr int NR1 = 32; @@ -4311,10 +4318,10 @@ kernel void kernel_mul_mm_id_iq2_xxs_pair_swiglu_f16( const uint64_t offset0 = im*args.nb02 + i13*args.nb03; const short offset1 = il0/QK_NL; - device const block_iq2_xxs * xg = - (device const block_iq2_xxs *)(src0_gate + args.nb01*(r0 + lr0) + offset0) + offset1; - device const block_iq2_xxs * xu = - (device const block_iq2_xxs *)(src0_up + args.nb01*(r0 + lr0) + offset0) + offset1; + device const block_q * xg = + (device const block_q *)(src0_gate + args.nb01*(r0 + lr0) + offset0) + offset1; + device const block_q * xu = + (device const block_q *)(src0_up + args.nb01*(r0 + lr0) + offset0) + offset1; const short iy = 8*(tiitg % NL1); @@ -4324,7 +4331,8 @@ kernel void kernel_mul_mm_id_iq2_xxs_pair_swiglu_f16( + args.nb11*i11 + args.nb10*iy); - simdgroup_half8x8 ma[4]; + simdgroup_half8x8 ma_g[4]; + simdgroup_half8x8 ma_u[4]; simdgroup_half8x8 mb[2]; simdgroup_float8x8 mc_gate[8]; @@ -4336,83 +4344,47 @@ kernel void kernel_mul_mm_id_iq2_xxs_pair_swiglu_f16( } for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) { - const short sx_b = (tiitg%NL1); - const short sy_b = (tiitg/NL1)/8; - const short ly_b = (tiitg/NL1)%8; - const short ib_b = 4*sx_b + sy_b; - *(threadgroup half2x4 *)(sb + 64*ib_b + 8*ly_b) = - (half2x4)(*((device float2x4 *) y)); - half4x4 temp_gate; - dequantize_iq2_xxs(xg, il, temp_gate); - - threadgroup_barrier(mem_flags::mem_threadgroup); - - FOR_UNROLL (short i = 0; i < 16; i++) { - const short sx = 2*il0 + i/8; - const short sy = (tiitg/NL0)/8; - const short lx = (tiitg/NL0)%8; - const short ly = i%8; - const short ib = 8*sx + sy; - *(sa + 64*ib + 8*ly + lx) = temp_gate[i/4][i%4]; - } + dequantize_func(xg, il, temp_gate); + half4x4 temp_up; + dequantize_func(xu, il, temp_up); threadgroup_barrier(mem_flags::mem_threadgroup); - threadgroup const half * lsma_gate = (sa + 4*64*(sgitg%2)); - threadgroup const half * lsmb = (sb + 2*64*(sgitg/2)); - - FOR_UNROLL (short ik = 0; ik < NK/8; ik++) { - simdgroup_barrier(mem_flags::mem_none); - - FOR_UNROLL (short i = 0; i < 4; i++) { - simdgroup_load(ma[i], lsma_gate + 64*i, 8, 0, false); - } - - simdgroup_barrier(mem_flags::mem_none); - - FOR_UNROLL (short i = 0; i < 2; i++) { - simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false); - } - - simdgroup_barrier(mem_flags::mem_none); - - FOR_UNROLL (short i = 0; i < 8; i++) { - simdgroup_multiply_accumulate(mc_gate[i], mb[i/4], ma[i%4], mc_gate[i]); - } - - lsma_gate += 8*64; - lsmb += 4*64; + // Stage the shared B tile and both expert A tiles in one pass: two + // threadgroup barriers per k-step total instead of two per matrix. + { + const short sx_b = (tiitg%NL1); + const short sy_b = (tiitg/NL1)/8; + const short ly_b = (tiitg/NL1)%8; + const short ib_b = 4*sx_b + sy_b; + *(threadgroup half2x4 *)(sb + 64*ib_b + 8*ly_b) = + (half2x4)(*((device float2x4 *) y)); } - half4x4 temp_up; - dequantize_iq2_xxs(xu, il, temp_up); - - threadgroup_barrier(mem_flags::mem_threadgroup); - FOR_UNROLL (short i = 0; i < 16; i++) { const short sx = 2*il0 + i/8; const short sy = (tiitg/NL0)/8; const short lx = (tiitg/NL0)%8; const short ly = i%8; const short ib = 8*sx + sy; - *(sa + 64*ib + 8*ly + lx) = temp_up[i/4][i%4]; + *(sa_gate + 64*ib + 8*ly + lx) = temp_gate[i/4][i%4]; + *(sa_up + 64*ib + 8*ly + lx) = temp_up[i/4][i%4]; } threadgroup_barrier(mem_flags::mem_threadgroup); - threadgroup const half * lsma_up = (sa + 4*64*(sgitg%2)); - lsmb = (sb + 2*64*(sgitg/2)); + threadgroup const half * lsma_gate = (sa_gate + 4*64*(sgitg%2)); + threadgroup const half * lsma_up = (sa_up + 4*64*(sgitg%2)); + threadgroup const half * lsmb = (sb + 2*64*(sgitg/2)); FOR_UNROLL (short ik = 0; ik < NK/8; ik++) { simdgroup_barrier(mem_flags::mem_none); FOR_UNROLL (short i = 0; i < 4; i++) { - simdgroup_load(ma[i], lsma_up + 64*i, 8, 0, false); + simdgroup_load(ma_g[i], lsma_gate + 64*i, 8, 0, false); + simdgroup_load(ma_u[i], lsma_up + 64*i, 8, 0, false); } - - simdgroup_barrier(mem_flags::mem_none); - FOR_UNROLL (short i = 0; i < 2; i++) { simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false); } @@ -4420,11 +4392,13 @@ kernel void kernel_mul_mm_id_iq2_xxs_pair_swiglu_f16( simdgroup_barrier(mem_flags::mem_none); FOR_UNROLL (short i = 0; i < 8; i++) { - simdgroup_multiply_accumulate(mc_up[i], mb[i/4], ma[i%4], mc_up[i]); + simdgroup_multiply_accumulate(mc_gate[i], mb[i/4], ma_g[i%4], mc_gate[i]); + simdgroup_multiply_accumulate(mc_up[i], mb[i/4], ma_u[i%4], mc_up[i]); } - lsma_up += 8*64; - lsmb += 4*64; + lsma_gate += 8*64; + lsma_up += 8*64; + lsmb += 4*64; } il = (il + 2 < QK_NL) ? il + 2 : il % 2; @@ -4478,6 +4452,13 @@ kernel void kernel_mul_mm_id_iq2_xxs_pair_swiglu_f16( } } +typedef decltype(kernel_mul_mm_id_pair_swiglu_f16_impl) mul_mm_id_pair_swiglu_f16_iq2; +typedef decltype(kernel_mul_mm_id_pair_swiglu_f16_impl) mul_mm_id_pair_swiglu_f16_q4; + +// Host-visible fused routed pair matmuls for the DS4 expert quant formats. +template [[host_name("kernel_mul_mm_id_iq2_xxs_pair_swiglu_f16")]] kernel mul_mm_id_pair_swiglu_f16_iq2 kernel_mul_mm_id_pair_swiglu_f16_impl; +template [[host_name("kernel_mul_mm_id_q4_K_pair_swiglu_f16")]] kernel mul_mm_id_pair_swiglu_f16_q4 kernel_mul_mm_id_pair_swiglu_f16_impl; + typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>) mul_mm_id; typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, half, half4x4, half, half2x4>) mul_mm_id_f16_rhs; @@ -4497,6 +4478,11 @@ template [[host_name("kernel_mul_mm_id_iq2_xxs_f16")]] kernel mul_mm_id_f16 // weights to a half tile, then let TensorOps read the dense head activations // directly. Only the 64-token direct-RHS instantiation is exported because the // staged-RHS and 32-token variants were benchmark-only experiments. +// +// Full tiles (the host dispatch guarantee for aligned batches) skip all bounds +// work. The weight tile is double-buffered: the next k-step's dequant +// overlaps the current cooperative matmul, so the k-loop needs one +// threadgroup barrier per step instead of two. template kernel void kernel_attn_out_low_q8_0_mpp_direct_rhs( constant ds4_metal_args_mul_mm_id & args, @@ -4524,7 +4510,8 @@ kernel void kernel_attn_out_low_q8_0_mpp_direct_rhs( const bool full_tile = r0 + NR0 <= M && r1 + NR1 <= N && (K % NK) == 0; threadgroup half *sa = (threadgroup half *)shmem; - auto tA = tensor(sa, dextents(NK, NR0)); + auto tA0 = tensor(sa, dextents(NK, NR0)); + auto tA1 = tensor(sa + NR0*NK, dextents(NK, NR0)); device float *ptrB = (device float *)(srcB + args.nb11*group); const int strideB = args.nb12/sizeof(float); @@ -4535,7 +4522,7 @@ kernel void kernel_attn_out_low_q8_0_mpp_direct_rhs( matmul2d_descriptor::mode::multiply_accumulate), execution_simdgroups<4>> mm; - auto cT = mm.template get_destination_cooperative_tensor(); + auto cT = mm.template get_destination_cooperative_tensor(); #pragma unroll for (uint16_t i = 0; i < cT.get_capacity(); ++i) { @@ -4544,37 +4531,57 @@ kernel void kernel_attn_out_low_q8_0_mpp_direct_rhs( } } - for (int loop_k = 0; loop_k < K; loop_k += NK) { + auto stage_tile = [&](const int loop_k, threadgroup half *buf) { for (int work = tiitg; work < NR0*NL; work += NUM_THREADS) { const int row = work/NL; const int k_chunk = work%NL; const int k_pos = loop_k + k_chunk*16; const short k_base = k_chunk*16; - if (full_tile || r0 + row < M) { + if (full_tile) { + device const block_q8_0 *row_ptr = + (device const block_q8_0 *)(srcA + args.nb01*(r0 + row) + group*args.nb02); + + half4x4 temp_a; + dequantize_q8_0_pairs(row_ptr + k_pos/32, (k_pos/16)%2, temp_a); + threadgroup half4 *dst4 = (threadgroup half4 *)(buf + row*NK + k_base); + dst4[0] = temp_a[0]; + dst4[1] = temp_a[1]; + dst4[2] = temp_a[2]; + dst4[3] = temp_a[3]; + } else if (r0 + row < M) { const int block_idx = k_pos/32; const short il = (k_pos/16)%2; device const block_q8_0 *row_ptr = (device const block_q8_0 *)(srcA + args.nb01*(r0 + row) + group*args.nb02); half4x4 temp_a; - dequantize_q8_0(row_ptr + block_idx, il, temp_a); + dequantize_q8_0_pairs(row_ptr + block_idx, il, temp_a); FOR_UNROLL (short i = 0; i < 16; i++) { - sa[row*NK + k_base + i] = (full_tile || k_pos + i < K) ? temp_a[i/4][i%4] : (half)0; + buf[row*NK + k_base + i] = (k_pos + i < K) ? temp_a[i/4][i%4] : (half)0; } } else { FOR_UNROLL (short i = 0; i < 16; i++) { - sa[row*NK + k_base + i] = (half)0; + buf[row*NK + k_base + i] = (half)0; } } } + }; - threadgroup_barrier(mem_flags::mem_threadgroup); + stage_tile(0, sa); + threadgroup_barrier(mem_flags::mem_threadgroup); - auto mA = tA.slice(0, 0); + uint buf_sel = 0; + for (int loop_k = 0; loop_k < K; loop_k += NK) { + auto mA = (buf_sel ? tA1 : tA0).slice(0, 0); auto mB = tB.slice(loop_k, r1); mm.run(mB, mA, cT); + const int next_k = loop_k + NK; + if (next_k < K) { + buf_sel ^= 1u; + stage_tile(next_k, buf_sel ? sa + NR0*NK : sa); + } threadgroup_barrier(mem_flags::mem_threadgroup); } diff --git a/metal/norm.metal b/metal/norm.metal index 892067043..04428987b 100644 --- a/metal/norm.metal +++ b/metal/norm.metal @@ -84,6 +84,57 @@ typedef decltype(kernel_rms_norm_fuse_impl) kernel_rms_norm_fuse_t; template [[host_name("kernel_rms_norm_f32_4")]] kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl; template [[host_name("kernel_rms_norm_mul_f32_4")]] kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl; +// RMSNorm reduction used when the following F16 matmul applies the row scale +// while staging its RHS tile. Keep the reduction and scale expression exactly +// aligned with kernel_rms_norm_fuse_impl; only the large normalized +// row store is replaced by one F32 scale per row. +kernel void kernel_rms_norm_scale_f32_4( + constant ds4_metal_args_norm & args, + device const char * src0, + device float * dst_scale, + threadgroup float * shmem_f32 [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort3 tpitg[[thread_position_in_threadgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + if (sgitg == 0) { + shmem_f32[tiisg] = 0.0f; + } + + const int i01 = tgpig.x; + const int i02 = tgpig.y; + const int i03 = tgpig.z; + + device const float4 * x = (device const float4 *) (src0 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]); + + float sumf = 0.0f; + + // Preserve the established per-lane accumulation and two-stage reduction. + for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) { + sumf += dot(x[i00], x[i00]); + } + sumf = simd_sum(sumf); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (tiisg == 0) { + shmem_f32[sgitg] = sumf; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + sumf = shmem_f32[tiisg]; + sumf = simd_sum(sumf); + + const float mean = sumf/args.ne00; + const float scale = 1.0f/sqrt(mean + args.eps); + + if (tpitg.x == 0) { + dst_scale[i01] = scale; + } +} + struct ds4_metal_args_qkv_rms_norm { int32_t q_n; int32_t q_n4; diff --git a/rocm/ds4_rocm_compressor.cuh b/rocm/ds4_rocm_compressor.cuh index e9e709501..49b0ca16c 100644 --- a/rocm/ds4_rocm_compressor.cuh +++ b/rocm/ds4_rocm_compressor.cuh @@ -260,7 +260,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 || !cuda_compressor_shape_supported(head_dim, ratio) || n_rot > head_dim || (n_rot & 1u) != 0 || @@ -284,10 +285,12 @@ extern "C" int ds4_gpu_compressor_update_tensor( (emit && !cuda_tensor_has_bytes(comp_cache, 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( diff --git a/rocm/ds4_rocm_matmul.cuh b/rocm/ds4_rocm_matmul.cuh index 0f9e06625..2f3cacbeb 100644 --- a/rocm/ds4_rocm_matmul.cuh +++ b/rocm/ds4_rocm_matmul.cuh @@ -753,6 +753,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 || in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) return 0; diff --git a/tests/ds4_test.c b/tests/ds4_test.c index ea1e52487..3fef60805 100644 --- a/tests/ds4_test.c +++ b/tests/ds4_test.c @@ -138,6 +138,69 @@ static uint64_t test_round_up_u64(uint64_t n, uint64_t align) { return (n + align - 1) & ~(align - 1); } +static uint32_t test_float_ordered_u32(float value) { + uint32_t bits; + memcpy(&bits, &value, sizeof(bits)); + return (bits & 0x80000000u) != 0 ? ~bits : bits | 0x80000000u; +} + +typedef struct { + size_t mismatch_count; + uint32_t max_ulp; + float max_abs; +} test_float_compare_stats; + +static test_float_compare_stats test_compare_float_bits( + const float *reference, + const float *actual, + size_t count) { + test_float_compare_stats stats = {0}; + for (size_t i = 0; i < count; i++) { + if (memcmp(&reference[i], &actual[i], sizeof(float)) != 0) { + stats.mismatch_count++; + } + + const uint32_t ref_ordered = test_float_ordered_u32(reference[i]); + const uint32_t actual_ordered = test_float_ordered_u32(actual[i]); + const uint32_t ulp = ref_ordered > actual_ordered + ? ref_ordered - actual_ordered + : actual_ordered - ref_ordered; + if (ulp > stats.max_ulp) stats.max_ulp = ulp; + + const float abs_error = fabsf(reference[i] - actual[i]); + if (abs_error > stats.max_abs) stats.max_abs = abs_error; + } + return stats; +} + +#if defined(__APPLE__) +static const uint32_t test_copy_f32_patterns[] = { + 0x00000000u, 0x80000000u, /* signed zero */ + 0x7f800000u, 0xff800000u, /* infinities */ + 0x7fc00000u, 0x7fc12345u, 0xffc54321u, 0x7fa00001u, /* NaNs */ + 0x00000001u, 0x007fffffu, 0x80000001u, 0x807fffffu, /* F32 subnormals */ + 0x32ffffffu, 0x33000000u, 0x33000001u, + 0x337fffffu, 0x33800000u, 0x33800001u, /* minimum F16 subnormal boundary */ + 0x387fbfffu, 0x387fc000u, 0x387fffffu, + 0x38800000u, 0x38800001u, /* maximum subnormal/minimum normal boundary */ + 0x3f800fffu, 0x3f801000u, 0x3f801001u, + 0x3f802fffu, 0x3f803000u, 0x3f803001u, /* round-to-nearest ties */ + 0x477fdfffu, 0x477fe000u, 0x477fefffu, + 0x477ff000u, 0x477ff001u, 0x47800000u, /* maximum/overflow boundary */ + 0xc77fe000u, 0xbf801000u, 0x3eaaaaabu, 0xbeaaaaabu, +}; + +static void test_fill_copy_f32_patterns(void *dst, uint32_t n, uint32_t salt) { + uint8_t *bytes = dst; + const uint32_t n_patterns = + (uint32_t)(sizeof(test_copy_f32_patterns) / sizeof(test_copy_f32_patterns[0])); + for (uint32_t i = 0; i < n; i++) { + const uint32_t bits = test_copy_f32_patterns[(i + salt) % n_patterns]; + memcpy(bytes + (uint64_t)i * sizeof(bits), &bits, sizeof(bits)); + } +} +#endif + static uint16_t test_float_to_f16(float f) { union { float f; @@ -194,7 +257,8 @@ static float test_f16_to_f32(uint16_t h) { static void test_fill_q8_0_weights(uint8_t *weights, uint32_t in_dim, - uint32_t out_dim) { + uint32_t out_dim, + uint32_t seed) { const uint32_t blocks = in_dim / 32u; const uint64_t row_bytes = (uint64_t)blocks * 34u; for (uint32_t o = 0; o < out_dim; o++) { @@ -204,7 +268,8 @@ static void test_fill_q8_0_weights(uint8_t *weights, float amax = 0.0f; for (uint32_t i = 0; i < 32; i++) { const uint32_t k = b * 32u + i; - const int v = (int)((o * 17u + k * 23u + (o ^ k) * 3u) % 67u) - 33; + const int v = (int)((o * 17u + k * 23u + (o ^ k) * 3u + + seed * 29u + ((o + seed) ^ k) * 5u) % 67u) - 33; vals[i] = (float)v / 96.0f; float av = fabsf(vals[i]); if (av > amax) amax = av; @@ -408,7 +473,7 @@ static void test_metal_q8_0_prefill_matmul(void) { uint8_t *weights = weights_raw; memset(weights, 0, (size_t)weight_alloc); - test_fill_q8_0_weights(weights, in_dim, out_dim); + test_fill_q8_0_weights(weights, in_dim, out_dim, 0); ds4_gpu_tensor *x = ds4_gpu_tensor_alloc(x_bytes); ds4_gpu_tensor *out = ds4_gpu_tensor_alloc(out_bytes); @@ -486,10 +551,4217 @@ static void test_metal_q8_0_prefill_matmul(void) { free(weights_raw); } +static void test_metal_q8_0_decode_pair_exact_case( + uint32_t out0_dim, + uint32_t out1_dim, + uint32_t seed0, + uint32_t seed1) { + /* Exercise the Q-A/KV contract with unequal, odd output extents and + * independently page-aligned model ranges. The paired kernel must be + * bit-identical to two standalone decode matvec dispatches. */ + const uint32_t in_dim = 4096; + const uint64_t page = (uint64_t)getpagesize(); + const uint64_t row_bytes = (uint64_t)(in_dim / 32u) * 34u; + const uint64_t weight0_bytes = (uint64_t)out0_dim * row_bytes; + const uint64_t weight1_bytes = (uint64_t)out1_dim * row_bytes; + const uint64_t weight1_offset = test_round_up_u64(weight0_bytes, page); + const uint64_t weight_alloc = + test_round_up_u64(weight1_offset + weight1_bytes, page); + + void *weights_raw = NULL; + TEST_ASSERT(posix_memalign(&weights_raw, (size_t)page, (size_t)weight_alloc) == 0); + if (!weights_raw) return; + memset(weights_raw, 0, (size_t)weight_alloc); + test_fill_q8_0_weights((uint8_t *)weights_raw, in_dim, out0_dim, seed0); + test_fill_q8_0_weights((uint8_t *)weights_raw + weight1_offset, + in_dim, out1_dim, seed1); + + const uint64_t x_bytes = (uint64_t)in_dim * sizeof(float); + const uint64_t out0_bytes = (uint64_t)out0_dim * sizeof(float); + const uint64_t out1_bytes = (uint64_t)out1_dim * sizeof(float); + ds4_gpu_tensor *x = ds4_gpu_tensor_alloc(x_bytes); + ds4_gpu_tensor *ref0 = ds4_gpu_tensor_alloc(out0_bytes); + ds4_gpu_tensor *ref1 = ds4_gpu_tensor_alloc(out1_bytes); + ds4_gpu_tensor *pair0 = ds4_gpu_tensor_alloc(out0_bytes); + ds4_gpu_tensor *pair1 = ds4_gpu_tensor_alloc(out1_bytes); + TEST_ASSERT(x != NULL); + TEST_ASSERT(ref0 != NULL); + TEST_ASSERT(ref1 != NULL); + TEST_ASSERT(pair0 != NULL); + TEST_ASSERT(pair1 != NULL); + if (!x || !ref0 || !ref1 || !pair0 || !pair1) { + ds4_gpu_tensor_free(x); + ds4_gpu_tensor_free(ref0); + ds4_gpu_tensor_free(ref1); + ds4_gpu_tensor_free(pair0); + ds4_gpu_tensor_free(pair1); + free(weights_raw); + return; + } + + float *x_host = malloc((size_t)x_bytes); + float *ref0_host = malloc((size_t)out0_bytes); + float *ref1_host = malloc((size_t)out1_bytes); + float *pair0_host = malloc((size_t)out0_bytes); + float *pair1_host = malloc((size_t)out1_bytes); + TEST_ASSERT(x_host != NULL); + TEST_ASSERT(ref0_host != NULL); + TEST_ASSERT(ref1_host != NULL); + TEST_ASSERT(pair0_host != NULL); + TEST_ASSERT(pair1_host != NULL); + if (!x_host || !ref0_host || !ref1_host || !pair0_host || !pair1_host) { + free(x_host); + free(ref0_host); + free(ref1_host); + free(pair0_host); + free(pair1_host); + ds4_gpu_tensor_free(x); + ds4_gpu_tensor_free(ref0); + ds4_gpu_tensor_free(ref1); + ds4_gpu_tensor_free(pair0); + ds4_gpu_tensor_free(pair1); + free(weights_raw); + return; + } + + for (uint32_t i = 0; i < in_dim; i++) { + const int v = (int)((i * 29u + (i ^ (i >> 3u)) * 7u) % 127u) - 63; + x_host[i] = (float)v / 72.0f; + } + + TEST_ASSERT(ds4_gpu_tensor_write(x, 0, x_host, x_bytes) != 0); + TEST_ASSERT(ds4_gpu_set_model_map(weights_raw, weight_alloc) != 0); + ds4_gpu_set_quality(false); + TEST_ASSERT(ds4_gpu_matmul_q8_0_tensor(ref0, weights_raw, weight_alloc, 0, + in_dim, out0_dim, x, 1) != 0); + TEST_ASSERT(ds4_gpu_matmul_q8_0_tensor(ref1, weights_raw, weight_alloc, + weight1_offset, + in_dim, out1_dim, x, 1) != 0); + TEST_ASSERT(ds4_gpu_matmul_q8_0_pair_tensor(pair0, pair1, + weights_raw, weight_alloc, + 0, weight1_offset, + in_dim, out0_dim, out1_dim, + x, 1) != 0); + TEST_ASSERT(ds4_gpu_tensor_read(ref0, 0, ref0_host, out0_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read(ref1, 0, ref1_host, out1_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read(pair0, 0, pair0_host, out0_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read(pair1, 0, pair1_host, out1_bytes) != 0); + + uint32_t mismatch0 = 0; + uint32_t mismatch1 = 0; + float max_abs0 = 0.0f; + float max_abs1 = 0.0f; + for (uint32_t i = 0; i < out0_dim; i++) { + if (memcmp(&ref0_host[i], &pair0_host[i], sizeof(float)) != 0) mismatch0++; + const float err = fabsf(ref0_host[i] - pair0_host[i]); + if (err > max_abs0) max_abs0 = err; + } + for (uint32_t i = 0; i < out1_dim; i++) { + if (memcmp(&ref1_host[i], &pair1_host[i], sizeof(float)) != 0) mismatch1++; + const float err = fabsf(ref1_host[i] - pair1_host[i]); + if (err > max_abs1) max_abs1 = err; + } + if (mismatch0 != 0 || mismatch1 != 0) { + fprintf(stderr, + "ds4-test: paired Q8_0 exactness mismatches=%u/%u max_abs=%g, %u/%u max_abs=%g\n", + mismatch0, out0_dim, max_abs0, + mismatch1, out1_dim, max_abs1); + } + TEST_ASSERT(mismatch0 == 0); + TEST_ASSERT(mismatch1 == 0); + + free(x_host); + free(ref0_host); + free(ref1_host); + free(pair0_host); + free(pair1_host); + ds4_gpu_tensor_free(x); + ds4_gpu_tensor_free(ref0); + ds4_gpu_tensor_free(ref1); + ds4_gpu_tensor_free(pair0); + ds4_gpu_tensor_free(pair1); + free(weights_raw); +} + +static void test_metal_q8_0_decode_pair_exact(void) { + /* Cover both possible one-bank tail directions. Distinct seeds ensure a + * mistaken A-for-B weight binding cannot compare equal by construction. */ + test_metal_q8_0_decode_pair_exact_case(77, 19, 11, 97); + test_metal_q8_0_decode_pair_exact_case(19, 77, 23, 131); +} + +#if defined(__APPLE__) +static void test_metal_q8_0_output_nr4_exact_case( + uint32_t in_dim, + uint32_t out_dim, + uint32_t seed) { + const uint64_t page = (uint64_t)getpagesize(); + const uint64_t row_bytes = (uint64_t)(in_dim / 32u) * 34u; + const uint64_t weight_bytes = (uint64_t)out_dim * row_bytes; + const uint64_t weight_alloc = + test_round_up_u64(weight_bytes, page); + const uint64_t x_bytes = (uint64_t)in_dim * sizeof(float); + const uint64_t out_bytes = (uint64_t)out_dim * sizeof(float); + + void *weights_raw = NULL; + TEST_ASSERT(posix_memalign(&weights_raw, (size_t)page, + (size_t)weight_alloc) == 0); + ds4_gpu_tensor *x = ds4_gpu_tensor_alloc(x_bytes); + ds4_gpu_tensor *reference = ds4_gpu_tensor_alloc(out_bytes); + ds4_gpu_tensor *candidate = ds4_gpu_tensor_alloc(out_bytes); + float *x_host = malloc((size_t)x_bytes); + float *reference_host = malloc((size_t)out_bytes); + float *candidate_host = malloc((size_t)out_bytes); + TEST_ASSERT(weights_raw != NULL); + TEST_ASSERT(x != NULL); + TEST_ASSERT(reference != NULL); + TEST_ASSERT(candidate != NULL); + TEST_ASSERT(x_host != NULL); + TEST_ASSERT(reference_host != NULL); + TEST_ASSERT(candidate_host != NULL); + + const char *force_env = "DS4_METAL_ENABLE_OUTPUT_Q8_NR4"; + const char *disable_env = "DS4_METAL_DISABLE_M3_OUTPUT_Q8_NR4"; + char *saved_force = test_save_env(force_env); + char *saved_disable = test_save_env(disable_env); + test_float_compare_stats stats = {0}; + + const bool allocated = weights_raw && x && reference && candidate && + x_host && reference_host && candidate_host; + if (allocated) { + memset(weights_raw, 0, (size_t)weight_alloc); + test_fill_q8_0_weights( + (uint8_t *)weights_raw, in_dim, out_dim, seed); + for (uint32_t i = 0; i < in_dim; i++) { + const int value = + (int)((i * 29u + (i ^ (i >> 3u)) * 7u + + seed * 17u) % 127u) - 63; + x_host[i] = (float)value / 72.0f; + } + for (uint32_t i = 0; i < out_dim; i++) { + const uint32_t poison = 0x7fc00001u + (i & 0x3ffu); + memcpy(reference_host + i, &poison, sizeof(poison)); + memcpy(candidate_host + i, &poison, sizeof(poison)); + } + TEST_ASSERT(ds4_gpu_tensor_write(x, 0, x_host, x_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + reference, 0, reference_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + candidate, 0, candidate_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_set_model_map( + weights_raw, weight_alloc) != 0); + ds4_gpu_set_quality(false); + + TEST_ASSERT(unsetenv(force_env) == 0); + TEST_ASSERT(setenv(disable_env, "1", 1) == 0); + TEST_ASSERT(ds4_gpu_matmul_q8_0_tensor( + reference, weights_raw, weight_alloc, 0, + in_dim, out_dim, x, 1) != 0); + + TEST_ASSERT(setenv(force_env, "1", 1) == 0); + TEST_ASSERT(unsetenv(disable_env) == 0); + TEST_ASSERT(ds4_gpu_matmul_q8_0_tensor( + candidate, weights_raw, weight_alloc, 0, + in_dim, out_dim, x, 1) != 0); + + TEST_ASSERT(ds4_gpu_tensor_read( + reference, 0, reference_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + candidate, 0, candidate_host, out_bytes) != 0); + stats = test_compare_float_bits( + reference_host, candidate_host, out_dim); + } + + test_restore_env(force_env, saved_force); + test_restore_env(disable_env, saved_disable); + fprintf(stderr, + "ds4-test: output Q8 NR4 exact in=%u out=%u nsg=%u " + "mismatch=%zu/%u max_ulp=%u max_abs=%g\n", + in_dim, out_dim, out_dim > 65536u ? 8u : 4u, + stats.mismatch_count, out_dim, stats.max_ulp, stats.max_abs); + TEST_ASSERT(stats.mismatch_count == 0); + + free(candidate_host); + free(reference_host); + free(x_host); + ds4_gpu_tensor_free(candidate); + ds4_gpu_tensor_free(reference); + ds4_gpu_tensor_free(x); + free(weights_raw); +} + +static void test_metal_q8_0_output_nr4_exact(void) { + test_metal_q8_0_output_nr4_exact_case(4096, 68, 83); + test_metal_q8_0_output_nr4_exact_case(128, 65540, 89); +} + +static void test_metal_f16_compressor_pair_state_store_exact_case( + uint32_t width, + uint32_t ratio, + uint32_t pos, + uint32_t ape_type, + uint32_t seed) { + const uint32_t in_dim = 4096u; + const uint32_t coff = ratio == 4u ? 2u : 1u; + const uint32_t head_dim = width / coff; + const uint32_t state_rows = coff * ratio; + const bool emit = ((pos + 1u) % ratio) == 0u; + const uint64_t page = (uint64_t)getpagesize(); + const uint64_t weight_bytes = + (uint64_t)width * in_dim * sizeof(uint16_t); + const uint64_t score_weight_offset = + test_round_up_u64(weight_bytes, page); + const uint64_t ape_offset = test_round_up_u64( + score_weight_offset + weight_bytes, page); + const uint64_t ape_elem_bytes = ape_type == 1u ? 2u : 4u; + const uint64_t ape_bytes = (uint64_t)ratio * width * ape_elem_bytes; + const uint64_t norm_offset = + test_round_up_u64(ape_offset + ape_bytes, page); + const uint64_t model_bytes = test_round_up_u64( + norm_offset + (uint64_t)head_dim * sizeof(float), page); + const uint64_t x_bytes = (uint64_t)in_dim * sizeof(float); + const uint64_t out_bytes = (uint64_t)width * sizeof(float); + const uint64_t state_count = (uint64_t)state_rows * width; + const uint64_t state_bytes = state_count * sizeof(float); + const uint64_t comp_bytes = (uint64_t)head_dim * sizeof(float); + + void *model_raw = NULL; + TEST_ASSERT(posix_memalign(&model_raw, (size_t)page, + (size_t)model_bytes) == 0); + ds4_gpu_tensor *x = ds4_gpu_tensor_alloc(x_bytes); + ds4_gpu_tensor *ref_kv = ds4_gpu_tensor_alloc(out_bytes); + ds4_gpu_tensor *ref_score = ds4_gpu_tensor_alloc(out_bytes); + ds4_gpu_tensor *fused_kv = ds4_gpu_tensor_alloc(out_bytes); + ds4_gpu_tensor *fused_score = ds4_gpu_tensor_alloc(out_bytes); + ds4_gpu_tensor *ref_state_kv = ds4_gpu_tensor_alloc(state_bytes); + ds4_gpu_tensor *ref_state_score = ds4_gpu_tensor_alloc(state_bytes); + ds4_gpu_tensor *fused_state_kv = ds4_gpu_tensor_alloc(state_bytes); + ds4_gpu_tensor *fused_state_score = ds4_gpu_tensor_alloc(state_bytes); + ds4_gpu_tensor *ref_comp = ds4_gpu_tensor_alloc(comp_bytes); + ds4_gpu_tensor *fused_comp = ds4_gpu_tensor_alloc(comp_bytes); + + float *x_host = malloc((size_t)x_bytes); + float *ref_kv_host = malloc((size_t)out_bytes); + float *ref_score_host = malloc((size_t)out_bytes); + float *fused_kv_host = malloc((size_t)out_bytes); + float *fused_score_host = malloc((size_t)out_bytes); + float *ref_state_kv_host = malloc((size_t)state_bytes); + float *ref_state_score_host = malloc((size_t)state_bytes); + float *fused_state_kv_host = malloc((size_t)state_bytes); + float *fused_state_score_host = malloc((size_t)state_bytes); + float *ref_comp_host = malloc((size_t)comp_bytes); + float *fused_comp_host = malloc((size_t)comp_bytes); + + TEST_ASSERT(model_raw != NULL); + TEST_ASSERT(x != NULL); + TEST_ASSERT(ref_kv != NULL); + TEST_ASSERT(ref_score != NULL); + TEST_ASSERT(fused_kv != NULL); + TEST_ASSERT(fused_score != NULL); + TEST_ASSERT(ref_state_kv != NULL); + TEST_ASSERT(ref_state_score != NULL); + TEST_ASSERT(fused_state_kv != NULL); + TEST_ASSERT(fused_state_score != NULL); + TEST_ASSERT(ref_comp != NULL); + TEST_ASSERT(fused_comp != NULL); + TEST_ASSERT(x_host != NULL); + TEST_ASSERT(ref_kv_host != NULL); + TEST_ASSERT(ref_score_host != NULL); + TEST_ASSERT(fused_kv_host != NULL); + TEST_ASSERT(fused_score_host != NULL); + TEST_ASSERT(ref_state_kv_host != NULL); + TEST_ASSERT(ref_state_score_host != NULL); + TEST_ASSERT(fused_state_kv_host != NULL); + TEST_ASSERT(fused_state_score_host != NULL); + TEST_ASSERT(ref_comp_host != NULL); + TEST_ASSERT(fused_comp_host != NULL); + + const bool allocated = model_raw && x && ref_kv && ref_score && fused_kv && + fused_score && ref_state_kv && ref_state_score && fused_state_kv && + fused_state_score && ref_comp && fused_comp && x_host && ref_kv_host && + ref_score_host && fused_kv_host && fused_score_host && + ref_state_kv_host && ref_state_score_host && fused_state_kv_host && + fused_state_score_host && ref_comp_host && fused_comp_host; + + const char *force_env = + "DS4_METAL_ENABLE_COMPRESSOR_PAIR_STATE_STORE"; + const char *disable_env = + "DS4_METAL_DISABLE_M3_COMPRESSOR_PAIR_STATE_STORE"; + const char *pair_disable_env = + "DS4_METAL_DISABLE_COMPRESSOR_PAIR_PROJ"; + const char *store_disable_env = + "DS4_METAL_DISABLE_COMPRESSOR_STORE_ONE"; + char *saved_force = test_save_env(force_env); + char *saved_disable = test_save_env(disable_env); + char *saved_pair_disable = test_save_env(pair_disable_env); + char *saved_store_disable = test_save_env(store_disable_env); + + test_float_compare_stats kv_stats = {0}; + test_float_compare_stats score_stats = {0}; + test_float_compare_stats state_kv_stats = {0}; + test_float_compare_stats state_score_stats = {0}; + test_float_compare_stats comp_stats = {0}; + + if (allocated) { + memset(model_raw, 0, (size_t)model_bytes); + uint16_t *kv_weights = model_raw; + uint16_t *score_weights = + (uint16_t *)((uint8_t *)model_raw + score_weight_offset); + for (uint32_t o = 0; o < width; o++) { + for (uint32_t i = 0; i < in_dim; i++) { + const int kv_value = + (int)((o * 17u + i * 23u + (o ^ i) * 3u + + seed * 29u) % 67u) - 33; + const int score_value = + (int)((o * 31u + i * 11u + (o ^ (i >> 2u)) * 5u + + seed * 19u) % 71u) - 35; + const uint64_t wi = (uint64_t)o * in_dim + i; + kv_weights[wi] = test_float_to_f16( + (float)kv_value / 96.0f); + score_weights[wi] = test_float_to_f16( + (float)score_value / 104.0f); + } + } + + if (ape_type == 1u) { + uint16_t *ape = + (uint16_t *)((uint8_t *)model_raw + ape_offset); + for (uint64_t i = 0; i < (uint64_t)ratio * width; i++) { + const int value = + (int)((i * 13u + (i ^ (i >> 3u)) * 7u + + seed * 17u) % 61u) - 30; + ape[i] = test_float_to_f16((float)value / 80.0f); + } + } else { + float *ape = (float *)((uint8_t *)model_raw + ape_offset); + for (uint64_t i = 0; i < (uint64_t)ratio * width; i++) { + const int value = + (int)((i * 13u + (i ^ (i >> 3u)) * 7u + + seed * 17u) % 61u) - 30; + ape[i] = (float)value / 80.0f; + } + } + float *norm = (float *)((uint8_t *)model_raw + norm_offset); + for (uint32_t i = 0; i < head_dim; i++) { + norm[i] = 0.75f + + (float)((i * 7u + seed * 3u) % 23u) / 64.0f; + } + + for (uint32_t i = 0; i < in_dim; i++) { + const int value = + (int)((i * 29u + (i ^ (i >> 4u)) * 9u + + seed * 11u) % 127u) - 63; + x_host[i] = (float)value / 88.0f; + } + for (uint32_t i = 0; i < width; i++) { + const uint32_t poison = 0x7fc00001u + (i & 0x3ffu); + memcpy(ref_kv_host + i, &poison, sizeof(poison)); + memcpy(ref_score_host + i, &poison, sizeof(poison)); + memcpy(fused_kv_host + i, &poison, sizeof(poison)); + memcpy(fused_score_host + i, &poison, sizeof(poison)); + } + for (uint64_t i = 0; i < state_count; i++) { + const int kv_value = + (int)((i * 5u + seed * 13u) % 97u) - 48; + const int score_value = + (int)((i * 7u + seed * 5u) % 101u) - 50; + ref_state_kv_host[i] = (float)kv_value / 64.0f; + fused_state_kv_host[i] = ref_state_kv_host[i]; + ref_state_score_host[i] = (float)score_value / 72.0f; + fused_state_score_host[i] = ref_state_score_host[i]; + } + for (uint32_t i = 0; i < head_dim; i++) { + const uint32_t poison = 0x7fc01001u + (i & 0x3ffu); + memcpy(ref_comp_host + i, &poison, sizeof(poison)); + memcpy(fused_comp_host + i, &poison, sizeof(poison)); + } + + TEST_ASSERT(ds4_gpu_tensor_write(x, 0, x_host, x_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + ref_kv, 0, ref_kv_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + ref_score, 0, ref_score_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_kv, 0, fused_kv_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_score, 0, fused_score_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + ref_state_kv, 0, ref_state_kv_host, state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + ref_state_score, 0, ref_state_score_host, + state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_state_kv, 0, fused_state_kv_host, + state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_state_score, 0, fused_state_score_host, + state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + ref_comp, 0, ref_comp_host, comp_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_comp, 0, fused_comp_host, comp_bytes) != 0); + TEST_ASSERT(ds4_gpu_set_model_map(model_raw, model_bytes) != 0); + ds4_gpu_set_quality(false); + + TEST_ASSERT(setenv(force_env, "1", 1) == 0); + TEST_ASSERT(unsetenv(disable_env) == 0); + TEST_ASSERT(unsetenv(pair_disable_env) == 0); + TEST_ASSERT(unsetenv(store_disable_env) == 0); + + TEST_ASSERT(ds4_gpu_matmul_f16_pair_tensor( + ref_kv, ref_score, model_raw, model_bytes, + 0, score_weight_offset, in_dim, width, x, 1) != 0); + TEST_ASSERT(ds4_gpu_compressor_update_tensor( + ref_kv, ref_score, ref_state_kv, ref_state_score, + ref_comp, model_raw, model_bytes, ape_offset, ape_type, + norm_offset, 0, head_dim, ratio, pos, 0, 0, 0, + 10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, + 1.0e-6f, false) != 0); + + TEST_ASSERT(ds4_gpu_matmul_f16_pair_compressor_store_tensor( + fused_kv, fused_score, + fused_state_kv, fused_state_score, + model_raw, model_bytes, 0, score_weight_offset, + ape_offset, ape_type, in_dim, width, x, + ratio, pos) == 1); + TEST_ASSERT(ds4_gpu_compressor_update_tensor( + fused_kv, fused_score, + fused_state_kv, fused_state_score, + fused_comp, model_raw, model_bytes, + ape_offset, ape_type, norm_offset, 0, + head_dim, ratio, pos, 0, 0, 0, + 10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, + 1.0e-6f, true) != 0); + + TEST_ASSERT(ds4_gpu_tensor_read( + ref_kv, 0, ref_kv_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + ref_score, 0, ref_score_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + fused_kv, 0, fused_kv_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + fused_score, 0, fused_score_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + ref_state_kv, 0, ref_state_kv_host, state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + ref_state_score, 0, ref_state_score_host, + state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + fused_state_kv, 0, fused_state_kv_host, + state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + fused_state_score, 0, fused_state_score_host, + state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + ref_comp, 0, ref_comp_host, comp_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + fused_comp, 0, fused_comp_host, comp_bytes) != 0); + + kv_stats = test_compare_float_bits( + ref_kv_host, fused_kv_host, width); + score_stats = test_compare_float_bits( + ref_score_host, fused_score_host, width); + state_kv_stats = test_compare_float_bits( + ref_state_kv_host, fused_state_kv_host, (size_t)state_count); + state_score_stats = test_compare_float_bits( + ref_state_score_host, fused_state_score_host, + (size_t)state_count); + comp_stats = test_compare_float_bits( + ref_comp_host, fused_comp_host, head_dim); + } + + test_restore_env(force_env, saved_force); + test_restore_env(disable_env, saved_disable); + test_restore_env(pair_disable_env, saved_pair_disable); + test_restore_env(store_disable_env, saved_store_disable); + + fprintf(stderr, + "ds4-test: compressor pair state-store exact width=%u ratio=%u " + "pos=%u emit=%u ape=%s proj=%zu/%zu state=%zu/%zu " + "comp=%zu max_ulp=%u/%u/%u/%u/%u\n", + width, ratio, pos, emit ? 1u : 0u, + ape_type == 1u ? "f16" : "f32", + kv_stats.mismatch_count, score_stats.mismatch_count, + state_kv_stats.mismatch_count, + state_score_stats.mismatch_count, + comp_stats.mismatch_count, + kv_stats.max_ulp, score_stats.max_ulp, + state_kv_stats.max_ulp, state_score_stats.max_ulp, + comp_stats.max_ulp); + TEST_ASSERT(kv_stats.mismatch_count == 0); + TEST_ASSERT(score_stats.mismatch_count == 0); + TEST_ASSERT(state_kv_stats.mismatch_count == 0); + TEST_ASSERT(state_score_stats.mismatch_count == 0); + TEST_ASSERT(comp_stats.mismatch_count == 0); + + free(fused_comp_host); + free(ref_comp_host); + free(fused_state_score_host); + free(fused_state_kv_host); + free(ref_state_score_host); + free(ref_state_kv_host); + free(fused_score_host); + free(fused_kv_host); + free(ref_score_host); + free(ref_kv_host); + free(x_host); + ds4_gpu_tensor_free(fused_comp); + ds4_gpu_tensor_free(ref_comp); + ds4_gpu_tensor_free(fused_state_score); + ds4_gpu_tensor_free(fused_state_kv); + ds4_gpu_tensor_free(ref_state_score); + ds4_gpu_tensor_free(ref_state_kv); + ds4_gpu_tensor_free(fused_score); + ds4_gpu_tensor_free(fused_kv); + ds4_gpu_tensor_free(ref_score); + ds4_gpu_tensor_free(ref_kv); + ds4_gpu_tensor_free(x); + free(model_raw); +} + +static void test_metal_f16_compressor_pair_state_store_exact(void) { + test_metal_f16_compressor_pair_state_store_exact_case( + 256, 4, 8, 0, 17); + test_metal_f16_compressor_pair_state_store_exact_case( + 1024, 4, 11, 1, 29); + test_metal_f16_compressor_pair_state_store_exact_case( + 512, 128, 255, 1, 43); +} + +static void test_metal_compressor_ape_add_exact_case( + uint32_t head_dim, + uint32_t ratio, + uint32_t pos0, + uint32_t n_tokens, + uint32_t ape_type, + uint32_t seed, + bool test_pack_fusion) { + const uint32_t coff = ratio == 4u ? 2u : 1u; + const uint32_t width = coff * head_dim; + const uint32_t state_rows = coff * ratio; + const uint32_t n_comp = n_tokens / ratio; + const uint64_t input_count = (uint64_t)n_tokens * width; + const uint64_t state_count = (uint64_t)state_rows * width; + const uint64_t comp_count = (uint64_t)n_comp * head_dim; + const uint64_t input_bytes = input_count * sizeof(float); + const uint64_t state_bytes = state_count * sizeof(float); + const uint64_t comp_bytes = comp_count * sizeof(float); + const uint64_t ape_elem_bytes = ape_type == 1u ? 2u : 4u; + const uint64_t ape_bytes = (uint64_t)width * ratio * ape_elem_bytes; + const uint64_t page = (uint64_t)getpagesize(); + const uint64_t norm_offset = test_round_up_u64(ape_bytes, page); + const uint64_t model_bytes = test_round_up_u64( + norm_offset + (uint64_t)head_dim * sizeof(float), page); + + ds4_gpu_tensor *kv = ds4_gpu_tensor_alloc(input_bytes); + ds4_gpu_tensor *sc = ds4_gpu_tensor_alloc(input_bytes); + ds4_gpu_tensor *ref_comp = ds4_gpu_tensor_alloc(comp_bytes); + ds4_gpu_tensor *fused_comp = ds4_gpu_tensor_alloc(comp_bytes); + ds4_gpu_tensor *ref_state_kv = ds4_gpu_tensor_alloc(state_bytes); + ds4_gpu_tensor *ref_state_score = ds4_gpu_tensor_alloc(state_bytes); + ds4_gpu_tensor *fused_state_kv = ds4_gpu_tensor_alloc(state_bytes); + ds4_gpu_tensor *fused_state_score = ds4_gpu_tensor_alloc(state_bytes); + TEST_ASSERT(kv != NULL); + TEST_ASSERT(sc != NULL); + TEST_ASSERT(ref_comp != NULL); + TEST_ASSERT(fused_comp != NULL); + TEST_ASSERT(ref_state_kv != NULL); + TEST_ASSERT(ref_state_score != NULL); + TEST_ASSERT(fused_state_kv != NULL); + TEST_ASSERT(fused_state_score != NULL); + + float *kv_host = malloc((size_t)input_bytes); + float *sc_host = malloc((size_t)input_bytes); + float *ref_comp_host = malloc((size_t)comp_bytes); + float *fused_comp_host = malloc((size_t)comp_bytes); + float *ref_state_kv_host = malloc((size_t)state_bytes); + float *ref_state_score_host = malloc((size_t)state_bytes); + float *fused_state_kv_host = malloc((size_t)state_bytes); + float *fused_state_score_host = malloc((size_t)state_bytes); + const uint64_t poison_count = input_count > state_count ? + input_count : state_count; + float *poison_host = test_pack_fusion ? + malloc((size_t)(poison_count * sizeof(float))) : NULL; + TEST_ASSERT(kv_host != NULL); + TEST_ASSERT(sc_host != NULL); + TEST_ASSERT(ref_comp_host != NULL); + TEST_ASSERT(fused_comp_host != NULL); + TEST_ASSERT(ref_state_kv_host != NULL); + TEST_ASSERT(ref_state_score_host != NULL); + TEST_ASSERT(fused_state_kv_host != NULL); + TEST_ASSERT(fused_state_score_host != NULL); + TEST_ASSERT(!test_pack_fusion || poison_host != NULL); + + void *model_raw = NULL; + TEST_ASSERT(posix_memalign(&model_raw, (size_t)page, + (size_t)model_bytes) == 0); + const bool allocated = kv && sc && ref_comp && fused_comp && ref_state_kv && + ref_state_score && fused_state_kv && fused_state_score && kv_host && + sc_host && ref_comp_host && fused_comp_host && ref_state_kv_host && + ref_state_score_host && fused_state_kv_host && fused_state_score_host && + (!test_pack_fusion || poison_host) && model_raw; + const char *force_env = "DS4_METAL_ENABLE_COMPRESSOR_APE_ADD"; + const char *disable_env = "DS4_METAL_DISABLE_M3_COMPRESSOR_APE_ADD"; + const char *pack_force_env = + "DS4_METAL_ENABLE_COMPRESSOR_RATIO4_PACK_FUSION"; + const char *pack_disable_env = + "DS4_METAL_DISABLE_M3_COMPRESSOR_RATIO4_PACK_FUSION"; + char *saved_force = test_save_env(force_env); + char *saved_disable = test_save_env(disable_env); + char *saved_pack_force = test_save_env(pack_force_env); + char *saved_pack_disable = test_save_env(pack_disable_env); + test_float_compare_stats comp_stats = {0}; + test_float_compare_stats state_kv_stats = {0}; + test_float_compare_stats state_score_stats = {0}; + + if (allocated) { + memset(model_raw, 0, (size_t)model_bytes); + if (ape_type == 1u) { + uint16_t *ape = model_raw; + for (uint64_t i = 0; i < (uint64_t)width * ratio; i++) { + const int value = (int)((i * 17u + (i ^ (i >> 4u)) * 5u + + seed * 13u) % 127u) - 63; + ape[i] = test_float_to_f16((float)value / 80.0f); + } + } else { + float *ape = model_raw; + for (uint64_t i = 0; i < (uint64_t)width * ratio; i++) { + const int value = (int)((i * 17u + (i ^ (i >> 4u)) * 5u + + seed * 13u) % 127u) - 63; + ape[i] = (float)value / 80.0f; + } + } + float *norm = (float *)((uint8_t *)model_raw + norm_offset); + for (uint32_t i = 0; i < head_dim; i++) { + norm[i] = 0.75f + (float)((i * 7u + seed) % 19u) / 64.0f; + } + for (uint64_t i = 0; i < input_count; i++) { + const int kv_value = (int)((i * 29u + (i ^ (i >> 5u)) * 3u + + seed * 23u) % 181u) - 90; + const int sc_value = (int)((i * 31u + (i ^ (i >> 3u)) * 11u + + seed * 17u) % 173u) - 86; + kv_host[i] = (float)kv_value / 112.0f; + sc_host[i] = (float)sc_value / 96.0f; + } + kv_host[0] = -0.0f; + sc_host[0] = -0.0f; + + // Exercise exact-add edge values in the first active APE row. Using + // bit patterns avoids host fast-math rewriting signed zeros or + // subnormals before the legacy and fused Metal paths see them. + const uint32_t cutoff = (n_tokens / ratio) * ratio; + const uint32_t edge_token = cutoff < n_tokens ? cutoff : cutoff - ratio; + const uint64_t edge_ape = + (uint64_t)((pos0 + edge_token) % ratio) * width; + float *edge_score = sc_host + (uint64_t)edge_token * width; + static const uint16_t edge_f16[] = { + 0x0000u, 0x8000u, 0x0001u, 0x8001u, + 0x3c00u, 0xbc00u, 0x7bffu, 0xfbffu, + }; + static const uint32_t edge_f32[] = { + 0x00000000u, 0x80000000u, 0x00000001u, 0x80000001u, + 0x3f800000u, 0xbf800000u, 0x7f7fffffu, 0xff7fffffu, + }; + static const uint32_t edge_score_f16[] = { + 0x80000000u, 0x00000000u, 0x00000000u, 0x80000000u, + 0xbf800000u, 0x3f800000u, 0xc77fe000u, 0x477fe000u, + }; + static const uint32_t edge_score_f32[] = { + 0x80000000u, 0x00000000u, 0x00000000u, 0x80000000u, + 0xbf800000u, 0x3f800000u, 0xff7fffffu, 0x7f7fffffu, + }; + if (ape_type == 1u) { + uint16_t *ape = model_raw; + memcpy(ape + edge_ape, edge_f16, sizeof(edge_f16)); + for (uint32_t i = 0; i < 8u; i++) { + memcpy(edge_score + i, edge_score_f16 + i, sizeof(uint32_t)); + } + } else { + uint32_t *ape = model_raw; + memcpy(ape + edge_ape, edge_f32, sizeof(edge_f32)); + for (uint32_t i = 0; i < 8u; i++) { + memcpy(edge_score + i, edge_score_f32 + i, sizeof(uint32_t)); + } + } + for (uint64_t i = 0; i < state_count; i++) { + ref_state_kv_host[i] = 1234.0f; + fused_state_kv_host[i] = 1234.0f; + ref_state_score_host[i] = -1234.0f; + fused_state_score_host[i] = -1234.0f; + } + + TEST_ASSERT(ds4_gpu_tensor_write(kv, 0, kv_host, input_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write(sc, 0, sc_host, input_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + ref_state_kv, 0, ref_state_kv_host, state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_state_kv, 0, fused_state_kv_host, state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + ref_state_score, 0, ref_state_score_host, state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_state_score, 0, fused_state_score_host, state_bytes) != 0); + TEST_ASSERT(ds4_gpu_set_model_map(model_raw, model_bytes) != 0); + ds4_gpu_set_quality(false); + + if (test_pack_fusion) { + TEST_ASSERT(setenv(force_env, "1", 1) == 0); + TEST_ASSERT(unsetenv(disable_env) == 0); + TEST_ASSERT(unsetenv(pack_force_env) == 0); + } else { + TEST_ASSERT(unsetenv(force_env) == 0); + TEST_ASSERT(setenv(disable_env, "1", 1) == 0); + } + TEST_ASSERT(setenv(pack_disable_env, "1", 1) == 0); + TEST_ASSERT(ds4_gpu_compressor_prefill_tensor( + ref_comp, ref_state_kv, ref_state_score, kv, sc, + model_raw, model_bytes, 0, ape_type, norm_offset, 0, + head_dim, ratio, pos0, n_tokens, 0, 0, false, + 10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, 1.0e-6f) != 0); + + if (test_pack_fusion) { + // Overwrite every persistent pack cell with qNaN payloads through + // the legacy replay path. This makes an omitted candidate write + // observable even for plane-zero padding, which normal legacy + // packing would otherwise leave at the correct 0/-inf values. + for (uint64_t i = 0; i < poison_count; i++) { + const uint32_t bits = + 0x7fc00001u + (uint32_t)(i & 0x3ffu); + memcpy(poison_host + i, &bits, sizeof(bits)); + } + TEST_ASSERT(ds4_gpu_tensor_write( + kv, 0, poison_host, input_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + sc, 0, poison_host, input_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_state_kv, 0, poison_host, + state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_state_score, 0, poison_host, + state_bytes) != 0); + TEST_ASSERT(ds4_gpu_compressor_prefill_ratio4_replay_tensor( + fused_comp, fused_state_kv, fused_state_score, kv, sc, + model_raw, model_bytes, 0, ape_type, norm_offset, 0, + head_dim, 0, n_comp * ratio, 0, 0, false, + 10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, 1.0e-6f) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + kv, 0, kv_host, input_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + sc, 0, sc_host, input_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_state_kv, 0, fused_state_kv_host, + state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_state_score, 0, fused_state_score_host, + state_bytes) != 0); + TEST_ASSERT(setenv(pack_force_env, "1", 1) == 0); + TEST_ASSERT(unsetenv(pack_disable_env) == 0); + } else { + TEST_ASSERT(setenv(force_env, "1", 1) == 0); + TEST_ASSERT(unsetenv(disable_env) == 0); + } + TEST_ASSERT(ds4_gpu_compressor_prefill_tensor( + fused_comp, fused_state_kv, fused_state_score, kv, sc, + model_raw, model_bytes, 0, ape_type, norm_offset, 0, + head_dim, ratio, pos0, n_tokens, 0, 0, false, + 10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, 1.0e-6f) != 0); + + if (test_pack_fusion) { + TEST_ASSERT(ds4_gpu_tensor_read( + kv, 0, poison_host, input_bytes) != 0); + TEST_ASSERT(memcmp(kv_host, poison_host, + (size_t)input_bytes) == 0); + TEST_ASSERT(ds4_gpu_tensor_read( + sc, 0, poison_host, input_bytes) != 0); + TEST_ASSERT(memcmp(sc_host, poison_host, + (size_t)input_bytes) == 0); + } + TEST_ASSERT(ds4_gpu_tensor_read( + ref_comp, 0, ref_comp_host, comp_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + fused_comp, 0, fused_comp_host, comp_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + ref_state_kv, 0, ref_state_kv_host, state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + fused_state_kv, 0, fused_state_kv_host, state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + ref_state_score, 0, ref_state_score_host, state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + fused_state_score, 0, fused_state_score_host, + state_bytes) != 0); + + comp_stats = test_compare_float_bits( + ref_comp_host, fused_comp_host, (size_t)comp_count); + state_kv_stats = test_compare_float_bits( + ref_state_kv_host, fused_state_kv_host, (size_t)state_count); + state_score_stats = test_compare_float_bits( + ref_state_score_host, fused_state_score_host, (size_t)state_count); + } + + test_restore_env(force_env, saved_force); + test_restore_env(disable_env, saved_disable); + test_restore_env(pack_force_env, saved_pack_force); + test_restore_env(pack_disable_env, saved_pack_disable); + fprintf(stderr, + "ds4-test: compressor %s exact head=%u ratio=%u pos=%u " + "tokens=%u ape=%s comp=%zu/%llu state_kv=%zu/%llu " + "state_score=%zu/%llu max_ulp=%u/%u/%u\n", + test_pack_fusion ? "ratio4 pack" : "APE add", + head_dim, ratio, pos0, n_tokens, ape_type == 1u ? "f16" : "f32", + comp_stats.mismatch_count, (unsigned long long)comp_count, + state_kv_stats.mismatch_count, (unsigned long long)state_count, + state_score_stats.mismatch_count, (unsigned long long)state_count, + comp_stats.max_ulp, state_kv_stats.max_ulp, + state_score_stats.max_ulp); + TEST_ASSERT(comp_stats.mismatch_count == 0); + TEST_ASSERT(state_kv_stats.mismatch_count == 0); + TEST_ASSERT(state_score_stats.mismatch_count == 0); + + free(model_raw); + free(poison_host); + free(fused_state_score_host); + free(fused_state_kv_host); + free(ref_state_score_host); + free(ref_state_kv_host); + free(fused_comp_host); + free(ref_comp_host); + free(sc_host); + free(kv_host); + ds4_gpu_tensor_free(fused_state_score); + ds4_gpu_tensor_free(fused_state_kv); + ds4_gpu_tensor_free(ref_state_score); + ds4_gpu_tensor_free(ref_state_kv); + ds4_gpu_tensor_free(fused_comp); + ds4_gpu_tensor_free(ref_comp); + ds4_gpu_tensor_free(sc); + ds4_gpu_tensor_free(kv); +} + +static void test_metal_compressor_ape_add_exact(void) { + test_metal_compressor_ape_add_exact_case(128, 4, 3, 16, 1, 7, false); + test_metal_compressor_ape_add_exact_case(512, 4, 0, 16, 1, 13, false); + test_metal_compressor_ape_add_exact_case(512, 4, 1, 14, 0, 19, false); + test_metal_compressor_ape_add_exact_case(512, 128, 127, 257, 1, 31, false); +} + +static void test_metal_compressor_ratio4_pack_exact(void) { + test_metal_compressor_ape_add_exact_case(128, 4, 0, 4, 1, 37, true); + test_metal_compressor_ape_add_exact_case(512, 4, 0, 8, 1, 41, true); + test_metal_compressor_ape_add_exact_case(512, 4, 1, 14, 0, 43, true); +} + +static void test_metal_compressor_ratio4_replay_pack_exact_case( + uint32_t head_dim, + uint32_t n_tokens, + uint32_t seed) { + const uint32_t width = 2u * head_dim; + const uint32_t n_comp = n_tokens / 4u; + const uint64_t input_count = (uint64_t)n_tokens * width; + const uint64_t state_count = (uint64_t)8u * width; + const uint64_t comp_count = (uint64_t)n_comp * head_dim; + const uint64_t input_bytes = input_count * sizeof(float); + const uint64_t state_bytes = state_count * sizeof(float); + const uint64_t comp_bytes = comp_count * sizeof(float); + const uint64_t ape_bytes = (uint64_t)4u * width * sizeof(uint16_t); + const uint64_t page = (uint64_t)getpagesize(); + const uint64_t norm_offset = test_round_up_u64(ape_bytes, page); + const uint64_t model_bytes = test_round_up_u64( + norm_offset + (uint64_t)head_dim * sizeof(float), page); + + ds4_gpu_tensor *kv = ds4_gpu_tensor_alloc(input_bytes); + ds4_gpu_tensor *sc = ds4_gpu_tensor_alloc(input_bytes); + ds4_gpu_tensor *ref_comp = ds4_gpu_tensor_alloc(comp_bytes); + ds4_gpu_tensor *fused_comp = ds4_gpu_tensor_alloc(comp_bytes); + ds4_gpu_tensor *ref_state_kv = ds4_gpu_tensor_alloc(state_bytes); + ds4_gpu_tensor *ref_state_score = ds4_gpu_tensor_alloc(state_bytes); + ds4_gpu_tensor *fused_state_kv = ds4_gpu_tensor_alloc(state_bytes); + ds4_gpu_tensor *fused_state_score = ds4_gpu_tensor_alloc(state_bytes); + float *kv_host = malloc((size_t)input_bytes); + float *sc_host = malloc((size_t)input_bytes); + float *state_kv_host = malloc((size_t)state_bytes); + float *state_score_host = malloc((size_t)state_bytes); + float *source_after_host = malloc((size_t)input_bytes); + float *ref_comp_host = malloc((size_t)comp_bytes); + float *fused_comp_host = malloc((size_t)comp_bytes); + float *ref_state_kv_host = malloc((size_t)state_bytes); + float *fused_state_kv_host = malloc((size_t)state_bytes); + float *ref_state_score_host = malloc((size_t)state_bytes); + float *fused_state_score_host = malloc((size_t)state_bytes); + void *model_raw = NULL; + TEST_ASSERT(posix_memalign( + &model_raw, (size_t)page, (size_t)model_bytes) == 0); + + TEST_ASSERT(kv != NULL); + TEST_ASSERT(sc != NULL); + TEST_ASSERT(ref_comp != NULL); + TEST_ASSERT(fused_comp != NULL); + TEST_ASSERT(ref_state_kv != NULL); + TEST_ASSERT(ref_state_score != NULL); + TEST_ASSERT(fused_state_kv != NULL); + TEST_ASSERT(fused_state_score != NULL); + TEST_ASSERT(kv_host != NULL); + TEST_ASSERT(sc_host != NULL); + TEST_ASSERT(state_kv_host != NULL); + TEST_ASSERT(state_score_host != NULL); + TEST_ASSERT(source_after_host != NULL); + TEST_ASSERT(ref_comp_host != NULL); + TEST_ASSERT(fused_comp_host != NULL); + TEST_ASSERT(ref_state_kv_host != NULL); + TEST_ASSERT(fused_state_kv_host != NULL); + TEST_ASSERT(ref_state_score_host != NULL); + TEST_ASSERT(fused_state_score_host != NULL); + TEST_ASSERT(model_raw != NULL); + + const char *ape_force_env = "DS4_METAL_ENABLE_COMPRESSOR_APE_ADD"; + const char *ape_disable_env = "DS4_METAL_DISABLE_M3_COMPRESSOR_APE_ADD"; + const char *pack_force_env = + "DS4_METAL_ENABLE_COMPRESSOR_RATIO4_PACK_FUSION"; + const char *pack_disable_env = + "DS4_METAL_DISABLE_M3_COMPRESSOR_RATIO4_PACK_FUSION"; + char *saved_ape_force = test_save_env(ape_force_env); + char *saved_ape_disable = test_save_env(ape_disable_env); + char *saved_pack_force = test_save_env(pack_force_env); + char *saved_pack_disable = test_save_env(pack_disable_env); + test_float_compare_stats comp_stats = {0}; + test_float_compare_stats state_kv_stats = {0}; + test_float_compare_stats state_score_stats = {0}; + + const bool allocated = kv && sc && ref_comp && fused_comp && + ref_state_kv && ref_state_score && fused_state_kv && + fused_state_score && kv_host && sc_host && state_kv_host && + state_score_host && source_after_host && ref_comp_host && fused_comp_host && + ref_state_kv_host && fused_state_kv_host && ref_state_score_host && + fused_state_score_host && model_raw; + if (allocated) { + memset(model_raw, 0, (size_t)model_bytes); + uint16_t *ape = model_raw; + for (uint64_t i = 0; i < (uint64_t)4u * width; i++) { + const int value = + (int)((i * 19u + (i ^ (i >> 3u)) * 7u + seed * 11u) % + 113u) - 56; + ape[i] = test_float_to_f16((float)value / 72.0f); + } + float *norm = (float *)((uint8_t *)model_raw + norm_offset); + for (uint32_t i = 0; i < head_dim; i++) { + norm[i] = 0.875f + (float)((i * 5u + seed) % 17u) / 64.0f; + } + for (uint64_t i = 0; i < input_count; i++) { + const int kv_value = + (int)((i * 31u + (i ^ (i >> 4u)) * 5u + seed * 13u) % + 193u) - 96; + const int sc_value = + (int)((i * 37u + (i ^ (i >> 5u)) * 9u + seed * 17u) % + 181u) - 90; + kv_host[i] = (float)kv_value / 104.0f; + sc_host[i] = (float)sc_value / 88.0f; + } + for (uint64_t i = 0; i < state_count; i++) { + const int kv_value = + (int)((i * 23u + (i ^ (i >> 2u)) * 3u + seed * 29u) % + 167u) - 83; + const int sc_value = + (int)((i * 41u + (i ^ (i >> 6u)) * 11u + seed * 7u) % + 157u) - 78; + state_kv_host[i] = (float)kv_value / 80.0f; + state_score_host[i] = (float)sc_value / 92.0f; + } + const uint32_t negative_zero = 0x80000000u; + memcpy(kv_host, &negative_zero, sizeof(negative_zero)); + memcpy(state_kv_host, &negative_zero, sizeof(negative_zero)); + + TEST_ASSERT(ds4_gpu_tensor_write(kv, 0, kv_host, input_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write(sc, 0, sc_host, input_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + ref_state_kv, 0, state_kv_host, state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_state_kv, 0, state_kv_host, state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + ref_state_score, 0, state_score_host, state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_state_score, 0, state_score_host, state_bytes) != 0); + TEST_ASSERT(ds4_gpu_set_model_map(model_raw, model_bytes) != 0); + ds4_gpu_set_quality(false); + TEST_ASSERT(setenv(ape_force_env, "1", 1) == 0); + TEST_ASSERT(unsetenv(ape_disable_env) == 0); + TEST_ASSERT(unsetenv(pack_force_env) == 0); + TEST_ASSERT(setenv(pack_disable_env, "1", 1) == 0); + + TEST_ASSERT(ds4_gpu_compressor_prefill_ratio4_replay_tensor( + ref_comp, ref_state_kv, ref_state_score, kv, sc, + model_raw, model_bytes, 0, 1, norm_offset, 0, + head_dim, 0, n_tokens, 0, 0, false, + 10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, 1.0e-6f) != 0); + + // Poison every persistent pack cell through the legacy full-fill path + // so a missing candidate write cannot inherit the reference value. + for (uint64_t i = 0; i < input_count; i++) { + const uint32_t bits = 0x7fc00001u + (uint32_t)(i & 0x3ffu); + memcpy(kv_host + i, &bits, sizeof(bits)); + memcpy(sc_host + i, &bits, sizeof(bits)); + } + for (uint64_t i = 0; i < state_count; i++) { + const uint32_t bits = 0x7fc00401u + (uint32_t)(i & 0x3ffu); + memcpy(ref_state_kv_host + i, &bits, sizeof(bits)); + memcpy(ref_state_score_host + i, &bits, sizeof(bits)); + } + TEST_ASSERT(ds4_gpu_tensor_write(kv, 0, kv_host, input_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write(sc, 0, sc_host, input_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_state_kv, 0, ref_state_kv_host, state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_state_score, 0, ref_state_score_host, + state_bytes) != 0); + TEST_ASSERT(ds4_gpu_compressor_prefill_ratio4_replay_tensor( + fused_comp, fused_state_kv, fused_state_score, kv, sc, + model_raw, model_bytes, 0, 1, norm_offset, 0, + head_dim, 0, n_tokens, 0, 0, false, + 10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, 1.0e-6f) != 0); + + // Restore finite sources and the candidate's original replay state. + for (uint64_t i = 0; i < input_count; i++) { + const int kv_value = + (int)((i * 31u + (i ^ (i >> 4u)) * 5u + seed * 13u) % + 193u) - 96; + const int sc_value = + (int)((i * 37u + (i ^ (i >> 5u)) * 9u + seed * 17u) % + 181u) - 90; + kv_host[i] = (float)kv_value / 104.0f; + sc_host[i] = (float)sc_value / 88.0f; + } + memcpy(kv_host, &negative_zero, sizeof(negative_zero)); + TEST_ASSERT(ds4_gpu_tensor_write(kv, 0, kv_host, input_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write(sc, 0, sc_host, input_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_state_kv, 0, state_kv_host, state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_state_score, 0, state_score_host, + state_bytes) != 0); + TEST_ASSERT(setenv(pack_force_env, "1", 1) == 0); + TEST_ASSERT(unsetenv(pack_disable_env) == 0); + TEST_ASSERT(ds4_gpu_compressor_prefill_ratio4_replay_tensor( + fused_comp, fused_state_kv, fused_state_score, kv, sc, + model_raw, model_bytes, 0, 1, norm_offset, 0, + head_dim, 0, n_tokens, 0, 0, false, + 10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, 1.0e-6f) != 0); + + TEST_ASSERT(ds4_gpu_tensor_read( + kv, 0, source_after_host, input_bytes) != 0); + TEST_ASSERT(memcmp(kv_host, source_after_host, + (size_t)input_bytes) == 0); + TEST_ASSERT(ds4_gpu_tensor_read( + sc, 0, source_after_host, input_bytes) != 0); + TEST_ASSERT(memcmp(sc_host, source_after_host, + (size_t)input_bytes) == 0); + TEST_ASSERT(ds4_gpu_tensor_read( + ref_comp, 0, ref_comp_host, comp_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + fused_comp, 0, fused_comp_host, comp_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + ref_state_kv, 0, ref_state_kv_host, state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + fused_state_kv, 0, fused_state_kv_host, + state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + ref_state_score, 0, ref_state_score_host, + state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + fused_state_score, 0, fused_state_score_host, + state_bytes) != 0); + + comp_stats = test_compare_float_bits( + ref_comp_host, fused_comp_host, (size_t)comp_count); + state_kv_stats = test_compare_float_bits( + ref_state_kv_host, fused_state_kv_host, (size_t)state_count); + state_score_stats = test_compare_float_bits( + ref_state_score_host, fused_state_score_host, + (size_t)state_count); + } + + test_restore_env(ape_force_env, saved_ape_force); + test_restore_env(ape_disable_env, saved_ape_disable); + test_restore_env(pack_force_env, saved_pack_force); + test_restore_env(pack_disable_env, saved_pack_disable); + fprintf(stderr, + "ds4-test: compressor ratio4 replay pack exact head=%u " + "tokens=%u comp=%zu/%llu state_kv=%zu/%llu " + "state_score=%zu/%llu max_ulp=%u/%u/%u\n", + head_dim, n_tokens, + comp_stats.mismatch_count, (unsigned long long)comp_count, + state_kv_stats.mismatch_count, (unsigned long long)state_count, + state_score_stats.mismatch_count, (unsigned long long)state_count, + comp_stats.max_ulp, state_kv_stats.max_ulp, + state_score_stats.max_ulp); + TEST_ASSERT(comp_stats.mismatch_count == 0); + TEST_ASSERT(state_kv_stats.mismatch_count == 0); + TEST_ASSERT(state_score_stats.mismatch_count == 0); + + free(model_raw); + free(fused_state_score_host); + free(ref_state_score_host); + free(fused_state_kv_host); + free(ref_state_kv_host); + free(fused_comp_host); + free(ref_comp_host); + free(source_after_host); + free(state_score_host); + free(state_kv_host); + free(sc_host); + free(kv_host); + ds4_gpu_tensor_free(fused_state_score); + ds4_gpu_tensor_free(fused_state_kv); + ds4_gpu_tensor_free(ref_state_score); + ds4_gpu_tensor_free(ref_state_kv); + ds4_gpu_tensor_free(fused_comp); + ds4_gpu_tensor_free(ref_comp); + ds4_gpu_tensor_free(sc); + ds4_gpu_tensor_free(kv); +} + +static void test_metal_compressor_ratio4_replay_pack_exact(void) { + test_metal_compressor_ratio4_replay_pack_exact_case(128, 4, 47); + test_metal_compressor_ratio4_replay_pack_exact_case(512, 8, 53); +} + +static void test_metal_compressor_ratio4_direct_pool_exact_case( + uint32_t head_dim, + uint32_t pos0, + uint32_t n_tokens, + uint32_t ape_type, + bool replay, + uint32_t seed) { + const uint32_t ratio = 4u; + const uint32_t width = 2u * head_dim; + const uint32_t state_rows = 8u; + const uint32_t n_comp = n_tokens / ratio; + const uint64_t input_count = (uint64_t)n_tokens * width; + const uint64_t state_count = (uint64_t)state_rows * width; + const uint64_t comp_count = (uint64_t)n_comp * head_dim; + const uint64_t input_bytes = input_count * sizeof(float); + const uint64_t state_bytes = state_count * sizeof(float); + const uint64_t comp_bytes = comp_count * sizeof(float); + const uint64_t ape_elem_bytes = ape_type == 1u ? 2u : 4u; + const uint64_t ape_bytes = (uint64_t)ratio * width * ape_elem_bytes; + const uint64_t page = (uint64_t)getpagesize(); + const uint64_t norm_offset = test_round_up_u64(ape_bytes, page); + const uint64_t model_bytes = test_round_up_u64( + norm_offset + (uint64_t)head_dim * sizeof(float), page); + + TEST_ASSERT(n_comp != 0); + TEST_ASSERT(!replay || ((pos0 & 3u) == 0u && (n_tokens & 3u) == 0u)); + + ds4_gpu_tensor *kv = ds4_gpu_tensor_alloc(input_bytes); + ds4_gpu_tensor *sc = ds4_gpu_tensor_alloc(input_bytes); + ds4_gpu_tensor *ref_comp = ds4_gpu_tensor_alloc(comp_bytes); + ds4_gpu_tensor *direct_comp = ds4_gpu_tensor_alloc(comp_bytes); + ds4_gpu_tensor *ref_state_kv = ds4_gpu_tensor_alloc(state_bytes); + ds4_gpu_tensor *ref_state_score = ds4_gpu_tensor_alloc(state_bytes); + ds4_gpu_tensor *direct_state_kv = ds4_gpu_tensor_alloc(state_bytes); + ds4_gpu_tensor *direct_state_score = ds4_gpu_tensor_alloc(state_bytes); + float *kv_host = malloc((size_t)input_bytes); + float *sc_host = malloc((size_t)input_bytes); + float *source_after_host = malloc((size_t)input_bytes); + float *state_kv_host = malloc((size_t)state_bytes); + float *state_score_host = malloc((size_t)state_bytes); + float *ref_comp_host = malloc((size_t)comp_bytes); + float *direct_comp_host = malloc((size_t)comp_bytes); + float *ref_state_kv_host = malloc((size_t)state_bytes); + float *direct_state_kv_host = malloc((size_t)state_bytes); + float *ref_state_score_host = malloc((size_t)state_bytes); + float *direct_state_score_host = malloc((size_t)state_bytes); + void *model_raw = NULL; + TEST_ASSERT(posix_memalign( + &model_raw, (size_t)page, (size_t)model_bytes) == 0); + + TEST_ASSERT(kv != NULL); + TEST_ASSERT(sc != NULL); + TEST_ASSERT(ref_comp != NULL); + TEST_ASSERT(direct_comp != NULL); + TEST_ASSERT(ref_state_kv != NULL); + TEST_ASSERT(ref_state_score != NULL); + TEST_ASSERT(direct_state_kv != NULL); + TEST_ASSERT(direct_state_score != NULL); + TEST_ASSERT(kv_host != NULL); + TEST_ASSERT(sc_host != NULL); + TEST_ASSERT(source_after_host != NULL); + TEST_ASSERT(state_kv_host != NULL); + TEST_ASSERT(state_score_host != NULL); + TEST_ASSERT(ref_comp_host != NULL); + TEST_ASSERT(direct_comp_host != NULL); + TEST_ASSERT(ref_state_kv_host != NULL); + TEST_ASSERT(direct_state_kv_host != NULL); + TEST_ASSERT(ref_state_score_host != NULL); + TEST_ASSERT(direct_state_score_host != NULL); + TEST_ASSERT(model_raw != NULL); + + const char *ape_force_env = "DS4_METAL_ENABLE_COMPRESSOR_APE_ADD"; + const char *ape_disable_env = "DS4_METAL_DISABLE_M3_COMPRESSOR_APE_ADD"; + const char *pack_force_env = + "DS4_METAL_ENABLE_COMPRESSOR_RATIO4_PACK_FUSION"; + const char *pack_disable_env = + "DS4_METAL_DISABLE_M3_COMPRESSOR_RATIO4_PACK_FUSION"; + const char *direct_force_env = + "DS4_METAL_ENABLE_COMPRESSOR_RATIO4_DIRECT_POOL"; + const char *direct_disable_env = + "DS4_METAL_DISABLE_M3_COMPRESSOR_RATIO4_DIRECT_POOL"; + char *saved_ape_force = test_save_env(ape_force_env); + char *saved_ape_disable = test_save_env(ape_disable_env); + char *saved_pack_force = test_save_env(pack_force_env); + char *saved_pack_disable = test_save_env(pack_disable_env); + char *saved_direct_force = test_save_env(direct_force_env); + char *saved_direct_disable = test_save_env(direct_disable_env); + test_float_compare_stats comp_stats = {0}; + test_float_compare_stats state_kv_stats = {0}; + test_float_compare_stats state_score_stats = {0}; + + const bool allocated = kv && sc && ref_comp && direct_comp && + ref_state_kv && ref_state_score && direct_state_kv && + direct_state_score && kv_host && sc_host && source_after_host && + state_kv_host && state_score_host && ref_comp_host && + direct_comp_host && ref_state_kv_host && direct_state_kv_host && + ref_state_score_host && direct_state_score_host && model_raw; + if (allocated) { + memset(model_raw, 0, (size_t)model_bytes); + if (ape_type == 1u) { + uint16_t *ape = model_raw; + for (uint64_t i = 0; i < (uint64_t)ratio * width; i++) { + const int value = + (int)((i * 17u + (i ^ (i >> 4u)) * 5u + seed * 13u) % + 127u) - 63; + ape[i] = test_float_to_f16((float)value / 80.0f); + } + } else { + float *ape = model_raw; + for (uint64_t i = 0; i < (uint64_t)ratio * width; i++) { + const int value = + (int)((i * 17u + (i ^ (i >> 4u)) * 5u + seed * 13u) % + 127u) - 63; + ape[i] = (float)value / 80.0f; + } + } + float *norm = (float *)((uint8_t *)model_raw + norm_offset); + for (uint32_t i = 0; i < head_dim; i++) { + norm[i] = 0.75f + (float)((i * 7u + seed) % 19u) / 64.0f; + } + + for (uint64_t i = 0; i < input_count; i++) { + const uint32_t token = (uint32_t)(i / width); + const uint32_t col = (uint32_t)(i - (uint64_t)token * width); + const int kv_value = + (int)(((uint64_t)token * 37u + (uint64_t)col * 29u + + (col >= head_dim ? 71u : 3u) + seed * 23u) % + 193u) - 96; + const int sc_value = + (int)(((uint64_t)token * 41u + (uint64_t)col * 31u + + (col >= head_dim ? 17u : 83u) + seed * 11u) % + 181u) - 90; + kv_host[i] = (float)kv_value / 104.0f; + sc_host[i] = (float)sc_value / 32.0f; + } + const uint32_t negative_zero = 0x80000000u; + memcpy(kv_host, &negative_zero, sizeof(negative_zero)); + memcpy(sc_host + width + head_dim, &negative_zero, + sizeof(negative_zero)); + + for (uint32_t row = 0; row < state_rows; row++) { + for (uint32_t col = 0; col < width; col++) { + const uint64_t i = (uint64_t)row * width + col; + if (replay && row < ratio && col < head_dim) { + const int kv_value = + (int)(((uint64_t)row * 43u + (uint64_t)col * 19u + + seed * 29u) % 167u) - 83; + const int sc_value = + (int)(((uint64_t)row * 47u + (uint64_t)col * 23u + + seed * 7u) % 157u) - 78; + state_kv_host[i] = (float)kv_value / 80.0f; + state_score_host[i] = (float)sc_value / 28.0f; + } else { + const uint32_t kv_bits = + 0x7fc00001u + (uint32_t)(i & 0x3ffu); + const uint32_t score_bits = + 0x7fc00401u + (uint32_t)(i & 0x3ffu); + memcpy(state_kv_host + i, &kv_bits, sizeof(kv_bits)); + memcpy(state_score_host + i, &score_bits, + sizeof(score_bits)); + } + } + } + + for (uint64_t i = 0; i < comp_count; i++) { + const uint32_t bits = 0x7fc00801u + (uint32_t)(i & 0x3ffu); + memcpy(ref_comp_host + i, &bits, sizeof(bits)); + memcpy(direct_comp_host + i, &bits, sizeof(bits)); + } + + TEST_ASSERT(ds4_gpu_tensor_write(kv, 0, kv_host, input_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write(sc, 0, sc_host, input_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + ref_comp, 0, ref_comp_host, comp_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + direct_comp, 0, direct_comp_host, comp_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + ref_state_kv, 0, state_kv_host, state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + direct_state_kv, 0, state_kv_host, state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + ref_state_score, 0, state_score_host, + state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + direct_state_score, 0, state_score_host, + state_bytes) != 0); + TEST_ASSERT(ds4_gpu_set_model_map(model_raw, model_bytes) != 0); + ds4_gpu_set_quality(false); + + TEST_ASSERT(setenv(ape_force_env, "1", 1) == 0); + TEST_ASSERT(unsetenv(ape_disable_env) == 0); + TEST_ASSERT(unsetenv(direct_force_env) == 0); + TEST_ASSERT(setenv(direct_disable_env, "1", 1) == 0); + TEST_ASSERT(setenv(pack_force_env, "1", 1) == 0); + TEST_ASSERT(unsetenv(pack_disable_env) == 0); + int ref_ok; + if (replay) { + ref_ok = ds4_gpu_compressor_prefill_ratio4_replay_tensor( + ref_comp, ref_state_kv, ref_state_score, kv, sc, + model_raw, model_bytes, 0, ape_type, norm_offset, 0, + head_dim, pos0, n_tokens, 0, 0, false, + 10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, 1.0e-6f); + } else { + ref_ok = ds4_gpu_compressor_prefill_tensor( + ref_comp, ref_state_kv, ref_state_score, kv, sc, + model_raw, model_bytes, 0, ape_type, norm_offset, 0, + head_dim, ratio, pos0, n_tokens, 0, 0, false, + 10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, 1.0e-6f); + } + TEST_ASSERT(ref_ok != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + kv, 0, source_after_host, input_bytes) != 0); + TEST_ASSERT(memcmp(kv_host, source_after_host, + (size_t)input_bytes) == 0); + TEST_ASSERT(ds4_gpu_tensor_read( + sc, 0, source_after_host, input_bytes) != 0); + TEST_ASSERT(memcmp(sc_host, source_after_host, + (size_t)input_bytes) == 0); + + TEST_ASSERT(setenv(direct_force_env, "1", 1) == 0); + TEST_ASSERT(unsetenv(direct_disable_env) == 0); + TEST_ASSERT(unsetenv(pack_force_env) == 0); + TEST_ASSERT(setenv(pack_disable_env, "1", 1) == 0); + int direct_ok; + if (replay) { + direct_ok = ds4_gpu_compressor_prefill_ratio4_replay_tensor( + direct_comp, direct_state_kv, direct_state_score, kv, sc, + model_raw, model_bytes, 0, ape_type, norm_offset, 0, + head_dim, pos0, n_tokens, 0, 0, false, + 10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, 1.0e-6f); + } else { + direct_ok = ds4_gpu_compressor_prefill_tensor( + direct_comp, direct_state_kv, direct_state_score, kv, sc, + model_raw, model_bytes, 0, ape_type, norm_offset, 0, + head_dim, ratio, pos0, n_tokens, 0, 0, false, + 10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f, 1.0e-6f); + } + TEST_ASSERT(direct_ok != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + kv, 0, source_after_host, input_bytes) != 0); + TEST_ASSERT(memcmp(kv_host, source_after_host, + (size_t)input_bytes) == 0); + TEST_ASSERT(ds4_gpu_tensor_read( + sc, 0, source_after_host, input_bytes) != 0); + TEST_ASSERT(memcmp(sc_host, source_after_host, + (size_t)input_bytes) == 0); + + TEST_ASSERT(ds4_gpu_tensor_read( + ref_comp, 0, ref_comp_host, comp_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + direct_comp, 0, direct_comp_host, comp_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + ref_state_kv, 0, ref_state_kv_host, + state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + direct_state_kv, 0, direct_state_kv_host, + state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + ref_state_score, 0, ref_state_score_host, + state_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + direct_state_score, 0, direct_state_score_host, + state_bytes) != 0); + + comp_stats = test_compare_float_bits( + ref_comp_host, direct_comp_host, (size_t)comp_count); + state_kv_stats = test_compare_float_bits( + ref_state_kv_host, direct_state_kv_host, (size_t)state_count); + state_score_stats = test_compare_float_bits( + ref_state_score_host, direct_state_score_host, + (size_t)state_count); + } + + test_restore_env(ape_force_env, saved_ape_force); + test_restore_env(ape_disable_env, saved_ape_disable); + test_restore_env(pack_force_env, saved_pack_force); + test_restore_env(pack_disable_env, saved_pack_disable); + test_restore_env(direct_force_env, saved_direct_force); + test_restore_env(direct_disable_env, saved_direct_disable); + fprintf(stderr, + "ds4-test: compressor ratio4 direct pool exact mode=%s " + "head=%u pos=%u tokens=%u comp_rows=%u ape=%s " + "comp=%zu/%llu state_kv=%zu/%llu state_score=%zu/%llu " + "max_ulp=%u/%u/%u\n", + replay ? "replay" : "prefill", + head_dim, pos0, n_tokens, n_comp, + ape_type == 1u ? "f16" : "f32", + comp_stats.mismatch_count, (unsigned long long)comp_count, + state_kv_stats.mismatch_count, (unsigned long long)state_count, + state_score_stats.mismatch_count, + (unsigned long long)state_count, + comp_stats.max_ulp, state_kv_stats.max_ulp, + state_score_stats.max_ulp); + TEST_ASSERT(comp_stats.mismatch_count == 0); + TEST_ASSERT(state_kv_stats.mismatch_count == 0); + TEST_ASSERT(state_score_stats.mismatch_count == 0); + + free(model_raw); + free(direct_state_score_host); + free(ref_state_score_host); + free(direct_state_kv_host); + free(ref_state_kv_host); + free(direct_comp_host); + free(ref_comp_host); + free(state_score_host); + free(state_kv_host); + free(source_after_host); + free(sc_host); + free(kv_host); + ds4_gpu_tensor_free(direct_state_score); + ds4_gpu_tensor_free(direct_state_kv); + ds4_gpu_tensor_free(ref_state_score); + ds4_gpu_tensor_free(ref_state_kv); + ds4_gpu_tensor_free(direct_comp); + ds4_gpu_tensor_free(ref_comp); + ds4_gpu_tensor_free(sc); + ds4_gpu_tensor_free(kv); +} + +static void test_metal_compressor_ratio4_direct_pool_exact(void) { + /* n_comp == 1 deliberately stays on the exact GGML reduction path. */ + test_metal_compressor_ratio4_direct_pool_exact_case( + 512, 0, 4, 1, false, 59); + test_metal_compressor_ratio4_direct_pool_exact_case( + 512, 1, 16, 0, false, 61); + test_metal_compressor_ratio4_direct_pool_exact_case( + 128, 3, 14, 1, false, 67); + test_metal_compressor_ratio4_direct_pool_exact_case( + 512, 2048, 8, 1, true, 71); + test_metal_compressor_ratio4_direct_pool_exact_case( + 128, 12, 12, 0, true, 73); + test_metal_compressor_ratio4_direct_pool_exact_case( + 128, 8, 4, 1, true, 79); +} + +static void test_metal_inplace_rope_pair_exact(void) { + typedef struct { + uint32_t head_dim; + uint32_t n_rot; + uint32_t n_head; + uint32_t n_tok; + uint32_t pos0; + bool inverse; + float ext_factor; + } rope_case; + static const rope_case cases[] = { + { 512, 64, 64, 1, UINT32_MAX, false, 1.0f }, + { 512, 64, 1, 1, 2047, false, 1.0f }, + { 512, 64, 64, 1, 65533, true, 1.0f }, + { 128, 64, 64, 1, 37, false, 0.0f }, + { 512, 64, 4, 32, 0, false, 0.0f }, + { 512, 64, 7, 33, 2047, true, 1.0f }, + { 128, 64, 64, 34, 65533, false, 1.0f }, + { 128, 64, 64, 35, 37, true, 0.0f }, + { 128, 64, 4, 35, UINT32_MAX - 16u, true, 1.0f }, + }; + const char *disable_env = "DS4_METAL_DISABLE_M3_INPLACE_ROPE_PAIR"; + const char *enable_env = "DS4_METAL_ENABLE_INPLACE_ROPE_PAIR"; + const char *shared_disable_env = + "DS4_METAL_DISABLE_M3_SHARED_ROPE_COEFF"; + const char *affine_enable_env = + "DS4_METAL_ENABLE_AFFINE_ROPE_PAIR"; + const char *affine_disable_env = + "DS4_METAL_DISABLE_M3_AFFINE_ROPE_PAIR"; + char *saved_disable = test_save_env(disable_env); + char *saved_enable = test_save_env(enable_env); + char *saved_shared_disable = test_save_env(shared_disable_env); + char *saved_affine_enable = test_save_env(affine_enable_env); + char *saved_affine_disable = test_save_env(affine_disable_env); + size_t total_pair_mismatch = 0; + size_t total_shared_mismatch = 0; + size_t total_affine_mismatch = 0; + size_t total_pair_prefix_mismatch = 0; + size_t total_pair_tail_mismatch = 0; + size_t total_shared_prefix_mismatch = 0; + size_t total_shared_tail_mismatch = 0; + size_t total_affine_prefix_mismatch = 0; + size_t total_affine_tail_mismatch = 0; + size_t total_elements = 0; + + for (size_t ci = 0; ci < sizeof(cases) / sizeof(cases[0]); ci++) { + const rope_case *c = &cases[ci]; + const size_t elements = + (size_t)c->n_tok * c->n_head * c->head_dim; + const uint64_t bytes = (uint64_t)elements * sizeof(float); + const uint32_t n_nope = c->head_dim - c->n_rot; + const float freq_base = c->ext_factor != 0.0f ? 160000.0f : 10000.0f; + const float freq_scale = c->ext_factor != 0.0f ? 1.0f / 16.0f : 1.0f; + const uint32_t n_ctx_orig = c->ext_factor != 0.0f ? 65536u : 0u; + float attn_factor = 1.0f; + if (c->ext_factor != 0.0f) { + attn_factor /= 1.0f + 0.1f * logf(1.0f / freq_scale); + } + + ds4_gpu_tensor *reference = ds4_gpu_tensor_alloc(bytes); + ds4_gpu_tensor *pair_candidate = ds4_gpu_tensor_alloc(bytes); + ds4_gpu_tensor *shared_candidate = ds4_gpu_tensor_alloc(bytes); + ds4_gpu_tensor *affine_candidate = ds4_gpu_tensor_alloc(bytes); + float *input = malloc((size_t)bytes); + float *reference_host = malloc((size_t)bytes); + float *pair_host = malloc((size_t)bytes); + float *shared_host = malloc((size_t)bytes); + float *affine_host = malloc((size_t)bytes); + TEST_ASSERT(reference != NULL); + TEST_ASSERT(pair_candidate != NULL); + TEST_ASSERT(shared_candidate != NULL); + TEST_ASSERT(affine_candidate != NULL); + TEST_ASSERT(input != NULL); + TEST_ASSERT(reference_host != NULL); + TEST_ASSERT(pair_host != NULL); + TEST_ASSERT(shared_host != NULL); + TEST_ASSERT(affine_host != NULL); + + const bool allocated = reference && pair_candidate && + shared_candidate && affine_candidate && input && reference_host && + pair_host && shared_host && affine_host; + if (allocated) { + for (size_t i = 0; i < elements; i++) { + const uint32_t key = + (uint32_t)(i * 37u + (i ^ (i >> 5u)) * 11u + ci * 101u); + const int value = (int)(key % 4093u) - 2046; + input[i] = (float)value / 1024.0f; + if ((i + ci * 17u) % 257u == 0u) { + const uint32_t negative_zero = 0x80000000u; + memcpy(&input[i], &negative_zero, sizeof(negative_zero)); + } + } + + TEST_ASSERT(ds4_gpu_tensor_write(reference, 0, input, bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + pair_candidate, 0, input, bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + shared_candidate, 0, input, bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + affine_candidate, 0, input, bytes) != 0); + ds4_gpu_set_quality(false); + + TEST_ASSERT(unsetenv(affine_enable_env) == 0); + TEST_ASSERT(setenv(affine_disable_env, "1", 1) == 0); + TEST_ASSERT(unsetenv(enable_env) == 0); + TEST_ASSERT(setenv(disable_env, "1", 1) == 0); + TEST_ASSERT(setenv(shared_disable_env, "1", 1) == 0); + TEST_ASSERT(ds4_gpu_rope_tail_tensor( + reference, + c->n_tok, + c->n_head, + c->head_dim, + c->n_rot, + c->pos0, + n_ctx_orig, + c->inverse, + freq_base, + freq_scale, + c->ext_factor, + attn_factor, + 32.0f, + 1.0f) != 0); + + TEST_ASSERT(unsetenv(disable_env) == 0); + TEST_ASSERT(setenv(enable_env, "1", 1) == 0); + TEST_ASSERT(ds4_gpu_rope_tail_tensor( + pair_candidate, + c->n_tok, + c->n_head, + c->head_dim, + c->n_rot, + c->pos0, + n_ctx_orig, + c->inverse, + freq_base, + freq_scale, + c->ext_factor, + attn_factor, + 32.0f, + 1.0f) != 0); + + TEST_ASSERT(unsetenv(shared_disable_env) == 0); + TEST_ASSERT(ds4_gpu_rope_tail_tensor( + shared_candidate, + c->n_tok, + c->n_head, + c->head_dim, + c->n_rot, + c->pos0, + n_ctx_orig, + c->inverse, + freq_base, + freq_scale, + c->ext_factor, + attn_factor, + 32.0f, + 1.0f) != 0); + + TEST_ASSERT(unsetenv(enable_env) == 0); + TEST_ASSERT(unsetenv(shared_disable_env) == 0); + TEST_ASSERT(unsetenv(affine_disable_env) == 0); + TEST_ASSERT(setenv(affine_enable_env, "1", 1) == 0); + TEST_ASSERT(ds4_gpu_rope_tail_tensor( + affine_candidate, + c->n_tok, + c->n_head, + c->head_dim, + c->n_rot, + c->pos0, + n_ctx_orig, + c->inverse, + freq_base, + freq_scale, + c->ext_factor, + attn_factor, + 32.0f, + 1.0f) != 0); + + TEST_ASSERT(ds4_gpu_tensor_read( + reference, 0, reference_host, bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + pair_candidate, 0, pair_host, bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + shared_candidate, 0, shared_host, bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + affine_candidate, 0, affine_host, bytes) != 0); + + const test_float_compare_stats pair_stats = + test_compare_float_bits( + reference_host, pair_host, elements); + const test_float_compare_stats shared_stats = + test_compare_float_bits( + pair_host, shared_host, elements); + const test_float_compare_stats affine_stats = + test_compare_float_bits( + pair_host, affine_host, elements); + size_t pair_prefix_mismatch = 0; + size_t pair_tail_mismatch = 0; + size_t shared_prefix_mismatch = 0; + size_t shared_tail_mismatch = 0; + size_t affine_prefix_mismatch = 0; + size_t affine_tail_mismatch = 0; + for (uint32_t t = 0; t < c->n_tok; t++) { + for (uint32_t h = 0; h < c->n_head; h++) { + const size_t row = + ((size_t)t * c->n_head + h) * c->head_dim; + for (uint32_t d = 0; d < n_nope; d++) { + if (memcmp(&input[row + d], + &pair_host[row + d], + sizeof(float)) != 0) { + pair_prefix_mismatch++; + } + if (memcmp(&input[row + d], + &shared_host[row + d], + sizeof(float)) != 0) { + shared_prefix_mismatch++; + } + if (memcmp(&input[row + d], + &affine_host[row + d], + sizeof(float)) != 0) { + affine_prefix_mismatch++; + } + } + for (uint32_t d = n_nope; d < c->head_dim; d++) { + if (memcmp(&reference_host[row + d], + &pair_host[row + d], + sizeof(float)) != 0) { + pair_tail_mismatch++; + } + if (memcmp(&pair_host[row + d], + &shared_host[row + d], + sizeof(float)) != 0) { + shared_tail_mismatch++; + } + if (memcmp(&pair_host[row + d], + &affine_host[row + d], + sizeof(float)) != 0) { + affine_tail_mismatch++; + } + } + } + } + + fprintf(stderr, + "ds4-test: in-place RoPE exactness case=%zu " + "shape=%ux%ux%u pos=%u inverse=%d ext=%g " + "pair=%zu/%zu shared=%zu/%zu affine=%zu/%zu " + "pair_prefix=%zu pair_tail=%zu " + "shared_prefix=%zu shared_tail=%zu " + "affine_prefix=%zu affine_tail=%zu " + "pair_max_ulp=%u shared_max_ulp=%u affine_max_ulp=%u " + "pair_max_abs=%g shared_max_abs=%g affine_max_abs=%g\n", + ci, + c->n_tok, + c->n_head, + c->head_dim, + c->pos0, + c->inverse ? 1 : 0, + c->ext_factor, + pair_stats.mismatch_count, + elements, + shared_stats.mismatch_count, + elements, + affine_stats.mismatch_count, + elements, + pair_prefix_mismatch, + pair_tail_mismatch, + shared_prefix_mismatch, + shared_tail_mismatch, + affine_prefix_mismatch, + affine_tail_mismatch, + pair_stats.max_ulp, + shared_stats.max_ulp, + affine_stats.max_ulp, + pair_stats.max_abs, + shared_stats.max_abs, + affine_stats.max_abs); + TEST_ASSERT(pair_stats.mismatch_count == 0); + TEST_ASSERT(shared_stats.mismatch_count == 0); + TEST_ASSERT(affine_stats.mismatch_count == 0); + TEST_ASSERT(pair_prefix_mismatch == 0); + TEST_ASSERT(pair_tail_mismatch == 0); + TEST_ASSERT(shared_prefix_mismatch == 0); + TEST_ASSERT(shared_tail_mismatch == 0); + TEST_ASSERT(affine_prefix_mismatch == 0); + TEST_ASSERT(affine_tail_mismatch == 0); + total_pair_mismatch += pair_stats.mismatch_count; + total_shared_mismatch += shared_stats.mismatch_count; + total_affine_mismatch += affine_stats.mismatch_count; + total_pair_prefix_mismatch += pair_prefix_mismatch; + total_pair_tail_mismatch += pair_tail_mismatch; + total_shared_prefix_mismatch += shared_prefix_mismatch; + total_shared_tail_mismatch += shared_tail_mismatch; + total_affine_prefix_mismatch += affine_prefix_mismatch; + total_affine_tail_mismatch += affine_tail_mismatch; + total_elements += elements; + } + + free(affine_host); + free(shared_host); + free(pair_host); + free(reference_host); + free(input); + ds4_gpu_tensor_free(affine_candidate); + ds4_gpu_tensor_free(shared_candidate); + ds4_gpu_tensor_free(pair_candidate); + ds4_gpu_tensor_free(reference); + } + + test_restore_env(disable_env, saved_disable); + test_restore_env(enable_env, saved_enable); + test_restore_env(shared_disable_env, saved_shared_disable); + test_restore_env(affine_enable_env, saved_affine_enable); + test_restore_env(affine_disable_env, saved_affine_disable); + fprintf(stderr, + "ds4-test: in-place RoPE total pair=%zu/%zu shared=%zu/%zu " + "affine=%zu/%zu " + "pair_prefix=%zu pair_tail=%zu " + "shared_prefix=%zu shared_tail=%zu " + "affine_prefix=%zu affine_tail=%zu\n", + total_pair_mismatch, + total_elements, + total_shared_mismatch, + total_elements, + total_affine_mismatch, + total_elements, + total_pair_prefix_mismatch, + total_pair_tail_mismatch, + total_shared_prefix_mismatch, + total_shared_tail_mismatch, + total_affine_prefix_mismatch, + total_affine_tail_mismatch); + TEST_ASSERT(total_pair_mismatch == 0); + TEST_ASSERT(total_shared_mismatch == 0); + TEST_ASSERT(total_affine_mismatch == 0); + TEST_ASSERT(total_pair_prefix_mismatch == 0); + TEST_ASSERT(total_pair_tail_mismatch == 0); + TEST_ASSERT(total_shared_prefix_mismatch == 0); + TEST_ASSERT(total_shared_tail_mismatch == 0); + TEST_ASSERT(total_affine_prefix_mismatch == 0); + TEST_ASSERT(total_affine_tail_mismatch == 0); +} + +static void test_metal_contiguous_f32_f16_roundtrip_exact(void) { + typedef struct { + uint32_t n; + uint32_t src_offset; + uint32_t dst_offset; + } copy_case; + static const copy_case cases[] = { + { 1, 0, 0 }, + { 3, 4, 2 }, + { 4, 16, 8 }, + { 5, 12, 6 }, + { 17, 20, 10 }, + { 65, 4, 2 }, + }; + const char *env_name = "DS4_METAL_DISABLE_CONTIG_F32_F16_COPY"; + char *saved_env = test_save_env(env_name); + size_t half_mismatch = 0; + size_t half_guard_mismatch = 0; + size_t roundtrip_mismatch = 0; + size_t roundtrip_guard_mismatch = 0; + size_t half_total = 0; + size_t roundtrip_total = 0; + + for (size_t ci = 0; ci < sizeof(cases) / sizeof(cases[0]); ci++) { + const uint32_t n = cases[ci].n; + const uint64_t src_bytes = cases[ci].src_offset + + (uint64_t)n * sizeof(float) + 16u; + const uint64_t half_bytes = cases[ci].dst_offset + + (uint64_t)n * sizeof(uint16_t) + 16u; + const uint32_t raw_cap = 3; + const uint32_t raw_row = 1; + const uint64_t raw_bytes = + (uint64_t)raw_cap * n * sizeof(float); + + ds4_gpu_tensor *src_base = ds4_gpu_tensor_alloc(src_bytes); + ds4_gpu_tensor *src_view = src_base + ? ds4_gpu_tensor_view(src_base, + cases[ci].src_offset, + (uint64_t)n * sizeof(float)) + : NULL; + ds4_gpu_tensor *half_ref = ds4_gpu_tensor_alloc(half_bytes); + ds4_gpu_tensor *half_vec = ds4_gpu_tensor_alloc(half_bytes); + ds4_gpu_tensor *raw_ref = ds4_gpu_tensor_alloc(raw_bytes); + ds4_gpu_tensor *raw_vec = ds4_gpu_tensor_alloc(raw_bytes); + TEST_ASSERT(src_base != NULL); + TEST_ASSERT(src_view != NULL); + TEST_ASSERT(half_ref != NULL); + TEST_ASSERT(half_vec != NULL); + TEST_ASSERT(raw_ref != NULL); + TEST_ASSERT(raw_vec != NULL); + + uint8_t *src_host = malloc((size_t)src_bytes); + uint16_t *half_init = malloc((size_t)half_bytes); + uint16_t *half_ref_host = malloc((size_t)half_bytes); + uint16_t *half_vec_host = malloc((size_t)half_bytes); + uint32_t *raw_init = malloc((size_t)raw_bytes); + uint32_t *raw_ref_host = malloc((size_t)raw_bytes); + uint32_t *raw_vec_host = malloc((size_t)raw_bytes); + TEST_ASSERT(src_host != NULL); + TEST_ASSERT(half_init != NULL); + TEST_ASSERT(half_ref_host != NULL); + TEST_ASSERT(half_vec_host != NULL); + TEST_ASSERT(raw_init != NULL); + TEST_ASSERT(raw_ref_host != NULL); + TEST_ASSERT(raw_vec_host != NULL); + + const bool allocated = src_base && src_view && half_ref && half_vec && + raw_ref && raw_vec && src_host && half_init && half_ref_host && + half_vec_host && raw_init && raw_ref_host && raw_vec_host; + if (allocated) { + memset(src_host, 0x6d, (size_t)src_bytes); + test_fill_copy_f32_patterns(src_host + cases[ci].src_offset, + n, + (uint32_t)(ci * 7u)); + const size_t half_words = (size_t)(half_bytes / sizeof(uint16_t)); + for (size_t i = 0; i < half_words; i++) { + half_init[i] = (uint16_t)(0xa55au ^ (uint16_t)(i * 73u)); + } + const size_t raw_words = (size_t)raw_cap * n; + for (size_t i = 0; i < raw_words; i++) { + raw_init[i] = 0x4a000000u + (uint32_t)i; + } + + TEST_ASSERT(ds4_gpu_tensor_write(src_base, 0, src_host, src_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write(half_ref, 0, half_init, half_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write(half_vec, 0, half_init, half_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write(raw_ref, 0, raw_init, raw_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write(raw_vec, 0, raw_init, raw_bytes) != 0); + + TEST_ASSERT(setenv(env_name, "1", 1) == 0); + TEST_ASSERT(ds4_gpu_tensor_copy_f32_to_f16( + half_ref, + cases[ci].dst_offset, + src_base, + cases[ci].src_offset, + n) != 0); + TEST_ASSERT(ds4_gpu_store_raw_kv_tensor( + raw_ref, src_view, raw_cap, raw_row, n) != 0); + + TEST_ASSERT(setenv(env_name, "0", 1) == 0); + TEST_ASSERT(ds4_gpu_tensor_copy_f32_to_f16( + half_vec, + cases[ci].dst_offset, + src_base, + cases[ci].src_offset, + n) != 0); + TEST_ASSERT(ds4_gpu_store_raw_kv_tensor( + raw_vec, src_view, raw_cap, raw_row, n) != 0); + + TEST_ASSERT(ds4_gpu_tensor_read( + half_ref, 0, half_ref_host, half_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + half_vec, 0, half_vec_host, half_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + raw_ref, 0, raw_ref_host, raw_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + raw_vec, 0, raw_vec_host, raw_bytes) != 0); + + const size_t half_first = cases[ci].dst_offset / sizeof(uint16_t); + const size_t half_last = half_first + n; + for (size_t i = 0; i < half_words; i++) { + if (half_ref_host[i] == half_vec_host[i]) continue; + if (i >= half_first && i < half_last) { + half_mismatch++; + } else { + half_guard_mismatch++; + } + } + + const size_t raw_first = (size_t)raw_row * n; + const size_t raw_last = raw_first + n; + for (size_t i = 0; i < raw_words; i++) { + if (raw_ref_host[i] == raw_vec_host[i]) continue; + if (i >= raw_first && i < raw_last) { + roundtrip_mismatch++; + } else { + roundtrip_guard_mismatch++; + } + } + half_total += n; + roundtrip_total += n; + } + + free(raw_vec_host); + free(raw_ref_host); + free(raw_init); + free(half_vec_host); + free(half_ref_host); + free(half_init); + free(src_host); + ds4_gpu_tensor_free(raw_vec); + ds4_gpu_tensor_free(raw_ref); + ds4_gpu_tensor_free(half_vec); + ds4_gpu_tensor_free(half_ref); + ds4_gpu_tensor_free(src_view); + ds4_gpu_tensor_free(src_base); + } + + test_restore_env(env_name, saved_env); + fprintf(stderr, + "ds4-test: contiguous conversion exactness " + "f32_f16=%zu/%zu guard=%zu, f16_f32_roundtrip=%zu/%zu guard=%zu\n", + half_mismatch, + half_total, + half_guard_mismatch, + roundtrip_mismatch, + roundtrip_total, + roundtrip_guard_mismatch); + TEST_ASSERT(half_mismatch == 0); + TEST_ASSERT(half_guard_mismatch == 0); + TEST_ASSERT(roundtrip_mismatch == 0); + TEST_ASSERT(roundtrip_guard_mismatch == 0); +} +#endif + +#if defined(__APPLE__) +static void test_metal_gathered_kv_stage_exact(void) { + const uint32_t head_dim = 512; + const uint32_t raw_cap = 7; + const uint32_t n_raw = 5; + const uint32_t n_comp = 3; + const uint32_t raw_starts[] = {0, 2, 5, 6}; + const uint64_t raw_bytes = + (uint64_t)raw_cap * head_dim * sizeof(float); + const uint64_t comp_bytes = + (uint64_t)n_comp * head_dim * sizeof(uint16_t); + const uint64_t payload_bytes = + ((uint64_t)n_raw + n_comp) * head_dim * sizeof(uint16_t); + const uint64_t raw_out_bytes = + (uint64_t)n_raw * head_dim * sizeof(uint16_t); + const uint64_t raw_view_offset = 4; + const uint64_t comp_view_offset = 2; + const uint64_t dst_view_offset = 6; + const uint64_t raw_base_bytes = raw_view_offset + raw_bytes + 12; + const uint64_t comp_base_bytes = comp_view_offset + comp_bytes + 14; + const uint64_t dst_base_bytes = dst_view_offset + payload_bytes + 10; + const char *envs[] = { + "DS4_METAL_ENABLE_GATHERED_KV_STAGE", + "DS4_METAL_DISABLE_M3_GATHERED_KV_STAGE", + "DS4_METAL_REQUIRE_GATHERED_KV_STAGE", + "DS4_METAL_DISABLE_CONTIG_F32_F16_COPY", + "DS4_METAL_DISABLE_CONTIG_F16_F16_COPY", + }; + char *saved[sizeof(envs)/sizeof(envs[0])]; + for (size_t i = 0; i < sizeof(envs)/sizeof(envs[0]); i++) { + saved[i] = test_save_env(envs[i]); + } + + ds4_gpu_tensor *raw_base = ds4_gpu_tensor_alloc(raw_base_bytes); + ds4_gpu_tensor *raw = raw_base + ? ds4_gpu_tensor_view(raw_base, raw_view_offset, raw_bytes) + : NULL; + ds4_gpu_tensor *comp_base = ds4_gpu_tensor_alloc(comp_base_bytes); + ds4_gpu_tensor *comp = comp_base + ? ds4_gpu_tensor_view(comp_base, comp_view_offset, comp_bytes) + : NULL; + ds4_gpu_tensor *ref_base = ds4_gpu_tensor_alloc(dst_base_bytes); + ds4_gpu_tensor *ref = ref_base + ? ds4_gpu_tensor_view(ref_base, dst_view_offset, payload_bytes) + : NULL; + ds4_gpu_tensor *fused_base = ds4_gpu_tensor_alloc(dst_base_bytes); + ds4_gpu_tensor *fused = fused_base + ? ds4_gpu_tensor_view(fused_base, dst_view_offset, payload_bytes) + : NULL; + TEST_ASSERT(raw_base != NULL); + TEST_ASSERT(raw != NULL); + TEST_ASSERT(comp_base != NULL); + TEST_ASSERT(comp != NULL); + TEST_ASSERT(ref_base != NULL); + TEST_ASSERT(ref != NULL); + TEST_ASSERT(fused_base != NULL); + TEST_ASSERT(fused != NULL); + + float *raw_host = malloc((size_t)raw_bytes); + uint16_t *comp_host = malloc((size_t)comp_bytes); + uint8_t *dst_init = malloc((size_t)dst_base_bytes); + uint8_t *ref_host = malloc((size_t)dst_base_bytes); + uint8_t *fused_host = malloc((size_t)dst_base_bytes); + TEST_ASSERT(raw_host != NULL); + TEST_ASSERT(comp_host != NULL); + TEST_ASSERT(dst_init != NULL); + TEST_ASSERT(ref_host != NULL); + TEST_ASSERT(fused_host != NULL); + + static const uint16_t half_patterns[] = { + 0x0000u, 0x8000u, 0x0001u, 0x03ffu, 0x0400u, + 0x3555u, 0x3c00u, 0x3c01u, 0x7bffu, 0xfbffu, + 0x7c00u, 0xfc00u, 0x7e00u, 0x7e01u, 0xfe55u, + }; + const bool allocated = raw_base && raw && comp_base && comp && + ref_base && ref && fused_base && fused && raw_host && comp_host && + dst_init && ref_host && fused_host; + size_t raw_mismatch = 0; + size_t comp_mismatch = 0; + size_t guard_mismatch = 0; + if (allocated) { + for (uint32_t row = 0; row < raw_cap; row++) { + test_fill_copy_f32_patterns( + raw_host + (uint64_t)row * head_dim, + head_dim, + row * 17u + 3u); + } + for (uint64_t i = 0; i < (uint64_t)n_comp * head_dim; i++) { + comp_host[i] = half_patterns[(i * 7u + (i >> 3u)) % + (sizeof(half_patterns)/sizeof(half_patterns[0]))]; + } + for (uint64_t i = 0; i < dst_base_bytes; i++) { + dst_init[i] = (uint8_t)(0xa5u ^ (uint8_t)(i * 37u)); + } + TEST_ASSERT(ds4_gpu_tensor_write( + raw, 0, raw_host, raw_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + comp, 0, comp_host, comp_bytes) != 0); + ds4_gpu_set_quality(false); + + for (size_t ci = 0; + ci < sizeof(raw_starts)/sizeof(raw_starts[0]); + ci++) { + TEST_ASSERT(ds4_gpu_tensor_write( + ref_base, 0, dst_init, dst_base_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_base, 0, dst_init, dst_base_bytes) != 0); + + TEST_ASSERT(unsetenv(envs[0]) == 0); + TEST_ASSERT(setenv(envs[1], "1", 1) == 0); + TEST_ASSERT(unsetenv(envs[2]) == 0); + TEST_ASSERT(unsetenv(envs[3]) == 0); + TEST_ASSERT(unsetenv(envs[4]) == 0); + TEST_ASSERT(ds4_gpu_flash_kv_stage_f16_tensor( + ref, raw, raw_cap, raw_starts[ci], n_raw, + comp, 1, n_comp, head_dim) != 0); + + TEST_ASSERT(setenv(envs[0], "1", 1) == 0); + TEST_ASSERT(unsetenv(envs[1]) == 0); + TEST_ASSERT(setenv(envs[2], "1", 1) == 0); + TEST_ASSERT(ds4_gpu_flash_kv_stage_f16_tensor( + fused, raw, raw_cap, raw_starts[ci], n_raw, + comp, 1, n_comp, head_dim) != 0); + + TEST_ASSERT(ds4_gpu_tensor_read( + ref_base, 0, ref_host, dst_base_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + fused_base, 0, fused_host, + dst_base_bytes) != 0); + for (uint64_t i = 0; i < dst_base_bytes; i++) { + if (i < dst_view_offset || + i >= dst_view_offset + payload_bytes) { + if (ref_host[i] != dst_init[i]) guard_mismatch++; + if (fused_host[i] != dst_init[i]) guard_mismatch++; + } else if (ref_host[i] == fused_host[i]) { + continue; + } else if (i < dst_view_offset + raw_out_bytes) { + raw_mismatch++; + } else { + comp_mismatch++; + } + } + } + + /* Explicit disable, either component-copy disable, and quality mode + * all win over force under strict selection. */ + TEST_ASSERT(setenv(envs[0], "1", 1) == 0); + TEST_ASSERT(setenv(envs[1], "1", 1) == 0); + TEST_ASSERT(setenv(envs[2], "1", 1) == 0); + TEST_ASSERT(ds4_gpu_flash_kv_stage_f16_tensor( + fused, raw, raw_cap, 5, n_raw, + comp, 1, n_comp, head_dim) == 0); + TEST_ASSERT(unsetenv(envs[1]) == 0); + TEST_ASSERT(setenv(envs[3], "1", 1) == 0); + TEST_ASSERT(ds4_gpu_flash_kv_stage_f16_tensor( + fused, raw, raw_cap, 5, n_raw, + comp, 1, n_comp, head_dim) == 0); + TEST_ASSERT(unsetenv(envs[3]) == 0); + TEST_ASSERT(setenv(envs[4], "1", 1) == 0); + TEST_ASSERT(ds4_gpu_flash_kv_stage_f16_tensor( + fused, raw, raw_cap, 5, n_raw, + comp, 1, n_comp, head_dim) == 0); + TEST_ASSERT(unsetenv(envs[4]) == 0); + ds4_gpu_set_quality(true); + TEST_ASSERT(ds4_gpu_flash_kv_stage_f16_tensor( + fused, raw, raw_cap, 5, n_raw, + comp, 1, n_comp, head_dim) == 0); + ds4_gpu_set_quality(false); + } + + for (size_t i = 0; i < sizeof(envs)/sizeof(envs[0]); i++) { + test_restore_env(envs[i], saved[i]); + } + fprintf(stderr, + "ds4-test: gathered KV staging exact cases=%zu " + "raw_bytes=%zu comp_bytes=%zu guard_bytes=%zu\n", + sizeof(raw_starts)/sizeof(raw_starts[0]), + raw_mismatch, comp_mismatch, guard_mismatch); + TEST_ASSERT(raw_mismatch == 0); + TEST_ASSERT(comp_mismatch == 0); + TEST_ASSERT(guard_mismatch == 0); + + free(fused_host); + free(ref_host); + free(dst_init); + free(comp_host); + free(raw_host); + ds4_gpu_tensor_free(fused); + ds4_gpu_tensor_free(fused_base); + ds4_gpu_tensor_free(ref); + ds4_gpu_tensor_free(ref_base); + ds4_gpu_tensor_free(comp); + ds4_gpu_tensor_free(comp_base); + ds4_gpu_tensor_free(raw); + ds4_gpu_tensor_free(raw_base); +} + +static void test_metal_contiguous_compressed_f16_attention_exact(void) { + const uint32_t head_dim = 512; + const uint32_t n_head = 2; + const uint32_t raw_cap = 7; + const uint32_t n_raw = 5; + const uint32_t raw_start = 5; + const uint32_t n_comp = 3; + const uint64_t raw_bytes = + (uint64_t)raw_cap * head_dim * sizeof(float); + const uint64_t comp_bytes = + (uint64_t)n_comp * head_dim * sizeof(uint16_t); + const uint64_t q_bytes = + (uint64_t)n_head * head_dim * sizeof(float); + const uint64_t page = (uint64_t)getpagesize(); + const char *env_name = "DS4_METAL_DISABLE_CONTIG_F16_F16_COPY"; + static const uint16_t half_patterns[] = { + 0x0000u, 0x8000u, 0x0001u, 0x03ffu, 0x0400u, + 0x1001u, 0x3555u, 0x3c00u, 0x3c01u, 0x4000u, + 0xbc00u, 0xc000u, 0x7bffu, 0xfbffu, + }; + + ds4_gpu_tensor *raw = ds4_gpu_tensor_alloc(raw_bytes); + ds4_gpu_tensor *comp_base = ds4_gpu_tensor_alloc(comp_bytes + 18u); + ds4_gpu_tensor *comp = comp_base + ? ds4_gpu_tensor_view(comp_base, 2u, comp_bytes) + : NULL; + ds4_gpu_tensor *q = ds4_gpu_tensor_alloc(q_bytes); + ds4_gpu_tensor *heads_blit = ds4_gpu_tensor_alloc(q_bytes); + ds4_gpu_tensor *heads_compute = ds4_gpu_tensor_alloc(q_bytes); + TEST_ASSERT(raw != NULL); + TEST_ASSERT(comp_base != NULL); + TEST_ASSERT(comp != NULL); + TEST_ASSERT(q != NULL); + TEST_ASSERT(heads_blit != NULL); + TEST_ASSERT(heads_compute != NULL); + + float *raw_host = malloc((size_t)raw_bytes); + uint16_t *comp_host = malloc((size_t)comp_bytes); + float *q_host = malloc((size_t)q_bytes); + float *blit_host = malloc((size_t)q_bytes); + float *compute_host = malloc((size_t)q_bytes); + TEST_ASSERT(raw_host != NULL); + TEST_ASSERT(comp_host != NULL); + TEST_ASSERT(q_host != NULL); + TEST_ASSERT(blit_host != NULL); + TEST_ASSERT(compute_host != NULL); + + void *model_raw = NULL; + TEST_ASSERT(posix_memalign(&model_raw, (size_t)page, (size_t)page) == 0); + const bool allocated = raw && comp_base && comp && q && heads_blit && + heads_compute && raw_host && comp_host && q_host && blit_host && + compute_host && model_raw; + char *saved_env = test_save_env(env_name); + test_float_compare_stats stats = {0}; + if (allocated) { + memset(model_raw, 0, (size_t)page); + float *sinks = model_raw; + sinks[0] = -0.375f; + sinks[1] = 0.1875f; + for (uint64_t i = 0; i < (uint64_t)raw_cap * head_dim; i++) { + const int value = + (int)((i * 19u + (i ^ (i >> 5u)) * 7u) % 193u) - 96; + raw_host[i] = (float)value / 128.0f; + } + for (uint64_t i = 0; i < (uint64_t)n_comp * head_dim; i++) { + comp_host[i] = half_patterns[(i * 5u + (i >> 4u)) % + (sizeof(half_patterns) / sizeof(half_patterns[0]))]; + } + for (uint32_t i = 0; i < n_head * head_dim; i++) { + const int value = (int)((i * 37u + (i ^ (i >> 3u)) * 11u) % 251u) - 125; + q_host[i] = (float)value / 96.0f; + } + + TEST_ASSERT(ds4_gpu_tensor_write(raw, 0, raw_host, raw_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write(comp, 0, comp_host, comp_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write(q, 0, q_host, q_bytes) != 0); + TEST_ASSERT(ds4_gpu_set_model_map(model_raw, page) != 0); + ds4_gpu_set_quality(false); + + TEST_ASSERT(setenv(env_name, "1", 1) == 0); + TEST_ASSERT(ds4_gpu_attention_decode_heads_tensor( + heads_blit, + model_raw, + page, + 0, + q, + raw, + n_raw, + raw_cap, + raw_start, + comp, + 1, + n_comp, + NULL, + 0, + n_head, + head_dim) != 0); + + TEST_ASSERT(setenv(env_name, "0", 1) == 0); + TEST_ASSERT(ds4_gpu_attention_decode_heads_tensor( + heads_compute, + model_raw, + page, + 0, + q, + raw, + n_raw, + raw_cap, + raw_start, + comp, + 1, + n_comp, + NULL, + 0, + n_head, + head_dim) != 0); + + TEST_ASSERT(ds4_gpu_tensor_read( + heads_blit, 0, blit_host, q_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + heads_compute, 0, compute_host, q_bytes) != 0); + stats = test_compare_float_bits( + blit_host, compute_host, (size_t)n_head * head_dim); + } + test_restore_env(env_name, saved_env); + fprintf(stderr, + "ds4-test: contiguous compressed-F16 staging exactness " + "mismatches=%zu/%u max_ulp=%u max_abs=%g\n", + stats.mismatch_count, + n_head * head_dim, + stats.max_ulp, + stats.max_abs); + TEST_ASSERT(stats.mismatch_count == 0); + + free(model_raw); + free(compute_host); + free(blit_host); + free(q_host); + free(comp_host); + free(raw_host); + ds4_gpu_tensor_free(heads_compute); + ds4_gpu_tensor_free(heads_blit); + ds4_gpu_tensor_free(q); + ds4_gpu_tensor_free(comp); + ds4_gpu_tensor_free(comp_base); + ds4_gpu_tensor_free(raw); +} + +static void test_metal_persistent_zero_attention_mask_exact_case( + uint32_t raw_cap, + uint32_t n_raw, + uint32_t raw_start, + uint32_t n_comp, + uint32_t seed) { + const uint32_t head_dim = 512; + const uint32_t n_head = 2; + const uint64_t raw_bytes = + (uint64_t)raw_cap * head_dim * sizeof(float); + const uint64_t comp_bytes = + (uint64_t)n_comp * head_dim * sizeof(uint16_t); + const uint64_t q_bytes = + (uint64_t)n_head * head_dim * sizeof(float); + const uint64_t mask_bytes = (uint64_t)n_comp * sizeof(float); + const uint64_t page = (uint64_t)getpagesize(); + + ds4_gpu_tensor *raw = ds4_gpu_tensor_alloc(raw_bytes); + ds4_gpu_tensor *comp = ds4_gpu_tensor_alloc(comp_bytes); + ds4_gpu_tensor *q = ds4_gpu_tensor_alloc(q_bytes); + ds4_gpu_tensor *comp_mask = ds4_gpu_tensor_alloc(mask_bytes); + ds4_gpu_tensor *legacy = ds4_gpu_tensor_alloc(q_bytes); + ds4_gpu_tensor *persistent = ds4_gpu_tensor_alloc(q_bytes); + ds4_gpu_tensor *masked = ds4_gpu_tensor_alloc(q_bytes); + ds4_gpu_tensor *pad_legacy = ds4_gpu_tensor_alloc(q_bytes); + ds4_gpu_tensor *after_mask = ds4_gpu_tensor_alloc(q_bytes); + TEST_ASSERT(raw != NULL); + TEST_ASSERT(comp != NULL); + TEST_ASSERT(q != NULL); + TEST_ASSERT(comp_mask != NULL); + TEST_ASSERT(legacy != NULL); + TEST_ASSERT(persistent != NULL); + TEST_ASSERT(masked != NULL); + TEST_ASSERT(pad_legacy != NULL); + TEST_ASSERT(after_mask != NULL); + + float *raw_host = malloc((size_t)raw_bytes); + uint16_t *comp_host = malloc((size_t)comp_bytes); + float *q_host = malloc((size_t)q_bytes); + float *mask_host = malloc((size_t)mask_bytes); + float *legacy_host = malloc((size_t)q_bytes); + float *persistent_host = malloc((size_t)q_bytes); + float *masked_host = malloc((size_t)q_bytes); + float *pad_legacy_host = malloc((size_t)q_bytes); + float *after_mask_host = malloc((size_t)q_bytes); + TEST_ASSERT(raw_host != NULL); + TEST_ASSERT(comp_host != NULL); + TEST_ASSERT(q_host != NULL); + TEST_ASSERT(mask_host != NULL); + TEST_ASSERT(legacy_host != NULL); + TEST_ASSERT(persistent_host != NULL); + TEST_ASSERT(masked_host != NULL); + TEST_ASSERT(pad_legacy_host != NULL); + TEST_ASSERT(after_mask_host != NULL); + + void *model_raw = NULL; + TEST_ASSERT(posix_memalign(&model_raw, (size_t)page, (size_t)page) == 0); + const bool allocated = raw && comp && q && comp_mask && legacy && + persistent && masked && pad_legacy && after_mask && raw_host && + comp_host && q_host && mask_host && legacy_host && persistent_host && + masked_host && pad_legacy_host && after_mask_host && + model_raw; + const char *force_env = + "DS4_METAL_ENABLE_PERSISTENT_ZERO_ATTN_MASK"; + const char *disable_env = + "DS4_METAL_DISABLE_M3_PERSISTENT_ZERO_ATTN_MASK"; + const char *pad_disable_env = + "DS4_METAL_DISABLE_M3_GATHERED_KV_PAD_FUSION"; + const char *shared_pad_disable_env = + "DS4_METAL_DISABLE_M3_SHARED_KV_PAD"; + char *saved_force = test_save_env(force_env); + char *saved_disable = test_save_env(disable_env); + char *saved_pad_disable = test_save_env(pad_disable_env); + char *saved_shared_pad_disable = test_save_env(shared_pad_disable_env); + test_float_compare_stats persistent_stats = {0}; + test_float_compare_stats masked_stats = {0}; + test_float_compare_stats pad_stats = {0}; + test_float_compare_stats after_mask_stats = {0}; + + if (allocated) { + memset(model_raw, 0, (size_t)page); + float *sinks = model_raw; + sinks[0] = -0.3125f; + sinks[1] = 0.21875f; + for (uint64_t i = 0; i < (uint64_t)raw_cap * head_dim; i++) { + const int value = (int)((i * 17u + (i ^ (i >> 5u)) * 11u + + seed * 13u) % 211u) - 105; + raw_host[i] = (float)value / 128.0f; + } + for (uint64_t i = 0; i < (uint64_t)n_comp * head_dim; i++) { + const int value = (int)((i * 23u + (i ^ (i >> 4u)) * 7u + + seed * 19u) % 193u) - 96; + comp_host[i] = test_float_to_f16((float)value / 112.0f); + } + for (uint32_t i = 0; i < n_head * head_dim; i++) { + const int value = (int)((i * 31u + (i ^ (i >> 3u)) * 5u + + seed * 29u) % 227u) - 113; + q_host[i] = (float)value / 104.0f; + } + for (uint32_t i = 0; i < n_comp; i++) { + mask_host[i] = i == 0 ? -8.0f : -(float)(i + 1u) / 8.0f; + } + + TEST_ASSERT(ds4_gpu_tensor_write(raw, 0, raw_host, raw_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write(comp, 0, comp_host, comp_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write(q, 0, q_host, q_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + comp_mask, 0, mask_host, mask_bytes) != 0); + TEST_ASSERT(ds4_gpu_set_model_map(model_raw, page) != 0); + ds4_gpu_set_quality(false); + + unsetenv(force_env); + TEST_ASSERT(setenv(disable_env, "1", 1) == 0); + TEST_ASSERT(ds4_gpu_attention_decode_heads_tensor( + legacy, model_raw, page, 0, q, raw, + n_raw, raw_cap, raw_start, comp, 1, n_comp, + NULL, 0, n_head, head_dim) != 0); + + TEST_ASSERT(setenv(force_env, "1", 1) == 0); + unsetenv(disable_env); + unsetenv(pad_disable_env); + unsetenv(shared_pad_disable_env); + TEST_ASSERT(ds4_gpu_attention_decode_heads_tensor( + persistent, model_raw, page, 0, q, raw, + n_raw, raw_cap, raw_start, comp, 1, n_comp, + NULL, 0, n_head, head_dim) != 0); + + TEST_ASSERT(ds4_gpu_attention_decode_heads_tensor( + masked, model_raw, page, 0, q, raw, + n_raw, raw_cap, raw_start, comp, 1, n_comp, + comp_mask, 1, n_head, head_dim) != 0); + + TEST_ASSERT(setenv(pad_disable_env, "1", 1) == 0); + TEST_ASSERT(setenv(shared_pad_disable_env, "1", 1) == 0); + TEST_ASSERT(ds4_gpu_attention_decode_heads_tensor( + pad_legacy, model_raw, page, 0, q, raw, + n_raw, raw_cap, raw_start, comp, 1, n_comp, + comp_mask, 1, n_head, head_dim) != 0); + unsetenv(pad_disable_env); + unsetenv(shared_pad_disable_env); + + TEST_ASSERT(ds4_gpu_attention_decode_heads_tensor( + after_mask, model_raw, page, 0, q, raw, + n_raw, raw_cap, raw_start, comp, 1, n_comp, + NULL, 0, n_head, head_dim) != 0); + + TEST_ASSERT(ds4_gpu_tensor_read( + legacy, 0, legacy_host, q_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + persistent, 0, persistent_host, q_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + masked, 0, masked_host, q_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + pad_legacy, 0, pad_legacy_host, q_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + after_mask, 0, after_mask_host, q_bytes) != 0); + persistent_stats = test_compare_float_bits( + legacy_host, persistent_host, (size_t)n_head * head_dim); + masked_stats = test_compare_float_bits( + legacy_host, masked_host, (size_t)n_head * head_dim); + pad_stats = test_compare_float_bits( + pad_legacy_host, masked_host, (size_t)n_head * head_dim); + after_mask_stats = test_compare_float_bits( + legacy_host, after_mask_host, (size_t)n_head * head_dim); + } + + test_restore_env(force_env, saved_force); + test_restore_env(disable_env, saved_disable); + test_restore_env(pad_disable_env, saved_pad_disable); + test_restore_env(shared_pad_disable_env, saved_shared_pad_disable); + fprintf(stderr, + "ds4-test: persistent zero attention mask exact keys=%u " + "candidate=%zu/%u max_ulp=%u masked_diff=%zu/%u " + "pad_fusion=%zu/%u max_ulp=%u " + "after_mask=%zu/%u max_ulp=%u\n", + n_raw + n_comp, + persistent_stats.mismatch_count, n_head * head_dim, + persistent_stats.max_ulp, + masked_stats.mismatch_count, n_head * head_dim, + pad_stats.mismatch_count, n_head * head_dim, + pad_stats.max_ulp, + after_mask_stats.mismatch_count, n_head * head_dim, + after_mask_stats.max_ulp); + TEST_ASSERT(persistent_stats.mismatch_count == 0); + TEST_ASSERT(masked_stats.mismatch_count != 0); + TEST_ASSERT(pad_stats.mismatch_count == 0); + TEST_ASSERT(after_mask_stats.mismatch_count == 0); + + free(model_raw); + free(after_mask_host); + free(pad_legacy_host); + free(masked_host); + free(persistent_host); + free(legacy_host); + free(mask_host); + free(q_host); + free(comp_host); + free(raw_host); + ds4_gpu_tensor_free(after_mask); + ds4_gpu_tensor_free(pad_legacy); + ds4_gpu_tensor_free(masked); + ds4_gpu_tensor_free(persistent); + ds4_gpu_tensor_free(legacy); + ds4_gpu_tensor_free(comp_mask); + ds4_gpu_tensor_free(q); + ds4_gpu_tensor_free(comp); + ds4_gpu_tensor_free(raw); +} + +static void test_metal_persistent_zero_attention_mask_exact(void) { + test_metal_persistent_zero_attention_mask_exact_case(7, 5, 5, 3, 11); + test_metal_persistent_zero_attention_mask_exact_case(37, 29, 35, 3, 23); +} + +typedef enum { + TEST_METAL_PREFILL_MASK_CACHE_RAW = 1, + TEST_METAL_PREFILL_MASK_CACHE_RATIO4 = 2, + TEST_METAL_PREFILL_MASK_CACHE_RATIO128 = 3, +} test_metal_prefill_mask_cache_kind; + +typedef struct { + uint32_t n_tokens; + uint32_t n_comp; + uint32_t window; + uint32_t ratio; +} test_metal_prefill_mask_cache_shape; + +static int test_metal_zero_prefix_prefill_mask_cache_call( + test_metal_prefill_mask_cache_kind kind, + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw, + const ds4_gpu_tensor *comp, + const ds4_gpu_tensor *comp_mask, + const test_metal_prefill_mask_cache_shape *shape, + bool masked, + uint32_t n_head, + uint32_t head_dim) { + if (kind == TEST_METAL_PREFILL_MASK_CACHE_RAW) { + if (masked) return 0; + return ds4_gpu_attention_prefill_raw_heads_tensor( + heads, model_map, model_size, 0, q, raw, + shape->n_tokens, shape->window, n_head, head_dim); + } + + if (masked) { + return ds4_gpu_attention_prefill_masked_mixed_heads_tensor( + heads, model_map, model_size, 0, q, raw, comp, 1, comp_mask, + shape->n_tokens, shape->n_comp, shape->window, shape->ratio, + n_head, head_dim); + } + + return ds4_gpu_attention_prefill_static_mixed_heads_tensor( + heads, model_map, model_size, 0, q, raw, comp, 1, + shape->n_tokens, shape->n_comp, shape->window, shape->ratio, + n_head, head_dim); +} + +static bool test_metal_zero_prefix_prefill_mask_cache_run( + test_metal_prefill_mask_cache_kind kind, + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw, + const ds4_gpu_tensor *comp, + const ds4_gpu_tensor *comp_mask, + const test_metal_prefill_mask_cache_shape *shape, + bool masked, + uint32_t n_head, + uint32_t head_dim, + float *host) { + const uint64_t bytes = + (uint64_t)shape->n_tokens * n_head * head_dim * sizeof(float); + const int call_ok = test_metal_zero_prefix_prefill_mask_cache_call( + kind, heads, model_map, model_size, q, raw, comp, comp_mask, + shape, masked, n_head, head_dim); + TEST_ASSERT(call_ok != 0); + if (!call_ok) return false; + + const int read_ok = ds4_gpu_tensor_read(heads, 0, host, bytes); + TEST_ASSERT(read_ok != 0); + return read_ok != 0; +} + +static void test_metal_zero_prefix_prefill_mask_cache_compare( + const float *expected, + const float *actual, + size_t count, + size_t *total_mismatches, + uint32_t *max_ulp) { + const test_float_compare_stats stats = + test_compare_float_bits(expected, actual, count); + *total_mismatches += stats.mismatch_count; + if (stats.max_ulp > *max_ulp) *max_ulp = stats.max_ulp; +} + +static void test_metal_zero_prefix_prefill_mask_cache_exact_kind( + test_metal_prefill_mask_cache_kind kind, + uint32_t seed) { + const uint32_t head_dim = 512; + const uint32_t n_head = 1; + const uint32_t max_tokens = 129; + const uint32_t max_comp = 32; + const uint64_t raw_count = (uint64_t)max_tokens * head_dim; + const uint64_t comp_count = (uint64_t)max_comp * head_dim; + const uint64_t q_count = (uint64_t)max_tokens * n_head * head_dim; + const uint64_t mask_count = (uint64_t)max_tokens * max_comp; + const uint64_t raw_bytes = raw_count * sizeof(float); + const uint64_t comp_bytes = comp_count * sizeof(uint16_t); + const uint64_t q_bytes = q_count * sizeof(float); + const uint64_t mask_bytes = mask_count * sizeof(float); + const uint64_t page = (uint64_t)getpagesize(); + + const test_metal_prefill_mask_cache_shape shape_a = { + .n_tokens = 128, + .n_comp = kind == TEST_METAL_PREFILL_MASK_CACHE_RAW ? 0u : + (kind == TEST_METAL_PREFILL_MASK_CACHE_RATIO4 ? 32u : 1u), + .window = 128, + .ratio = kind == TEST_METAL_PREFILL_MASK_CACHE_RATIO4 ? 4u : + (kind == TEST_METAL_PREFILL_MASK_CACHE_RATIO128 ? 128u : 0u), + }; + const test_metal_prefill_mask_cache_shape shape_b = { + .n_tokens = 129, + .n_comp = kind == TEST_METAL_PREFILL_MASK_CACHE_RAW ? 0u : + (kind == TEST_METAL_PREFILL_MASK_CACHE_RATIO4 ? 31u : 2u), + .window = 63, + .ratio = shape_a.ratio, + }; + const size_t count_a = + (size_t)shape_a.n_tokens * n_head * head_dim; + const size_t count_b = + (size_t)shape_b.n_tokens * n_head * head_dim; + + ds4_gpu_tensor *raw = ds4_gpu_tensor_alloc(raw_bytes); + ds4_gpu_tensor *comp = ds4_gpu_tensor_alloc(comp_bytes); + ds4_gpu_tensor *q = ds4_gpu_tensor_alloc(q_bytes); + ds4_gpu_tensor *comp_mask = ds4_gpu_tensor_alloc(mask_bytes); + ds4_gpu_tensor *heads = ds4_gpu_tensor_alloc(q_bytes); + float *raw_host = malloc((size_t)raw_bytes); + uint16_t *comp_host = malloc((size_t)comp_bytes); + float *q_host = calloc((size_t)q_count, sizeof(float)); + float *mask_host = malloc((size_t)mask_bytes); + float *ref_a = malloc((size_t)q_bytes); + float *ref_b = malloc((size_t)q_bytes); + float *actual = malloc((size_t)q_bytes); + float *masked_actual = malloc((size_t)q_bytes); + void *model_raw = NULL; + TEST_ASSERT(posix_memalign(&model_raw, (size_t)page, (size_t)page) == 0); + + TEST_ASSERT(raw != NULL); + TEST_ASSERT(comp != NULL); + TEST_ASSERT(q != NULL); + TEST_ASSERT(comp_mask != NULL); + TEST_ASSERT(heads != NULL); + TEST_ASSERT(raw_host != NULL); + TEST_ASSERT(comp_host != NULL); + TEST_ASSERT(q_host != NULL); + TEST_ASSERT(mask_host != NULL); + TEST_ASSERT(ref_a != NULL); + TEST_ASSERT(ref_b != NULL); + TEST_ASSERT(actual != NULL); + TEST_ASSERT(masked_actual != NULL); + TEST_ASSERT(model_raw != NULL); + + const char *force_env = + "DS4_METAL_ENABLE_ZERO_PREFIX_PREFILL_MASK_CACHE"; + const char *disable_env = + "DS4_METAL_DISABLE_M3_ZERO_PREFIX_PREFILL_MASK_CACHE"; + char *saved_force = test_save_env(force_env); + char *saved_disable = test_save_env(disable_env); + size_t total_mismatches = 0; + uint32_t max_ulp = 0; + size_t key_difference = 0; + size_t masked_difference = 0; + + const bool allocated = raw && comp && q && comp_mask && heads && + raw_host && comp_host && q_host && mask_host && ref_a && ref_b && + actual && masked_actual && model_raw; + if (allocated) { + memset(model_raw, 0, (size_t)page); + ((float *)model_raw)[0] = -1.0f; + + for (uint32_t row = 0; row < max_tokens; row++) { + for (uint32_t col = 0; col < head_dim; col++) { + const int value = (int)((row * 37u + col * 17u + + (col ^ (col >> 3u)) * 5u + + seed * 13u) % 257u) - 128; + raw_host[(uint64_t)row * head_dim + col] = + (float)value / 256.0f; + } + } + for (uint32_t row = 0; row < max_comp; row++) { + for (uint32_t col = 0; col < head_dim; col++) { + const int value = (int)((row * 29u + col * 11u + + (col ^ (col >> 4u)) * 7u + + seed * 19u) % 193u) - 96; + comp_host[(uint64_t)row * head_dim + col] = + test_float_to_f16(0.375f + (float)value / 384.0f); + } + } + for (uint64_t i = 0; i < mask_count; i++) { + mask_host[i] = -65504.0f; + } + if (shape_a.n_comp != 0) { + for (uint32_t row = 0; row < shape_a.n_tokens; row++) { + const uint32_t visible = (row + 1u) / shape_a.ratio; + for (uint32_t col = 0; col < shape_a.n_comp; col++) { + if (col < visible) { + mask_host[(uint64_t)row * shape_a.n_comp + col] = + (col & 1u) != 0u ? -2.0f : -65504.0f; + } + } + } + } + + TEST_ASSERT(ds4_gpu_tensor_write(raw, 0, raw_host, raw_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write(comp, 0, comp_host, comp_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write(q, 0, q_host, q_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + comp_mask, 0, mask_host, mask_bytes) != 0); + TEST_ASSERT(ds4_gpu_set_model_map(model_raw, page) != 0); + ds4_gpu_set_quality(false); + + TEST_ASSERT(unsetenv(force_env) == 0); + TEST_ASSERT(setenv(disable_env, "1", 1) == 0); + const bool have_ref_a = + test_metal_zero_prefix_prefill_mask_cache_run( + kind, heads, model_raw, page, q, raw, comp, comp_mask, + &shape_a, false, n_head, head_dim, ref_a); + const bool have_ref_b = + test_metal_zero_prefix_prefill_mask_cache_run( + kind, heads, model_raw, page, q, raw, comp, comp_mask, + &shape_b, false, n_head, head_dim, ref_b); + + if (have_ref_a && have_ref_b) { + const test_float_compare_stats key_stats = + test_compare_float_bits(ref_a, ref_b, count_a); + key_difference = key_stats.mismatch_count; + TEST_ASSERT(key_difference != 0); + } + + TEST_ASSERT(setenv(force_env, "1", 1) == 0); + TEST_ASSERT(unsetenv(disable_env) == 0); + + if (test_metal_zero_prefix_prefill_mask_cache_run( + kind, heads, model_raw, page, q, raw, comp, comp_mask, + &shape_a, false, n_head, head_dim, actual) && have_ref_a) { + test_metal_zero_prefix_prefill_mask_cache_compare( + ref_a, actual, count_a, &total_mismatches, &max_ulp); + } + if (test_metal_zero_prefix_prefill_mask_cache_run( + kind, heads, model_raw, page, q, raw, comp, comp_mask, + &shape_a, false, n_head, head_dim, actual) && have_ref_a) { + test_metal_zero_prefix_prefill_mask_cache_compare( + ref_a, actual, count_a, &total_mismatches, &max_ulp); + } + + if (kind != TEST_METAL_PREFILL_MASK_CACHE_RAW) { + if (test_metal_zero_prefix_prefill_mask_cache_run( + kind, heads, model_raw, page, q, raw, comp, comp_mask, + &shape_a, true, n_head, head_dim, masked_actual) && + have_ref_a) { + const test_float_compare_stats masked_stats = + test_compare_float_bits(ref_a, masked_actual, count_a); + masked_difference = masked_stats.mismatch_count; + TEST_ASSERT(masked_difference != 0); + } + if (test_metal_zero_prefix_prefill_mask_cache_run( + kind, heads, model_raw, page, q, raw, comp, comp_mask, + &shape_a, false, n_head, head_dim, actual) && have_ref_a) { + test_metal_zero_prefix_prefill_mask_cache_compare( + ref_a, actual, count_a, + &total_mismatches, &max_ulp); + } + } + + if (test_metal_zero_prefix_prefill_mask_cache_run( + kind, heads, model_raw, page, q, raw, comp, comp_mask, + &shape_b, false, n_head, head_dim, actual) && have_ref_b) { + test_metal_zero_prefix_prefill_mask_cache_compare( + ref_b, actual, count_b, &total_mismatches, &max_ulp); + } + if (test_metal_zero_prefix_prefill_mask_cache_run( + kind, heads, model_raw, page, q, raw, comp, comp_mask, + &shape_a, false, n_head, head_dim, actual) && have_ref_a) { + test_metal_zero_prefix_prefill_mask_cache_compare( + ref_a, actual, count_a, &total_mismatches, &max_ulp); + } + } + + test_restore_env(force_env, saved_force); + test_restore_env(disable_env, saved_disable); + const char *kind_name = kind == TEST_METAL_PREFILL_MASK_CACHE_RAW ? "raw" : + (kind == TEST_METAL_PREFILL_MASK_CACHE_RATIO4 ? "ratio4" : "ratio128"); + fprintf(stderr, + "ds4-test: zero-prefix prefill mask cache %s exact " + "mismatches=%zu max_ulp=%u key_diff=%zu masked_diff=%zu\n", + kind_name, total_mismatches, max_ulp, + key_difference, masked_difference); + TEST_ASSERT(total_mismatches == 0); + TEST_ASSERT(max_ulp == 0); + + free(model_raw); + free(masked_actual); + free(actual); + free(ref_b); + free(ref_a); + free(mask_host); + free(q_host); + free(comp_host); + free(raw_host); + ds4_gpu_tensor_free(heads); + ds4_gpu_tensor_free(comp_mask); + ds4_gpu_tensor_free(q); + ds4_gpu_tensor_free(comp); + ds4_gpu_tensor_free(raw); +} + +static void test_metal_zero_prefix_prefill_mask_cache_exact(void) { + test_metal_zero_prefix_prefill_mask_cache_exact_kind( + TEST_METAL_PREFILL_MASK_CACHE_RAW, 41); + test_metal_zero_prefix_prefill_mask_cache_exact_kind( + TEST_METAL_PREFILL_MASK_CACHE_RATIO4, 43); + test_metal_zero_prefix_prefill_mask_cache_exact_kind( + TEST_METAL_PREFILL_MASK_CACHE_RATIO128, 47); +} +#endif + +#if defined(__APPLE__) +static void test_metal_hc_split_weighted_sum_norm_batch_exact(void) { + /* Compare the batched HC+RMSNorm fusion against the exact two-dispatch + * sequence used by the reference path at DS4's production dimensions. */ + const uint32_t n_embd = 7168; + const uint32_t n_hc = 4; + const uint32_t n_rows = 3; + const uint32_t sinkhorn_iters = 20; + const float hc_eps = 1.0e-6f; + const float norm_eps = 1.0e-6f; + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + const uint64_t page = (uint64_t)getpagesize(); + const uint64_t scale_offset = 0; + const uint64_t base_offset = test_round_up_u64(3u * sizeof(float), page); + const uint64_t norm_weight_offset = + test_round_up_u64(base_offset + mix_hc * sizeof(float), page); + const uint64_t model_alloc = test_round_up_u64( + norm_weight_offset + (uint64_t)n_embd * sizeof(float), page); + const uint64_t mix_bytes = (uint64_t)n_rows * mix_hc * sizeof(float); + const uint64_t residual_bytes = + (uint64_t)n_rows * n_hc * n_embd * sizeof(float); + const uint64_t out_bytes = (uint64_t)n_rows * n_embd * sizeof(float); + + void *model_raw = NULL; + TEST_ASSERT(posix_memalign(&model_raw, (size_t)page, (size_t)model_alloc) == 0); + if (!model_raw) return; + memset(model_raw, 0, (size_t)model_alloc); + + float *scale = (float *)((uint8_t *)model_raw + scale_offset); + float *base = (float *)((uint8_t *)model_raw + base_offset); + float *norm_weight = (float *)((uint8_t *)model_raw + norm_weight_offset); + scale[0] = 0.625f; + scale[1] = -0.75f; + scale[2] = 0.4375f; + for (uint32_t i = 0; i < mix_hc; i++) { + const int value = (int)((i * 17u + 5u) % 29u) - 14; + base[i] = (float)value / 16.0f; + } + for (uint32_t i = 0; i < n_embd; i++) { + norm_weight[i] = 0.5f + (float)((i * 13u + 7u) % 31u) / 32.0f; + } + + ds4_gpu_tensor *mix = ds4_gpu_tensor_alloc(mix_bytes); + ds4_gpu_tensor *residual = ds4_gpu_tensor_alloc(residual_bytes); + ds4_gpu_tensor *ref_split = ds4_gpu_tensor_alloc(mix_bytes); + ds4_gpu_tensor *fused_split = ds4_gpu_tensor_alloc(mix_bytes); + ds4_gpu_tensor *ref_out = ds4_gpu_tensor_alloc(out_bytes); + ds4_gpu_tensor *fused_out = ds4_gpu_tensor_alloc(out_bytes); + ds4_gpu_tensor *ref_norm = ds4_gpu_tensor_alloc(out_bytes); + ds4_gpu_tensor *fused_norm = ds4_gpu_tensor_alloc(out_bytes); + TEST_ASSERT(mix != NULL); + TEST_ASSERT(residual != NULL); + TEST_ASSERT(ref_split != NULL); + TEST_ASSERT(fused_split != NULL); + TEST_ASSERT(ref_out != NULL); + TEST_ASSERT(fused_out != NULL); + TEST_ASSERT(ref_norm != NULL); + TEST_ASSERT(fused_norm != NULL); + if (!mix || !residual || !ref_split || !fused_split || + !ref_out || !fused_out || !ref_norm || !fused_norm) { + ds4_gpu_tensor_free(mix); + ds4_gpu_tensor_free(residual); + ds4_gpu_tensor_free(ref_split); + ds4_gpu_tensor_free(fused_split); + ds4_gpu_tensor_free(ref_out); + ds4_gpu_tensor_free(fused_out); + ds4_gpu_tensor_free(ref_norm); + ds4_gpu_tensor_free(fused_norm); + free(model_raw); + return; + } + + float *mix_host = malloc((size_t)mix_bytes); + float *residual_host = malloc((size_t)residual_bytes); + float *ref_split_host = malloc((size_t)mix_bytes); + float *fused_split_host = malloc((size_t)mix_bytes); + float *ref_out_host = malloc((size_t)out_bytes); + float *fused_out_host = malloc((size_t)out_bytes); + float *ref_norm_host = malloc((size_t)out_bytes); + float *fused_norm_host = malloc((size_t)out_bytes); + TEST_ASSERT(mix_host != NULL); + TEST_ASSERT(residual_host != NULL); + TEST_ASSERT(ref_split_host != NULL); + TEST_ASSERT(fused_split_host != NULL); + TEST_ASSERT(ref_out_host != NULL); + TEST_ASSERT(fused_out_host != NULL); + TEST_ASSERT(ref_norm_host != NULL); + TEST_ASSERT(fused_norm_host != NULL); + if (!mix_host || !residual_host || !ref_split_host || !fused_split_host || + !ref_out_host || !fused_out_host || !ref_norm_host || !fused_norm_host) { + free(mix_host); + free(residual_host); + free(ref_split_host); + free(fused_split_host); + free(ref_out_host); + free(fused_out_host); + free(ref_norm_host); + free(fused_norm_host); + ds4_gpu_tensor_free(mix); + ds4_gpu_tensor_free(residual); + ds4_gpu_tensor_free(ref_split); + ds4_gpu_tensor_free(fused_split); + ds4_gpu_tensor_free(ref_out); + ds4_gpu_tensor_free(fused_out); + ds4_gpu_tensor_free(ref_norm); + ds4_gpu_tensor_free(fused_norm); + free(model_raw); + return; + } + + for (uint32_t row = 0; row < n_rows; row++) { + for (uint32_t i = 0; i < mix_hc; i++) { + const int value = + (int)(((row + 1u) * 19u + i * 11u + (i ^ row) * 3u) % 47u) - 23; + mix_host[(uint64_t)row * mix_hc + i] = (float)value / 9.0f; + } + for (uint32_t hc = 0; hc < n_hc; hc++) { + for (uint32_t d = 0; d < n_embd; d++) { + const uint32_t key = + d * 37u + hc * 173u + row * 997u + ((d >> 3u) ^ (d * 7u)); + const int value = (int)(key % 2047u) - 1023; + residual_host[((uint64_t)row * n_hc + hc) * n_embd + d] = + (float)value / 512.0f; + } + } + } + + TEST_ASSERT(ds4_gpu_tensor_write(mix, 0, mix_host, mix_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write(residual, 0, residual_host, residual_bytes) != 0); + TEST_ASSERT(ds4_gpu_set_model_map(model_raw, model_alloc) != 0); + TEST_ASSERT(ds4_gpu_hc_split_weighted_sum_tensor( + ref_out, ref_split, mix, residual, + model_raw, model_alloc, scale_offset, base_offset, + n_embd, n_hc, sinkhorn_iters, hc_eps) != 0); + TEST_ASSERT(ds4_gpu_rms_norm_weight_rows_tensor( + ref_norm, ref_out, model_raw, model_alloc, norm_weight_offset, + n_embd, n_rows, norm_eps) != 0); + TEST_ASSERT(ds4_gpu_hc_split_weighted_sum_norm_tensor( + fused_out, fused_norm, fused_split, mix, residual, + model_raw, model_alloc, scale_offset, base_offset, norm_weight_offset, + n_embd, n_hc, sinkhorn_iters, hc_eps, norm_eps) != 0); + + TEST_ASSERT(ds4_gpu_tensor_read(ref_split, 0, ref_split_host, mix_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read(fused_split, 0, fused_split_host, mix_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read(ref_out, 0, ref_out_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read(fused_out, 0, fused_out_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read(ref_norm, 0, ref_norm_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read(fused_norm, 0, fused_norm_host, out_bytes) != 0); + + const test_float_compare_stats split_stats = test_compare_float_bits( + ref_split_host, fused_split_host, (size_t)(n_rows * mix_hc)); + const test_float_compare_stats out_stats = test_compare_float_bits( + ref_out_host, fused_out_host, (size_t)n_rows * n_embd); + const test_float_compare_stats norm_stats = test_compare_float_bits( + ref_norm_host, fused_norm_host, (size_t)n_rows * n_embd); + fprintf(stderr, + "ds4-test: batch HC+RMSNorm exactness rows=%u " + "split=%zu/%llu max_ulp=%u max_abs=%g, " + "collapse=%zu/%llu max_ulp=%u max_abs=%g, " + "norm=%zu/%llu max_ulp=%u max_abs=%g\n", + n_rows, + split_stats.mismatch_count, + (unsigned long long)(n_rows * mix_hc), + split_stats.max_ulp, + split_stats.max_abs, + out_stats.mismatch_count, + (unsigned long long)((uint64_t)n_rows * n_embd), + out_stats.max_ulp, + out_stats.max_abs, + norm_stats.mismatch_count, + (unsigned long long)((uint64_t)n_rows * n_embd), + norm_stats.max_ulp, + norm_stats.max_abs); + TEST_ASSERT(split_stats.mismatch_count == 0); + TEST_ASSERT(out_stats.mismatch_count == 0); + TEST_ASSERT(norm_stats.mismatch_count == 0); + + free(mix_host); + free(residual_host); + free(ref_split_host); + free(fused_split_host); + free(ref_out_host); + free(fused_out_host); + free(ref_norm_host); + free(fused_norm_host); + ds4_gpu_tensor_free(mix); + ds4_gpu_tensor_free(residual); + ds4_gpu_tensor_free(ref_split); + ds4_gpu_tensor_free(fused_split); + ds4_gpu_tensor_free(ref_out); + ds4_gpu_tensor_free(fused_out); + ds4_gpu_tensor_free(ref_norm); + ds4_gpu_tensor_free(fused_norm); + free(model_raw); +} + +static void test_metal_output_hc_weights4_exact(void) { + const uint32_t n_hc = 4; + const float eps = 1.0e-6f; + const uint64_t page = (uint64_t)getpagesize(); + const uint64_t scale_offset = 0; + const uint64_t base_offset = page; + const uint64_t model_alloc = 2u * page; + const uint64_t bytes = n_hc * sizeof(float); + const char *force_env = "DS4_METAL_ENABLE_OUTPUT_HC_WEIGHTS4"; + const char *disable_env = "DS4_METAL_DISABLE_M3_OUTPUT_HC_WEIGHTS4"; + const char *require_env = "DS4_METAL_REQUIRE_OUTPUT_HC_WEIGHTS4"; + char *saved_force = test_save_env(force_env); + char *saved_disable = test_save_env(disable_env); + char *saved_require = test_save_env(require_env); + + void *model_raw = NULL; + TEST_ASSERT(posix_memalign( + &model_raw, (size_t)page, (size_t)model_alloc) == 0); + ds4_gpu_tensor *pre = ds4_gpu_tensor_alloc(bytes); + ds4_gpu_tensor *reference = ds4_gpu_tensor_alloc(bytes); + ds4_gpu_tensor *candidate = ds4_gpu_tensor_alloc(bytes); + TEST_ASSERT(model_raw != NULL); + TEST_ASSERT(pre != NULL); + TEST_ASSERT(reference != NULL); + TEST_ASSERT(candidate != NULL); + + size_t total_mismatch = 0; + uint32_t max_ulp = 0; + const bool allocated = model_raw && pre && reference && candidate; + if (allocated) { + memset(model_raw, 0, (size_t)model_alloc); + float *scale = (float *)((uint8_t *)model_raw + scale_offset); + float *base = (float *)((uint8_t *)model_raw + base_offset); + + for (uint32_t ci = 0; ci < 4; ci++) { + float pre_host[4]; + float ref_host[4]; + float candidate_host[4]; + for (uint32_t i = 0; i < 4; i++) { + pre_host[i] = + ((float)((int)(ci * 17u + i * 11u) - 23)) / 8.0f; + base[i] = + ((float)((int)(ci * 13u + i * 7u) - 19)) / 16.0f; + } + scale[0] = 0.62500012f + (float)ci * 0.125f; + + if (ci == 1) { + /* The first lane distinguishes the required two-rounding + * sequence from an illegally contracted multiply-add. */ + const uint32_t pre_bits = 0x4620a541u; + const uint32_t scale_bits = 0x462483bau; + const uint32_t base_bits = 0xccce790eu; + memcpy(&pre_host[0], &pre_bits, sizeof(pre_bits)); + memcpy(&scale[0], &scale_bits, sizeof(scale_bits)); + memcpy(&base[0], &base_bits, sizeof(base_bits)); + } else if (ci == 2) { + const uint32_t pre_bits[4] = { + 0x00000000u, 0x80000000u, + 0x00000001u, 0x80000001u, + }; + const uint32_t base_bits[4] = { + 0x80000000u, 0x00000000u, + 0x00800000u, 0x80800000u, + }; + for (uint32_t i = 0; i < 4; i++) { + memcpy(&pre_host[i], &pre_bits[i], sizeof(uint32_t)); + memcpy(&base[i], &base_bits[i], sizeof(uint32_t)); + } + scale[0] = -1.00000012f; + } else if (ci == 3) { + pre_host[0] = 100.0f; + pre_host[1] = -100.0f; + pre_host[2] = 88.0f; + pre_host[3] = -88.0f; + scale[0] = 1.0f; + for (uint32_t i = 0; i < 4; i++) base[i] = 0.0f; + } + + memset(ref_host, 0xa5, sizeof(ref_host)); + memset(candidate_host, 0xa5, sizeof(candidate_host)); + TEST_ASSERT(ds4_gpu_tensor_write(pre, 0, pre_host, bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + reference, 0, ref_host, bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + candidate, 0, candidate_host, bytes) != 0); + TEST_ASSERT(ds4_gpu_set_model_map(model_raw, model_alloc) != 0); + ds4_gpu_set_quality(false); + + TEST_ASSERT(unsetenv(force_env) == 0); + TEST_ASSERT(setenv(disable_env, "1", 1) == 0); + TEST_ASSERT(unsetenv(require_env) == 0); + TEST_ASSERT(ds4_gpu_output_hc_weights_tensor( + reference, pre, model_raw, model_alloc, + scale_offset, base_offset, n_hc, eps) != 0); + + TEST_ASSERT(setenv(force_env, "1", 1) == 0); + TEST_ASSERT(unsetenv(disable_env) == 0); + TEST_ASSERT(setenv(require_env, "1", 1) == 0); + TEST_ASSERT(ds4_gpu_output_hc_weights_tensor( + candidate, pre, model_raw, model_alloc, + scale_offset, base_offset, n_hc, eps) != 0); + + TEST_ASSERT(ds4_gpu_tensor_read( + reference, 0, ref_host, bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + candidate, 0, candidate_host, bytes) != 0); + const test_float_compare_stats stats = + test_compare_float_bits(ref_host, candidate_host, n_hc); + fprintf(stderr, + "ds4-test: output HC weights4 exact case=%u " + "mismatch=%zu/%u max_ulp=%u max_abs=%g\n", + ci, stats.mismatch_count, n_hc, + stats.max_ulp, stats.max_abs); + TEST_ASSERT(stats.mismatch_count == 0); + total_mismatch += stats.mismatch_count; + if (stats.max_ulp > max_ulp) max_ulp = stats.max_ulp; + } + + /* Explicit disable and quality mode win over force under strict + * selection, proving these calls cannot silently use the fast path. */ + TEST_ASSERT(setenv(force_env, "1", 1) == 0); + TEST_ASSERT(setenv(disable_env, "1", 1) == 0); + TEST_ASSERT(setenv(require_env, "1", 1) == 0); + TEST_ASSERT(ds4_gpu_output_hc_weights_tensor( + candidate, pre, model_raw, model_alloc, + scale_offset, base_offset, n_hc, eps) == 0); + TEST_ASSERT(unsetenv(disable_env) == 0); + ds4_gpu_set_quality(true); + TEST_ASSERT(ds4_gpu_output_hc_weights_tensor( + candidate, pre, model_raw, model_alloc, + scale_offset, base_offset, n_hc, eps) == 0); + ds4_gpu_set_quality(false); + } + + test_restore_env(force_env, saved_force); + test_restore_env(disable_env, saved_disable); + test_restore_env(require_env, saved_require); + fprintf(stderr, + "ds4-test: output HC weights4 total mismatch=%zu/16 max_ulp=%u\n", + total_mismatch, max_ulp); + TEST_ASSERT(total_mismatch == 0); + TEST_ASSERT(max_ulp == 0); + + ds4_gpu_tensor_free(candidate); + ds4_gpu_tensor_free(reference); + ds4_gpu_tensor_free(pre); + free(model_raw); +} + +static void test_metal_output_hc_sum_norm_exact_case( + uint32_t n_embd, + uint32_t seed) { + const uint32_t n_hc = 4; + const float norm_eps = 1.0e-6f; + const uint64_t page = (uint64_t)getpagesize(); + const uint64_t out_bytes = (uint64_t)n_embd * sizeof(float); + const uint64_t residual_bytes = (uint64_t)n_hc * out_bytes; + const uint64_t weight_bytes = (uint64_t)n_hc * sizeof(float); + const uint64_t norm_weight_offset = page; + const uint64_t model_alloc = test_round_up_u64( + norm_weight_offset + out_bytes, page); + + void *model_raw = NULL; + TEST_ASSERT(posix_memalign( + &model_raw, (size_t)page, (size_t)model_alloc) == 0); + ds4_gpu_tensor *residual = ds4_gpu_tensor_alloc(residual_bytes); + ds4_gpu_tensor *weights = ds4_gpu_tensor_alloc(weight_bytes); + ds4_gpu_tensor *ref_out = ds4_gpu_tensor_alloc(out_bytes); + ds4_gpu_tensor *fused_out = ds4_gpu_tensor_alloc(out_bytes); + ds4_gpu_tensor *ref_norm = ds4_gpu_tensor_alloc(out_bytes); + ds4_gpu_tensor *fused_norm = ds4_gpu_tensor_alloc(out_bytes); + float *residual_host = malloc((size_t)residual_bytes); + float *ref_out_host = malloc((size_t)out_bytes); + float *fused_out_host = malloc((size_t)out_bytes); + float *ref_norm_host = malloc((size_t)out_bytes); + float *fused_norm_host = malloc((size_t)out_bytes); + + TEST_ASSERT(model_raw != NULL); + TEST_ASSERT(residual != NULL); + TEST_ASSERT(weights != NULL); + TEST_ASSERT(ref_out != NULL); + TEST_ASSERT(fused_out != NULL); + TEST_ASSERT(ref_norm != NULL); + TEST_ASSERT(fused_norm != NULL); + TEST_ASSERT(residual_host != NULL); + TEST_ASSERT(ref_out_host != NULL); + TEST_ASSERT(fused_out_host != NULL); + TEST_ASSERT(ref_norm_host != NULL); + TEST_ASSERT(fused_norm_host != NULL); + + const bool allocated = model_raw && residual && weights && ref_out && + fused_out && ref_norm && fused_norm && residual_host && ref_out_host && + fused_out_host && ref_norm_host && fused_norm_host; + test_float_compare_stats out_stats = {0}; + test_float_compare_stats norm_stats = {0}; + if (allocated) { + float *norm_weight = (float *)( + (uint8_t *)model_raw + norm_weight_offset); + for (uint32_t i = 0; i < n_embd; i++) { + norm_weight[i] = + 0.5f + (float)((i * 29u + seed * 17u) % 67u) / 64.0f; + } + + for (uint32_t h = 0; h < n_hc; h++) { + for (uint32_t d = 0; d < n_embd; d++) { + const uint32_t key = d * 73u + h * 1009u + seed * 131u + + ((d >> 5u) ^ (h * 37u)); + const int value = (int)(key % 4093u) - 2046; + float v = (float)value / 1024.0f; + if ((key % 521u) == 0u) { + const uint32_t bits = (key & 1u) ? 0x80000000u : 0u; + memcpy(&v, &bits, sizeof(bits)); + } else if ((key % 523u) == 0u) { + const uint32_t bits = (key & 1u) ? 0x80000001u : 1u; + memcpy(&v, &bits, sizeof(bits)); + } + residual_host[(uint64_t)h * n_embd + d] = v; + } + } + const float weight_host[4] = { + 0.12500012f, 0.37500024f, 0.62500036f, 0.87500048f, + }; + memset(ref_out_host, 0xa5, (size_t)out_bytes); + memset(fused_out_host, 0xa5, (size_t)out_bytes); + memset(ref_norm_host, 0xa5, (size_t)out_bytes); + memset(fused_norm_host, 0xa5, (size_t)out_bytes); + + TEST_ASSERT(ds4_gpu_tensor_write( + residual, 0, residual_host, residual_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + weights, 0, weight_host, weight_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + ref_out, 0, ref_out_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_out, 0, fused_out_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + ref_norm, 0, ref_norm_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_norm, 0, fused_norm_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_set_model_map(model_raw, model_alloc) != 0); + ds4_gpu_set_quality(false); + + TEST_ASSERT(setenv( + "DS4_METAL_DISABLE_M3_OUTPUT_HC_SUM_NORM_FUSION", "1", 1) == 0); + TEST_ASSERT(unsetenv( + "DS4_METAL_ENABLE_OUTPUT_HC_SUM_NORM_FUSION") == 0); + TEST_ASSERT(unsetenv( + "DS4_METAL_REQUIRE_OUTPUT_HC_SUM_NORM_FUSION") == 0); + TEST_ASSERT(ds4_gpu_hc_weighted_sum_tensor( + ref_out, residual, weights, n_embd, n_hc) != 0); + TEST_ASSERT(ds4_gpu_rms_norm_weight_tensor( + ref_norm, ref_out, model_raw, model_alloc, + norm_weight_offset, + n_embd, norm_eps) != 0); + + TEST_ASSERT(setenv( + "DS4_METAL_ENABLE_OUTPUT_HC_SUM_NORM_FUSION", "1", 1) == 0); + TEST_ASSERT(unsetenv( + "DS4_METAL_DISABLE_M3_OUTPUT_HC_SUM_NORM_FUSION") == 0); + TEST_ASSERT(setenv( + "DS4_METAL_REQUIRE_OUTPUT_HC_SUM_NORM_FUSION", "1", 1) == 0); + TEST_ASSERT(ds4_gpu_hc_weighted_sum_norm_tensor( + fused_out, fused_norm, residual, weights, + model_raw, model_alloc, norm_weight_offset, + n_embd, n_hc, norm_eps) != 0); + + /* Force never overrides the explicit disable, quality mode, or the + * production-shape restriction. These calls must not dispatch. */ + TEST_ASSERT(unsetenv( + "DS4_METAL_REQUIRE_OUTPUT_HC_SUM_NORM_FUSION") == 0); + TEST_ASSERT(setenv( + "DS4_METAL_DISABLE_M3_OUTPUT_HC_SUM_NORM_FUSION", "1", 1) == 0); + TEST_ASSERT(ds4_gpu_hc_weighted_sum_norm_tensor( + fused_out, fused_norm, residual, weights, + model_raw, model_alloc, norm_weight_offset, + n_embd, n_hc, norm_eps) == 0); + TEST_ASSERT(unsetenv( + "DS4_METAL_DISABLE_M3_OUTPUT_HC_SUM_NORM_FUSION") == 0); + ds4_gpu_set_quality(true); + TEST_ASSERT(ds4_gpu_hc_weighted_sum_norm_tensor( + fused_out, fused_norm, residual, weights, + model_raw, model_alloc, norm_weight_offset, + n_embd, n_hc, norm_eps) == 0); + ds4_gpu_set_quality(false); + TEST_ASSERT(ds4_gpu_hc_weighted_sum_norm_tensor( + fused_out, fused_norm, residual, weights, + model_raw, model_alloc, norm_weight_offset, + 2048, n_hc, norm_eps) == 0); + + TEST_ASSERT(ds4_gpu_tensor_read( + ref_out, 0, ref_out_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + fused_out, 0, fused_out_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + ref_norm, 0, ref_norm_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + fused_norm, 0, fused_norm_host, out_bytes) != 0); + out_stats = test_compare_float_bits( + ref_out_host, fused_out_host, n_embd); + norm_stats = test_compare_float_bits( + ref_norm_host, fused_norm_host, n_embd); + } + + fprintf(stderr, + "ds4-test: output HC sum+RMSNorm exact K=%u " + "collapse=%zu/%u max_ulp=%u max_abs=%g " + "norm=%zu/%u max_ulp=%u max_abs=%g\n", + n_embd, + out_stats.mismatch_count, n_embd, + out_stats.max_ulp, out_stats.max_abs, + norm_stats.mismatch_count, n_embd, + norm_stats.max_ulp, norm_stats.max_abs); + TEST_ASSERT(out_stats.mismatch_count == 0); + TEST_ASSERT(norm_stats.mismatch_count == 0); + + free(fused_norm_host); + free(ref_norm_host); + free(fused_out_host); + free(ref_out_host); + free(residual_host); + ds4_gpu_tensor_free(fused_norm); + ds4_gpu_tensor_free(ref_norm); + ds4_gpu_tensor_free(fused_out); + ds4_gpu_tensor_free(ref_out); + ds4_gpu_tensor_free(weights); + ds4_gpu_tensor_free(residual); + free(model_raw); +} + +static void test_metal_output_hc_sum_norm_exact(void) { + const char *force_env = + "DS4_METAL_ENABLE_OUTPUT_HC_SUM_NORM_FUSION"; + const char *disable_env = + "DS4_METAL_DISABLE_M3_OUTPUT_HC_SUM_NORM_FUSION"; + const char *require_env = + "DS4_METAL_REQUIRE_OUTPUT_HC_SUM_NORM_FUSION"; + char *saved_force = test_save_env(force_env); + char *saved_disable = test_save_env(disable_env); + char *saved_require = test_save_env(require_env); + + test_metal_output_hc_sum_norm_exact_case(4096, 83); + test_metal_output_hc_sum_norm_exact_case(7168, 89); + + test_restore_env(force_env, saved_force); + test_restore_env(disable_env, saved_disable); + test_restore_env(require_env, saved_require); +} + +static void test_metal_hc_rms_scale_project_f16_exact_shape( + uint32_t in_dim, + uint32_t seed) { + const uint32_t out_dim = 24; + /* One full 32-row matmul tile plus a tail row covers both load paths. */ + const uint32_t n_rows = 33; + const float eps = 1.0e-6f; + const uint64_t page = (uint64_t)getpagesize(); + const uint64_t weight_offset = page; + const uint64_t weight_bytes = + (uint64_t)in_dim * out_dim * sizeof(uint16_t); + const uint64_t model_alloc = test_round_up_u64( + weight_offset + weight_bytes, page); + const uint64_t x_count = (uint64_t)in_dim * n_rows; + const uint64_t out_count = (uint64_t)out_dim * n_rows; + const uint64_t x_bytes = x_count * sizeof(float); + const uint64_t out_bytes = out_count * sizeof(float); + const uint64_t scale_bytes = (uint64_t)n_rows * sizeof(float); + + void *model_raw = NULL; + TEST_ASSERT(posix_memalign( + &model_raw, (size_t)page, (size_t)model_alloc) == 0); + ds4_gpu_tensor *x = ds4_gpu_tensor_alloc(x_bytes); + ds4_gpu_tensor *ref_norm = ds4_gpu_tensor_alloc(x_bytes); + ds4_gpu_tensor *ref_out = ds4_gpu_tensor_alloc(out_bytes); + ds4_gpu_tensor *scaled_out = ds4_gpu_tensor_alloc(out_bytes); + /* Deliberately too small for the full-RMS fallback. */ + ds4_gpu_tensor *scale_scratch = ds4_gpu_tensor_alloc(scale_bytes); + float *x_host = malloc((size_t)x_bytes); + float *ref_norm_host = malloc((size_t)x_bytes); + float *ref_out_host = malloc((size_t)out_bytes); + float *scaled_out_host = malloc((size_t)out_bytes); + float *scale_host = malloc((size_t)scale_bytes); + float *expected_scale = malloc((size_t)scale_bytes); + + TEST_ASSERT(model_raw != NULL); + TEST_ASSERT(x != NULL); + TEST_ASSERT(ref_norm != NULL); + TEST_ASSERT(ref_out != NULL); + TEST_ASSERT(scaled_out != NULL); + TEST_ASSERT(scale_scratch != NULL); + TEST_ASSERT(x_host != NULL); + TEST_ASSERT(ref_norm_host != NULL); + TEST_ASSERT(ref_out_host != NULL); + TEST_ASSERT(scaled_out_host != NULL); + TEST_ASSERT(scale_host != NULL); + TEST_ASSERT(expected_scale != NULL); + + const bool allocated = model_raw && x && ref_norm && ref_out && + scaled_out && scale_scratch && x_host && ref_norm_host && + ref_out_host && scaled_out_host && scale_host && expected_scale; + test_float_compare_stats scale_stats = {0}; + test_float_compare_stats out_stats = {0}; + if (allocated) { + memset(model_raw, 0, (size_t)model_alloc); + uint16_t *weights = (uint16_t *)((uint8_t *)model_raw + weight_offset); + for (uint32_t o = 0; o < out_dim; o++) { + for (uint32_t i = 0; i < in_dim; i++) { + const uint32_t key = i * 37u + o * 1009u + seed * 53u + + ((i >> 4u) ^ (o * 17u)); + uint16_t bits; + if (key % 257u == 0u) { + bits = (key & 1u) ? 0x8000u : 0x0000u; + } else if (key % 263u == 0u) { + bits = (key & 1u) ? 0x8001u : 0x0001u; + } else if (key % 269u == 0u) { + bits = (key & 1u) ? 0x8400u : 0x0400u; + } else { + const int value = (int)(key % 127u) - 63; + bits = test_float_to_f16((float)value / 128.0f); + } + weights[(uint64_t)o * in_dim + i] = bits; + } + } + + static const uint32_t rounding_bits[] = { + 0x3f800fffu, 0x3f801000u, 0x3f801001u, + 0x3f802fffu, 0x3f803000u, 0x3f803001u, + 0xbf800fffu, 0xbf801000u, 0xbf801001u, + 0x3eaaaaabu, 0xbeaaaaabu, + }; + static const int sentinel_exp[] = { -2, -10, 0, -4, -6, -1 }; + for (uint32_t row = 0; row < n_rows; row++) { + for (uint32_t i = 0; i < in_dim; i++) { + const uint32_t key = i * 131u + row * 977u + seed * 71u + + ((i >> 3u) ^ (row * 29u)); + const float sign = (key & 1u) ? -1.0f : 1.0f; + float value; + switch (row % 6u) { + case 0: + value = (float)((int)(key % 4093u) - 2046) / 512.0f; + break; + case 1: + value = sign * ldexpf( + (float)(1u + (key & 7u)) / 8.0f, -19); + break; + case 2: + value = i % 257u == row % 257u + ? sign * (16.0f + (float)(key & 7u)) + : sign * (float)(1u + (key & 31u)) / 8192.0f; + break; + case 3: { + const uint32_t bits = rounding_bits[ + key % (sizeof(rounding_bits) / + sizeof(rounding_bits[0]))]; + memcpy(&value, &bits, sizeof(value)); + break; + } + case 4: + if ((key & 7u) == 0u) { + const uint32_t bits = (key & 8u) ? 0x80000000u : 0u; + memcpy(&value, &bits, sizeof(value)); + } else { + value = sign * ldexpf( + 1.0f + (float)(key & 3u) * 0.25f, + (int)((key >> 4u) % 14u) - 10); + } + break; + default: + value = sign * (float)(1u + (key % 251u)) / 256.0f; + break; + } + x_host[(uint64_t)row * in_dim + i] = value; + } + x_host[(uint64_t)row * in_dim] = + ldexpf(1.0f, sentinel_exp[row % 6u]); + } + + for (uint64_t i = 0; i < out_count; i++) { + const uint32_t bits = 0x7fc01000u + (uint32_t)(i & 0xfffu); + memcpy(scaled_out_host + i, &bits, sizeof(bits)); + } + for (uint32_t row = 0; row < n_rows; row++) { + const uint32_t bits = 0x7fc02000u + row; + memcpy(scale_host + row, &bits, sizeof(bits)); + } + + TEST_ASSERT(ds4_gpu_tensor_write(x, 0, x_host, x_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + scaled_out, 0, scaled_out_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + scale_scratch, 0, scale_host, scale_bytes) != 0); + TEST_ASSERT(ds4_gpu_set_model_map(model_raw, model_alloc) != 0); + ds4_gpu_set_quality(false); + + int ref_begun = ds4_gpu_begin_commands(); + int ref_ok = ref_begun; + if (ref_ok) ref_ok = ds4_gpu_rms_norm_plain_rows_tensor( + ref_norm, x, in_dim, n_rows, eps); + if (ref_ok) ref_ok = ds4_gpu_matmul_f16_tensor( + ref_out, model_raw, model_alloc, weight_offset, + in_dim, out_dim, ref_norm, n_rows); + const int ref_end = ref_begun ? ds4_gpu_end_commands() : 0; + TEST_ASSERT(ref_ok != 0); + TEST_ASSERT(ref_end != 0); + + int scaled_begun = ds4_gpu_begin_commands(); + int scaled_ok = scaled_begun; + if (scaled_ok) scaled_ok = ds4_gpu_hc_rms_scale_project_f16_tensor( + scaled_out, scale_scratch, model_raw, model_alloc, weight_offset, + in_dim, out_dim, x, n_rows, eps); + const int scaled_end = scaled_begun ? ds4_gpu_end_commands() : 0; + TEST_ASSERT(scaled_ok != 0); + TEST_ASSERT(scaled_end != 0); + + TEST_ASSERT(ds4_gpu_tensor_read( + ref_norm, 0, ref_norm_host, x_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + ref_out, 0, ref_out_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + scaled_out, 0, scaled_out_host, out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + scale_scratch, 0, scale_host, scale_bytes) != 0); + + for (uint32_t row = 0; row < n_rows; row++) { + expected_scale[row] = ldexpf( + ref_norm_host[(uint64_t)row * in_dim], + -sentinel_exp[row % 6u]); + TEST_ASSERT(isfinite(expected_scale[row])); + TEST_ASSERT(isfinite(scale_host[row])); + } + scale_stats = test_compare_float_bits( + expected_scale, scale_host, n_rows); + out_stats = test_compare_float_bits( + ref_out_host, scaled_out_host, (size_t)out_count); + } + + fprintf(stderr, + "ds4-test: HC RMS scale F16 projection exact K=%u rows=%u " + "scale=%zu/%u max_ulp=%u projection=%zu/%llu max_ulp=%u\n", + in_dim, n_rows, + scale_stats.mismatch_count, n_rows, scale_stats.max_ulp, + out_stats.mismatch_count, (unsigned long long)out_count, + out_stats.max_ulp); + TEST_ASSERT(scale_stats.mismatch_count == 0); + TEST_ASSERT(scale_stats.max_ulp == 0); + TEST_ASSERT(out_stats.mismatch_count == 0); + TEST_ASSERT(out_stats.max_ulp == 0); + + free(expected_scale); + free(scale_host); + free(scaled_out_host); + free(ref_out_host); + free(ref_norm_host); + free(x_host); + ds4_gpu_tensor_free(scale_scratch); + ds4_gpu_tensor_free(scaled_out); + ds4_gpu_tensor_free(ref_out); + ds4_gpu_tensor_free(ref_norm); + ds4_gpu_tensor_free(x); + free(model_raw); +} + +static void test_metal_hc_rms_scale_project_f16_exact(void) { + const char *force_env = "DS4_METAL_ENABLE_HC_RMS_SCALE_PROJ"; + const char *disable_env = "DS4_METAL_DISABLE_M3_HC_RMS_SCALE_PROJ"; + char *saved_force = test_save_env(force_env); + char *saved_disable = test_save_env(disable_env); + + TEST_ASSERT(setenv(force_env, "1", 1) == 0); + TEST_ASSERT(unsetenv(disable_env) == 0); + test_metal_hc_rms_scale_project_f16_exact_shape(16384u, 59u); + test_metal_hc_rms_scale_project_f16_exact_shape(28672u, 61u); + + test_restore_env(force_env, saved_force); + test_restore_env(disable_env, saved_disable); +} + +static void test_metal_router_simd_finalize_exact(void) { + typedef struct { + const char *name; + bool has_bias; + uint32_t pattern; + } router_case; + static const router_case cases[] = { + { "unique", false, 0 }, + { "bias", true, 0 }, + { "top6-ties", false, 1 }, + { "signed-zero-extremes", false, 2 }, + { "clamp-underflow", false, 3 }, + { "sum-rounding", true, 4 }, + }; + const uint32_t n_expert = 256; + const uint32_t n_used = 6; + const uint32_t repeats = 2; + const uint64_t probs_bytes = (uint64_t)n_expert * sizeof(float); + const uint64_t selected_bytes = (uint64_t)n_used * sizeof(int32_t); + const uint64_t weights_bytes = (uint64_t)n_used * sizeof(float); + const uint64_t page = (uint64_t)getpagesize(); + const char *force_env = "DS4_METAL_ENABLE_ROUTER_SIMD_FINALIZE"; + const char *disable_env = + "DS4_METAL_DISABLE_M3_ROUTER_SIMD_FINALIZE"; + const char *weights_force_env = + "DS4_METAL_ENABLE_ROUTER_SIMD_WEIGHTS_FUSION"; + const char *weights_disable_env = + "DS4_METAL_DISABLE_M3_ROUTER_SIMD_WEIGHTS_FUSION"; + const char *select_disable_env = + "DS4_METAL_DISABLE_ROUTER_SELECT_FUSION"; + + void *model_raw = NULL; + TEST_ASSERT(posix_memalign(&model_raw, (size_t)page, (size_t)page) == 0); + ds4_gpu_tensor *logits = ds4_gpu_tensor_alloc(probs_bytes); + ds4_gpu_tensor *ref_selected = ds4_gpu_tensor_alloc(selected_bytes); + ds4_gpu_tensor *simd_selected = ds4_gpu_tensor_alloc(selected_bytes); + ds4_gpu_tensor *ref_weights = ds4_gpu_tensor_alloc(weights_bytes); + ds4_gpu_tensor *simd_weights = ds4_gpu_tensor_alloc(weights_bytes); + ds4_gpu_tensor *ref_probs = ds4_gpu_tensor_alloc(probs_bytes); + ds4_gpu_tensor *simd_probs = ds4_gpu_tensor_alloc(probs_bytes); + float *logits_host = malloc((size_t)probs_bytes); + float *ref_probs_host = malloc((size_t)probs_bytes); + float *simd_probs_host = malloc((size_t)probs_bytes); + int32_t ref_selected_host[6]; + int32_t simd_selected_host[6]; + int32_t unique_selected_host[6]; + float ref_weights_host[6]; + float simd_weights_host[6]; + TEST_ASSERT(model_raw != NULL); + TEST_ASSERT(logits != NULL); + TEST_ASSERT(ref_selected != NULL); + TEST_ASSERT(simd_selected != NULL); + TEST_ASSERT(ref_weights != NULL); + TEST_ASSERT(simd_weights != NULL); + TEST_ASSERT(ref_probs != NULL); + TEST_ASSERT(simd_probs != NULL); + TEST_ASSERT(logits_host != NULL); + TEST_ASSERT(ref_probs_host != NULL); + TEST_ASSERT(simd_probs_host != NULL); + + char *saved_force = test_save_env(force_env); + char *saved_disable = test_save_env(disable_env); + char *saved_weights_force = test_save_env(weights_force_env); + char *saved_weights_disable = test_save_env(weights_disable_env); + char *saved_select_disable = test_save_env(select_disable_env); + size_t total_selected_mismatch = 0; + size_t total_weights_mismatch = 0; + size_t total_probs_mismatch = 0; + const bool allocated = model_raw && logits && ref_selected && + simd_selected && ref_weights && simd_weights && ref_probs && + simd_probs && logits_host && ref_probs_host && simd_probs_host; + if (allocated) { + memset(model_raw, 0, (size_t)page); + float *bias = model_raw; + for (uint32_t i = 0; i < n_used; i++) { + bias[200u + i] = 16.0f - (float)i; + } + TEST_ASSERT(ds4_gpu_set_model_map(model_raw, page) != 0); + ds4_gpu_set_quality(false); + TEST_ASSERT(setenv(force_env, "1", 1) == 0); + TEST_ASSERT(setenv(weights_force_env, "1", 1) == 0); + TEST_ASSERT(unsetenv(select_disable_env) == 0); + + for (size_t ci = 0; ci < sizeof(cases) / sizeof(cases[0]); ci++) { + const router_case *c = &cases[ci]; + for (uint32_t i = 0; i < n_expert; i++) { + const int value = + (int)((i * 47u + (i ^ (i >> 3u)) * 13u) % 257u) - 128; + logits_host[i] = (float)value / 32.0f; + } + if (c->pattern == 1) { + static const uint32_t tied[] = { + 7u, 19u, 43u, 71u, 103u, 149u, 211u, 239u, + }; + for (uint32_t i = 0; i < n_expert; i++) { + logits_host[i] = -7.0f - (float)(i % 17u) / 64.0f; + } + for (size_t i = 0; i < sizeof(tied) / sizeof(tied[0]); i++) { + logits_host[tied[i]] = 1.0f; + } + } else if (c->pattern == 2) { + for (uint32_t i = 0; i < n_expert; i++) { + const uint32_t zero_bits = (i & 1u) ? 0x80000000u : 0u; + memcpy(&logits_host[i], &zero_bits, sizeof(zero_bits)); + } + logits_host[3] = 80.0f; + logits_host[17] = 40.0f; + logits_host[61] = 20.0f; + logits_host[127] = -20.0f; + logits_host[193] = -40.0f; + logits_host[251] = -80.0f; + } else if (c->pattern == 3) { + for (uint32_t i = 0; i < n_expert; i++) { + logits_host[i] = -30.0f - (float)(i % 11u); + } + } else if (c->pattern == 4) { + static const float rounding_logits[6] = { + -0.6356699467f, + -0.8182631135f, + -2.7906901836f, + -3.4414808750f, + -3.4991359711f, + -3.1251864433f, + }; + for (uint32_t i = 0; i < n_expert; i++) { + logits_host[i] = -20.0f; + } + for (uint32_t i = 0; i < n_used; i++) { + logits_host[200u + i] = rounding_logits[i]; + } + } + TEST_ASSERT(ds4_gpu_tensor_write( + logits, 0, logits_host, probs_bytes) != 0); + + TEST_ASSERT(setenv(disable_env, "1", 1) == 0); + TEST_ASSERT(ds4_gpu_router_select_tensor( + ref_selected, ref_weights, ref_probs, + model_raw, page, 0, 0, 1, 0, + n_expert, n_used, 1.5f, 1, 0, + c->has_bias, false, logits) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + ref_selected, 0, ref_selected_host, selected_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + ref_weights, 0, ref_weights_host, weights_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + ref_probs, 0, ref_probs_host, probs_bytes) != 0); + + if (ci == 0) { + memcpy(unique_selected_host, + ref_selected_host, + sizeof(unique_selected_host)); + } else if (c->has_bias) { + TEST_ASSERT(memcmp(unique_selected_host, + ref_selected_host, + sizeof(unique_selected_host)) != 0); + } + + TEST_ASSERT(unsetenv(disable_env) == 0); + for (uint32_t rep = 0; rep <= repeats; rep++) { + const bool fused_weights = rep != 0; + if (fused_weights) { + TEST_ASSERT(unsetenv(weights_disable_env) == 0); + } else { + TEST_ASSERT(setenv(weights_disable_env, "1", 1) == 0); + } + TEST_ASSERT(ds4_gpu_router_select_tensor( + simd_selected, simd_weights, simd_probs, + model_raw, page, 0, 0, 1, 0, + n_expert, n_used, 1.5f, 1, 0, + c->has_bias, false, logits) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + simd_selected, 0, simd_selected_host, + selected_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + simd_weights, 0, simd_weights_host, + weights_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + simd_probs, 0, simd_probs_host, probs_bytes) != 0); + + size_t selected_mismatch = 0; + for (uint32_t i = 0; i < n_used; i++) { + if (ref_selected_host[i] != simd_selected_host[i]) { + selected_mismatch++; + } + } + const test_float_compare_stats weights_stats = + test_compare_float_bits( + ref_weights_host, simd_weights_host, n_used); + const test_float_compare_stats probs_stats = + test_compare_float_bits( + ref_probs_host, simd_probs_host, n_expert); + fprintf(stderr, + "ds4-test: router SIMD finalize exactness " + "case=%s mode=%s rep=%u selected=%zu/%u weights=%zu/%u " + "probs=%zu/%u max_weight_ulp=%u max_prob_ulp=%u\n", + c->name, + fused_weights ? "fused-weights" : "split-weights", + fused_weights ? rep - 1u : 0u, + selected_mismatch, + n_used, + weights_stats.mismatch_count, + n_used, + probs_stats.mismatch_count, + n_expert, + weights_stats.max_ulp, + probs_stats.max_ulp); + TEST_ASSERT(selected_mismatch == 0); + TEST_ASSERT(weights_stats.mismatch_count == 0); + TEST_ASSERT(probs_stats.mismatch_count == 0); + total_selected_mismatch += selected_mismatch; + total_weights_mismatch += weights_stats.mismatch_count; + total_probs_mismatch += probs_stats.mismatch_count; + } + } + } + + test_restore_env(force_env, saved_force); + test_restore_env(disable_env, saved_disable); + test_restore_env(weights_force_env, saved_weights_force); + test_restore_env(weights_disable_env, saved_weights_disable); + test_restore_env(select_disable_env, saved_select_disable); + fprintf(stderr, + "ds4-test: router SIMD finalize total selected=%zu/%zu " + "weights=%zu/%zu probs=%zu/%zu\n", + total_selected_mismatch, + (sizeof(cases) / sizeof(cases[0])) * (repeats + 1u) * (size_t)n_used, + total_weights_mismatch, + (sizeof(cases) / sizeof(cases[0])) * (repeats + 1u) * (size_t)n_used, + total_probs_mismatch, + (sizeof(cases) / sizeof(cases[0])) * (repeats + 1u) * (size_t)n_expert); + TEST_ASSERT(total_selected_mismatch == 0); + TEST_ASSERT(total_weights_mismatch == 0); + TEST_ASSERT(total_probs_mismatch == 0); + + free(simd_probs_host); + free(ref_probs_host); + free(logits_host); + ds4_gpu_tensor_free(simd_probs); + ds4_gpu_tensor_free(ref_probs); + ds4_gpu_tensor_free(simd_weights); + ds4_gpu_tensor_free(ref_weights); + ds4_gpu_tensor_free(simd_selected); + ds4_gpu_tensor_free(ref_selected); + ds4_gpu_tensor_free(logits); + free(model_raw); +} + +static void test_metal_router_weights_batch_exact(void) { + typedef struct { + const char *name; + uint32_t n_tokens; + bool has_bias; + bool hash_mode; + uint32_t pattern; + } router_batch_case; + static const router_batch_case cases[] = { + { "rows2-unique", 2, false, false, 0 }, + { "rows17-bias-ties", 17, true, false, 1 }, + { "rows129-hash-duplicates", 129, false, true, 2 }, + { "rows2048-typical", 2048, false, false, 0 }, + { "rows3-clamp-underflow", 3, false, false, 3 }, + }; + const uint32_t n_expert = 256; + const uint32_t n_used = 6; + const uint32_t max_tokens = 2048; + const uint32_t hash_rows = 64; + const uint32_t repeats = 2; + const uint64_t logits_bytes = + (uint64_t)max_tokens * n_expert * sizeof(float); + const uint64_t selected_bytes = + (uint64_t)max_tokens * n_used * sizeof(int32_t); + const uint64_t weights_bytes = + (uint64_t)max_tokens * n_used * sizeof(float); + const uint64_t tokens_bytes = (uint64_t)max_tokens * sizeof(int32_t); + const uint64_t page = (uint64_t)getpagesize(); + const uint64_t bias_offset = 0; + const uint64_t hash_offset = 2048; + const char *force_env = + "DS4_METAL_ENABLE_ROUTER_WEIGHTS_BATCH_FUSION"; + const char *disable_env = + "DS4_METAL_DISABLE_M3_ROUTER_WEIGHTS_BATCH_FUSION"; + const char *select_disable_env = + "DS4_METAL_DISABLE_ROUTER_SELECT_FUSION"; + + void *model_raw = NULL; + TEST_ASSERT(posix_memalign(&model_raw, (size_t)page, (size_t)page) == 0); + ds4_gpu_tensor *logits = ds4_gpu_tensor_alloc(logits_bytes); + ds4_gpu_tensor *tokens = ds4_gpu_tensor_alloc(tokens_bytes); + ds4_gpu_tensor *ref_selected = ds4_gpu_tensor_alloc(selected_bytes); + ds4_gpu_tensor *batch_selected = ds4_gpu_tensor_alloc(selected_bytes); + ds4_gpu_tensor *ref_weights = ds4_gpu_tensor_alloc(weights_bytes); + ds4_gpu_tensor *batch_weights = ds4_gpu_tensor_alloc(weights_bytes); + ds4_gpu_tensor *ref_probs = ds4_gpu_tensor_alloc(logits_bytes); + ds4_gpu_tensor *batch_probs = ds4_gpu_tensor_alloc(logits_bytes); + float *logits_host = malloc((size_t)logits_bytes); + int32_t *tokens_host = malloc((size_t)tokens_bytes); + int32_t *ref_selected_host = malloc((size_t)selected_bytes); + int32_t *batch_selected_host = malloc((size_t)selected_bytes); + float *ref_weights_host = malloc((size_t)weights_bytes); + float *batch_weights_host = malloc((size_t)weights_bytes); + float *ref_probs_host = malloc((size_t)logits_bytes); + float *batch_probs_host = malloc((size_t)logits_bytes); + TEST_ASSERT(model_raw != NULL); + TEST_ASSERT(logits != NULL); + TEST_ASSERT(tokens != NULL); + TEST_ASSERT(ref_selected != NULL); + TEST_ASSERT(batch_selected != NULL); + TEST_ASSERT(ref_weights != NULL); + TEST_ASSERT(batch_weights != NULL); + TEST_ASSERT(ref_probs != NULL); + TEST_ASSERT(batch_probs != NULL); + TEST_ASSERT(logits_host != NULL); + TEST_ASSERT(tokens_host != NULL); + TEST_ASSERT(ref_selected_host != NULL); + TEST_ASSERT(batch_selected_host != NULL); + TEST_ASSERT(ref_weights_host != NULL); + TEST_ASSERT(batch_weights_host != NULL); + TEST_ASSERT(ref_probs_host != NULL); + TEST_ASSERT(batch_probs_host != NULL); + + char *saved_force = test_save_env(force_env); + char *saved_disable = test_save_env(disable_env); + char *saved_select_disable = test_save_env(select_disable_env); + size_t total_selected_mismatch = 0; + size_t total_weights_mismatch = 0; + size_t total_probs_mismatch = 0; + const bool allocated = model_raw && logits && tokens && ref_selected && + batch_selected && ref_weights && batch_weights && ref_probs && + batch_probs && logits_host && tokens_host && ref_selected_host && + batch_selected_host && ref_weights_host && batch_weights_host && + ref_probs_host && batch_probs_host; + if (allocated) { + memset(model_raw, 0, (size_t)page); + float *bias = (float *)((uint8_t *)model_raw + bias_offset); + int32_t *hash = (int32_t *)((uint8_t *)model_raw + hash_offset); + for (uint32_t i = 0; i < n_expert; i++) { + bias[i] = (float)((int)((i * 29u) % 67u) - 33) / 16.0f; + } + for (uint32_t row = 0; row < hash_rows; row++) { + const int32_t base = (int32_t)((row * 37u) % n_expert); + hash[(uint64_t)row * n_used + 0u] = base; + hash[(uint64_t)row * n_used + 1u] = base; + hash[(uint64_t)row * n_used + 2u] = (base + 19) % (int32_t)n_expert; + hash[(uint64_t)row * n_used + 3u] = (base + 43) % (int32_t)n_expert; + hash[(uint64_t)row * n_used + 4u] = (base + 43) % (int32_t)n_expert; + hash[(uint64_t)row * n_used + 5u] = (base + 101) % (int32_t)n_expert; + } + for (uint32_t row = 0; row < max_tokens; row++) { + tokens_host[row] = (int32_t)((row * 13u + 7u) % hash_rows); + } + + TEST_ASSERT(ds4_gpu_set_model_map(model_raw, page) != 0); + ds4_gpu_set_quality(false); + TEST_ASSERT(ds4_gpu_tensor_write( + tokens, 0, tokens_host, tokens_bytes) != 0); + TEST_ASSERT(setenv(force_env, "1", 1) == 0); + TEST_ASSERT(unsetenv(select_disable_env) == 0); + + for (size_t ci = 0; ci < sizeof(cases) / sizeof(cases[0]); ci++) { + const router_batch_case *c = &cases[ci]; + const size_t prob_count = (size_t)c->n_tokens * n_expert; + const size_t route_count = (size_t)c->n_tokens * n_used; + const uint64_t case_probs_bytes = prob_count * sizeof(float); + const uint64_t case_selected_bytes = route_count * sizeof(int32_t); + const uint64_t case_weights_bytes = route_count * sizeof(float); + for (uint32_t row = 0; row < c->n_tokens; row++) { + for (uint32_t expert = 0; expert < n_expert; expert++) { + const size_t index = (size_t)row * n_expert + expert; + const int value = + (int)((row * 131u + expert * 47u + + (expert ^ (expert >> 3u)) * 13u) % 513u) - + 256; + logits_host[index] = (float)value / 64.0f; + if (c->pattern == 1) { + logits_host[index] = + -7.0f - (float)((row + expert) % 17u) / 64.0f; + } else if (c->pattern == 3) { + logits_host[index] = + -30.0f - (float)((row * 17u + expert * 29u) % 11u); + } + } + if (c->pattern == 1) { + static const uint32_t tied[] = { + 7u, 19u, 43u, 71u, 103u, 149u, 211u, 239u, + }; + for (size_t i = 0; i < sizeof(tied) / sizeof(tied[0]); i++) { + logits_host[(size_t)row * n_expert + tied[i]] = 1.0f; + } + } + } + TEST_ASSERT(ds4_gpu_tensor_write( + logits, 0, logits_host, case_probs_bytes) != 0); + + TEST_ASSERT(setenv(disable_env, "1", 1) == 0); + TEST_ASSERT(ds4_gpu_router_select_batch_tensor( + ref_selected, ref_weights, ref_probs, + model_raw, page, bias_offset, hash_offset, hash_rows, + 1, 0, c->has_bias, c->hash_mode, logits, tokens, + n_expert, n_used, 1.5f, c->n_tokens) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + ref_selected, 0, ref_selected_host, case_selected_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + ref_weights, 0, ref_weights_host, case_weights_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + ref_probs, 0, ref_probs_host, case_probs_bytes) != 0); + + if (c->hash_mode) { + for (uint32_t row = 0; row < c->n_tokens; row++) { + const uint32_t hash_row = (uint32_t)tokens_host[row]; + for (uint32_t lane = 0; lane < n_used; lane++) { + TEST_ASSERT( + ref_selected_host[(size_t)row * n_used + lane] == + hash[(uint64_t)hash_row * n_used + lane]); + } + } + } + if (c->pattern == 3) { + for (uint32_t row = 0; row < c->n_tokens; row++) { + float sum = 0.0f; + for (uint32_t lane = 0; lane < n_used; lane++) { + const int32_t expert = + ref_selected_host[(size_t)row * n_used + lane]; + sum += ref_probs_host[(size_t)row * n_expert + + (uint32_t)expert]; + } + TEST_ASSERT(sum < 6.103515625e-5f); + } + } + + TEST_ASSERT(unsetenv(disable_env) == 0); + for (uint32_t rep = 0; rep < repeats; rep++) { + TEST_ASSERT(ds4_gpu_router_select_batch_tensor( + batch_selected, batch_weights, batch_probs, + model_raw, page, bias_offset, hash_offset, hash_rows, + 1, 0, c->has_bias, c->hash_mode, logits, tokens, + n_expert, n_used, 1.5f, c->n_tokens) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + batch_selected, 0, batch_selected_host, + case_selected_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + batch_weights, 0, batch_weights_host, + case_weights_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + batch_probs, 0, batch_probs_host, + case_probs_bytes) != 0); + + size_t selected_mismatch = 0; + for (size_t i = 0; i < route_count; i++) { + if (ref_selected_host[i] != batch_selected_host[i]) { + selected_mismatch++; + } + } + const test_float_compare_stats weights_stats = + test_compare_float_bits( + ref_weights_host, batch_weights_host, route_count); + const test_float_compare_stats probs_stats = + test_compare_float_bits( + ref_probs_host, batch_probs_host, prob_count); + fprintf(stderr, + "ds4-test: router batch weights exactness " + "case=%s rep=%u selected=%zu/%zu weights=%zu/%zu " + "probs=%zu/%zu max_weight_ulp=%u max_prob_ulp=%u\n", + c->name, + rep, + selected_mismatch, + route_count, + weights_stats.mismatch_count, + route_count, + probs_stats.mismatch_count, + prob_count, + weights_stats.max_ulp, + probs_stats.max_ulp); + TEST_ASSERT(selected_mismatch == 0); + TEST_ASSERT(weights_stats.mismatch_count == 0); + TEST_ASSERT(probs_stats.mismatch_count == 0); + total_selected_mismatch += selected_mismatch; + total_weights_mismatch += weights_stats.mismatch_count; + total_probs_mismatch += probs_stats.mismatch_count; + } + } + } + + test_restore_env(force_env, saved_force); + test_restore_env(disable_env, saved_disable); + test_restore_env(select_disable_env, saved_select_disable); + fprintf(stderr, + "ds4-test: router batch weights total selected=%zu " + "weights=%zu probs=%zu\n", + total_selected_mismatch, + total_weights_mismatch, + total_probs_mismatch); + TEST_ASSERT(total_selected_mismatch == 0); + TEST_ASSERT(total_weights_mismatch == 0); + TEST_ASSERT(total_probs_mismatch == 0); + + free(batch_probs_host); + free(ref_probs_host); + free(batch_weights_host); + free(ref_weights_host); + free(batch_selected_host); + free(ref_selected_host); + free(tokens_host); + free(logits_host); + ds4_gpu_tensor_free(batch_probs); + ds4_gpu_tensor_free(ref_probs); + ds4_gpu_tensor_free(batch_weights); + ds4_gpu_tensor_free(ref_weights); + ds4_gpu_tensor_free(batch_selected); + ds4_gpu_tensor_free(ref_selected); + ds4_gpu_tensor_free(tokens); + ds4_gpu_tensor_free(logits); + free(model_raw); +} +#endif + static void test_metal_kernel_group(void) { test_metal_f16_matvec_fast_nr0_4(); test_metal_f16_prefill_matmul(); test_metal_q8_0_prefill_matmul(); + test_metal_q8_0_decode_pair_exact(); +#if defined(__APPLE__) + test_metal_q8_0_output_nr4_exact(); + test_metal_f16_compressor_pair_state_store_exact(); + test_metal_compressor_ape_add_exact(); + test_metal_compressor_ratio4_pack_exact(); + test_metal_compressor_ratio4_replay_pack_exact(); + test_metal_compressor_ratio4_direct_pool_exact(); + test_metal_inplace_rope_pair_exact(); + test_metal_contiguous_f32_f16_roundtrip_exact(); + test_metal_gathered_kv_stage_exact(); + test_metal_contiguous_compressed_f16_attention_exact(); + test_metal_persistent_zero_attention_mask_exact(); + test_metal_zero_prefix_prefill_mask_cache_exact(); + test_metal_hc_split_weighted_sum_norm_batch_exact(); + test_metal_output_hc_weights4_exact(); + test_metal_output_hc_sum_norm_exact(); + test_metal_hc_rms_scale_project_f16_exact(); + test_metal_router_simd_finalize_exact(); + test_metal_router_weights_batch_exact(); +#endif } static void test_metal_short_prefill_ratio4(void) {