diff --git a/ds4.c b/ds4.c index b54075539..d30db5a70 100644 --- a/ds4.c +++ b/ds4.c @@ -16109,6 +16109,11 @@ typedef struct { bool prefill_has_visual; const ds4_vision_span *prefill_vision_spans; size_t prefill_vision_span_count; + /* Greedy chain decode: when non-NULL, the one-token router select reads + * the token id from this device-resident view (a greedy-chain ring slot) + * instead of the encode-time constant, and the host-side hash-selected + * override is skipped (it is unused by the resident fixed-route MoE). */ + ds4_gpu_tensor *chain_token_view; } ds4_gpu_graph; /* Tensors that are temporary for chunked prefill and grouped multi-session @@ -22740,6 +22745,11 @@ static bool metal_graph_encode_decode_layer_phase( bool ok = true; const bool decode_stage_profile = metal_graph_decode_stage_profile_enabled(il); + /* Stage counters sample GPU timestamps inside the open encoder, so the + * schedule must stay the production one; only the end-and-wait profiler + * serializes and therefore disqualifies concurrent dispatch. */ + const bool decode_stage_serializing = + decode_stage_profile && !ds4_gpu_stage_counters_enabled(); double decode_stage_t0 = decode_stage_profile ? now_sec() : 0.0; const bool fuse_shared_gate_up = !g->quality && @@ -22762,7 +22772,7 @@ static bool metal_graph_encode_decode_layer_phase( !g->quality && g->tp_world < 2 && !g->ssd_streaming && !g->ssd_streaming_cold && !g->cuda_tp_decode && !g->cuda_tp_moe && !g->cuda_tp_shared && - !decode_stage_profile && + !decode_stage_serializing && !metal_graph_directional_steering_ffn_enabled(g) && metal_graph_debug_get_config()->prefix == NULL && getenv("DS4_METAL_MOE_ONE_STAGE_PROFILE") == NULL && @@ -24673,7 +24683,25 @@ static bool metal_graph_encode_decode_layer_phase( if (ok && router_shared_done == 0 && router_only_done == 0) ok = metal_graph_matmul_plain_tensor(metal_graph_router_logits(g), model, layer->ffn_gate_inp, DS4_N_EMBD, DS4_N_EXPERT, metal_graph_ffn_norm(g), 1); - if (ok && !router_project_select_fused) + if (ok && !router_project_select_fused) { + if (g->chain_token_view) { + /* Greedy chain: the token id is device-resident; the select + * kernel reads it from the ring slot (same hash-row gather). */ + ok = ds4_gpu_router_select_tensor_devtoken(metal_graph_router_selected(g), metal_graph_router_weights(g), metal_graph_router_probs(g), + model->map, model->size, + layer->ffn_exp_probs_b ? layer->ffn_exp_probs_b->abs_offset : 0, + layer->ffn_gate_tid2eid ? layer->ffn_gate_tid2eid->abs_offset : 0, + layer->ffn_gate_tid2eid ? (uint32_t)layer->ffn_gate_tid2eid->dim[1] : 0, + g->chain_token_view, + DS4_N_EXPERT, + DS4_N_EXPERT_USED, + DS4_EXPERT_WEIGHT_SCALE, + 0, + 0, + layer->ffn_exp_probs_b != NULL, + layer->ffn_gate_tid2eid != NULL, + metal_graph_router_logits(g)) != 0; + } else { ok = ds4_gpu_router_select_tensor(metal_graph_router_selected(g), metal_graph_router_weights(g), metal_graph_router_probs(g), model->map, model->size, layer->ffn_exp_probs_b ? layer->ffn_exp_probs_b->abs_offset : 0, @@ -24688,8 +24716,10 @@ static bool metal_graph_encode_decode_layer_phase( layer->ffn_exp_probs_b != NULL, layer->ffn_gate_tid2eid != NULL, metal_graph_router_logits(g)) != 0; + } + } } - if (ok) ok = metal_graph_decode_set_hash_selected_override(model, + if (ok && !g->chain_token_view) ok = metal_graph_decode_set_hash_selected_override(model, layer, il, (uint32_t)token, @@ -24994,7 +25024,8 @@ static bool metal_graph_encode_decode_layer_phase( DS4_N_EXPERT, tp_experts, DS4_SWIGLU_CLAMP_EXP, - peer_ffn_norm, NULL, 0, false) != 0; + peer_ffn_norm, NULL, 0, false, + NULL, NULL, NULL, NULL, NULL) != 0; } } #if !defined(__APPLE__) @@ -25075,7 +25106,8 @@ static bool metal_graph_encode_decode_layer_phase( DS4_N_EXPERT, tp_experts, DS4_SWIGLU_CLAMP_EXP, - metal_graph_ffn_norm(g), NULL, 0, false) != 0; + metal_graph_ffn_norm(g), NULL, 0, false, + NULL, NULL, NULL, NULL, NULL) != 0; } } #if !defined(__APPLE__) @@ -25159,7 +25191,7 @@ static bool metal_graph_encode_decode_layer_phase( const bool overlap_selected_shared = ok && g->tp_world < 2 && - !decode_stage_profile && + !decode_stage_serializing && !metal_graph_decode_cpu_router_applicable(g, layer) && layer->ffn_gate_tid2eid == NULL && getenv("DS4_MOE_REPLAY_SELECTED_IDS") == NULL && @@ -25178,7 +25210,7 @@ static bool metal_graph_encode_decode_layer_phase( ok && g->tp_world < 2 && !overlap_selected_shared && - !decode_stage_profile && + !decode_stage_serializing && metal_graph_use_iq2_selected_readahead_shared_delay(g) && metal_graph_decode_iq2_selected_slots_expected(g, layer) && !metal_graph_decode_cpu_router_applicable(g, layer) && @@ -25261,7 +25293,8 @@ static bool metal_graph_encode_decode_layer_phase( DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, metal_graph_ffn_norm(g), NULL, il, - false) != 0; + false, + NULL, NULL, NULL, NULL, NULL) != 0; DS4_METAL_PROFILE_DECODE_STAGE("routed_moe"); if (ok) { metal_graph_debug_dump_tensor("ffn_moe_gate_clamped", metal_graph_routed_gate(g), @@ -25467,7 +25500,8 @@ static bool metal_graph_encode_decode_layer_phase( DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, metal_graph_ffn_norm(g), NULL, il, - false) != 0; + false, + NULL, NULL, NULL, NULL, NULL) != 0; DS4_METAL_PROFILE_DECODE_STAGE("routed_moe"); if (ok) { metal_graph_debug_dump_tensor("ffn_moe_gate_clamped", metal_graph_routed_gate(g), @@ -25563,6 +25597,50 @@ static bool metal_graph_encode_decode_layer_phase( parallel_full_ffn = false; #endif } + /* HC epilogue tail fusion: when the fused router+shared gate/up kernel + * already produced shared_mid ahead of the routed MoE + * (router_shared_done != 0), encode the plain shared-down Q8_0 matvec + * into shared_out first and let the routed down-sum6 kernel add it and + * expand the four HC streams in the same dispatch. Bit-exact with the + * unfused dispatch pair (the fused kernel mirrors + * kernel_dsv4_shared_down_hc_expand4_q8_0's epilogue). Gated like the + * other fusion-forcing escape hatches (keep_ffn_out, steering, TP/CUDA + * splits) plus the DS4_METAL_DISABLE_DECODE_MOE_HC_FUSION rollback. If + * the wrapper declines the fused kernel (diagnostic pipeline toggles), + * moe_down_hc_fused stays 0 and the legacy shared-down+HC dispatch below + * reruns on the materialized routed_out/shared_out. */ + const bool fuse_moe_down_hc = + fuse_shared_down_hc && + router_shared_done != 0 && + !parallel_full_ffn && + !cuda_tp_moe && + !metal_graph_directional_steering_ffn_enabled(g) && + ds4_gpu_device_is_pre_m5_apple_silicon() && + layer->ffn_gate_exps->type == DS4_TENSOR_MXFP4 && + layer->ffn_up_exps->type == DS4_TENSOR_MXFP4 && + layer->ffn_down_exps->type == DS4_TENSOR_MXFP4 && + expert_in_dim == 4096u && + down_in_dim == 2048u && + gate_row_bytes == 2176u && + down_row_bytes == 1088u && + down_expert_bytes == 4456448u && + routed_out_dim == 4096u && + DS4_N_EXPERT == 256u && + DS4_N_EXPERT_USED == 6u && + DS4_N_HC == 4u && + !g->quality && + !g->ssd_streaming && + getenv("DS4_METAL_DISABLE_DECODE_MOE_HC_FUSION") == NULL; + if (ok && fuse_moe_down_hc) { + ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_out(g), + model, + layer->ffn_down_shexp, + shared_dim, + DS4_N_EMBD, + metal_graph_shared_mid(g), + 1); + } + int moe_down_hc_fused = 0; if (ok && !tp_fold_ffn && !cuda_tp_moe) ok = ds4_gpu_routed_moe_one_tensor(metal_graph_routed_out(g), metal_graph_routed_gate(g), metal_graph_routed_up(g), @@ -25584,7 +25662,32 @@ static bool metal_graph_encode_decode_layer_phase( DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, metal_graph_ffn_norm(g), NULL, il, - false) != 0; + false, + fuse_moe_down_hc ? metal_graph_shared_out(g) : NULL, + fuse_moe_down_hc ? metal_graph_after_attn_hc(g) : NULL, + fuse_moe_down_hc ? metal_graph_hc_split(g) : NULL, + fuse_moe_down_hc ? metal_graph_after_ffn_hc(g) : NULL, + &moe_down_hc_fused) != 0; + /* The wrapper declined the fused down+HC kernel (diagnostic pipeline + * toggles): rerun the legacy shared-down+HC dispatch. Its Q8 reduction + * is bit-identical to the plain matvec already in shared_out and its HC + * epilogue matches the fused kernel's, so the result is unchanged. */ + if (ok && fuse_moe_down_hc && !moe_down_hc_fused) { + ok = ds4_gpu_shared_down_hc_expand_q8_0_tensor( + metal_graph_after_ffn_hc(g), + metal_graph_shared_out(g), + model->map, + model->size, + layer->ffn_down_shexp->abs_offset, + shared_dim, + DS4_N_EMBD, + metal_graph_shared_mid(g), + metal_graph_routed_out(g), + metal_graph_after_attn_hc(g), + metal_graph_hc_split(g), + DS4_N_EMBD, + DS4_N_HC) != 0; + } if (!ok && parallel_full_ffn) { #if defined(__APPLE__) ds4_gpu_parallel_ffn_abort(); @@ -25792,7 +25895,7 @@ static bool metal_graph_encode_decode_layer_phase( DS4_N_EMBD, DS4_N_HC) != 0; } - } else if (ok && fuse_shared_down_hc) { + } else if (ok && fuse_shared_down_hc && !fuse_moe_down_hc) { if (cuda_tp_moe_peer_tmp) { ok = ds4_gpu_shared_down_hc_expand_add_q8_0_tensor( metal_graph_after_ffn_hc(g), @@ -25835,7 +25938,7 @@ static bool metal_graph_encode_decode_layer_phase( DS4_N_EMBD, metal_graph_shared_mid(g), 0); - } else if (ok) { + } else if (ok && !fuse_moe_down_hc) { ok = metal_graph_matmul_dense_quant_tensor(metal_graph_shared_out(g), model, layer->ffn_down_shexp, @@ -25871,7 +25974,8 @@ static bool metal_graph_encode_decode_layer_phase( DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, metal_graph_ffn_norm(g), metal_graph_shared_out(g), il, - false) != 0; + false, + NULL, NULL, NULL, NULL, NULL) != 0; DS4_METAL_PROFILE_DECODE_STAGE("routed_moe_folded"); } ds4_gpu_tensor *tp_ffn_a = NULL; /* rank0/rank1 partials consumed */ @@ -28358,7 +28462,8 @@ static bool metal_graph_encode_token_raw_swa( int token, uint32_t pos, bool need_logits, - bool allow_split_flush) { + bool allow_split_flush, + const ds4_gpu_tensor *token_dev) { if (g->raw_cap == 0) { fprintf(stderr, "ds4: Metal graph raw KV cache is not allocated\n"); return false; @@ -28389,14 +28494,31 @@ static bool metal_graph_encode_token_raw_swa( (void)ds4_gpu_set_decode_pipeline_fast_lookup(1); } #endif - bool ok = ds4_gpu_embed_token_hc_tensor(metal_graph_cur_hc(g), - model->map, - model->size, - weights->token_embd->abs_offset, - (uint32_t)weights->token_embd->dim[1], - (uint32_t)token, - DS4_N_EMBD, - DS4_N_HC) != 0; + bool ok; + if (token_dev) { + /* Greedy chain decode: the token id arrives GPU-resident (written by + * the previous token's argmax), so the host never blocks on it. The + * batched embed path gathers the identical row through the same + * get_rows/repeat kernels, keeping the transcript bit-identical. */ + ok = ds4_gpu_embed_tokens_hc_tensor(metal_graph_cur_hc(g), + token_dev, + model->map, + model->size, + weights->token_embd->abs_offset, + (uint32_t)weights->token_embd->dim[1], + 1, + DS4_N_EMBD, + DS4_N_HC) != 0; + } else { + ok = ds4_gpu_embed_token_hc_tensor(metal_graph_cur_hc(g), + model->map, + model->size, + weights->token_embd->abs_offset, + (uint32_t)weights->token_embd->dim[1], + (uint32_t)token, + DS4_N_EMBD, + DS4_N_HC) != 0; + } /* * Start executing the prefix of the decode graph while the CPU is still @@ -28841,6 +28963,9 @@ static bool metal_graph_indexer_stage_profile_boundary( uint32_t n_tokens, uint32_t n_comp, double *stage_t0) { + if (ds4_gpu_stage_counters_enabled()) { + return ds4_gpu_stage_counter_sample(stage) != 0; + } if (ds4_gpu_end_commands() == 0) return false; const double now = now_sec(); if (stage != NULL) { @@ -28967,6 +29092,9 @@ static bool metal_graph_layer_stage_profile_boundary( uint32_t pos0, uint32_t n_tokens, double *stage_t0) { + if (ds4_gpu_stage_counters_enabled()) { + return ds4_gpu_stage_counter_sample(stage) != 0; + } if (ds4_gpu_end_commands() == 0) return false; const double now = now_sec(); if (stage != NULL) { @@ -28989,6 +29117,9 @@ static bool metal_graph_q_stage_profile_boundary( uint32_t pos0, uint32_t n_tokens, double *stage_t0) { + if (ds4_gpu_stage_counters_enabled()) { + return ds4_gpu_stage_counter_sample(stage) != 0; + } if (ds4_gpu_end_commands() == 0) return false; const double now = now_sec(); fprintf(stderr, @@ -29331,6 +29462,7 @@ static bool metal_graph_encode_layer_attention_batch( ok = false; } bool q_b_f16_out = false; + bool batch_qkv_finalized = false; if (ok && !q_path_debug && layer->attn_q_b->type == DS4_TENSOR_Q8_0) { q_b_f16_out = ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor(tp_q ? tp_q : metal_graph_batch_q(g), tp_q_half ? tp_q_half : g->batch_q_half, @@ -29381,8 +29513,59 @@ static bool metal_graph_encode_layer_attention_batch( DS4_METAL_PROFILE_Q_STAGE("q_b"); const bool q_norm_debug = metal_graph_debug_wants("Qnorm", il, pos0); - bool q_norm_rope_fused = false; - if (ok && !q_norm_debug) { +#if defined(__APPLE__) + const bool kv_rope_debug = + metal_graph_debug_wants("KVrope", il, pos0); + const uint32_t prior_raw = + pos0 < g->raw_window ? pos0 : g->raw_window; + const bool batch_qkv_finalize_eligible = + qkv_rms_fused && + !q_path_debug && + !kv_rope_debug && + !q_stage_profile && + !layer_stage_profile && + !tp_row_split_attn && + !g->ssd_streaming && + !g->quality && + g->tp_world < 2 && + n_tokens >= 128u && + n_tokens <= 4096u && + DS4_N_HEAD == 64u && + DS4_N_HEAD_KV == 1u && + DS4_N_HEAD_DIM == 512u && + DS4_N_ROT == 64u && + g->layer_raw_cache[il] != NULL && + (uint64_t)n_tokens + prior_raw <= g->raw_cap && + !metal_graph_use_reference_kv_decode() && + getenv("DS4_METAL_DISABLE_PRE_M5_BATCH_QKV_FINALIZE") == NULL && + ds4_gpu_device_is_pre_m5_apple_silicon() && + ds4_gpu_dsv4_batch_qnorm_rope_kv_finalize_available() != 0; + if (ok && batch_qkv_finalize_eligible) { + batch_qkv_finalized = + ds4_gpu_dsv4_batch_qnorm_rope_kv_finalize_tensor( + metal_graph_batch_q(g), + metal_graph_batch_kv(g), + g->layer_raw_cache[il], + g->raw_cap, + n_tokens, + DS4_N_HEAD, + DS4_N_HEAD_DIM, + DS4_N_ROT, + pos0, + compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, + false, + freq_base, + freq_scale, + ext_factor, + attn_factor, + DS4_ROPE_YARN_BETA_FAST, + DS4_ROPE_YARN_BETA_SLOW, + DS4_RMS_EPS) != 0; + if (!batch_qkv_finalized) ok = false; + } +#endif + bool q_norm_rope_fused = batch_qkv_finalized; + if (ok && !q_norm_debug && !q_norm_rope_fused) { q_norm_rope_fused = ds4_gpu_head_rms_norm_rope_tail_tensor( tp_q ? tp_q : metal_graph_batch_q(g), tp_rows, @@ -29465,7 +29648,7 @@ static bool metal_graph_encode_layer_attention_batch( (uint64_t)n_tokens * DS4_N_HEAD_DIM, il, pos0); } } - if (ok) ok = ds4_gpu_rope_tail_tensor(metal_graph_batch_kv(g), + if (ok && !batch_qkv_finalized) ok = ds4_gpu_rope_tail_tensor(metal_graph_batch_kv(g), n_tokens, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, @@ -29483,7 +29666,7 @@ static bool metal_graph_encode_layer_attention_batch( metal_graph_debug_dump_tensor("KVrope", metal_graph_batch_kv(g), (uint64_t)n_tokens * DS4_N_HEAD_DIM, il, pos0); } - if (ok) ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(metal_graph_batch_kv(g), + if (ok && !batch_qkv_finalized) ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(metal_graph_batch_kv(g), n_tokens, DS4_N_HEAD_DIM, DS4_N_ROT) != 0; @@ -29500,7 +29683,7 @@ static bool metal_graph_encode_layer_attention_batch( * sized to hold the current chunk plus the previous SWA window, while the * attention mask still enforces the 128-token logical window. */ - if (ok && zero_prefix) ok = ds4_gpu_store_raw_kv_batch_tensor(g->layer_raw_cache[il], + if (ok && zero_prefix && !batch_qkv_finalized) ok = ds4_gpu_store_raw_kv_batch_tensor(g->layer_raw_cache[il], metal_graph_batch_kv(g), g->raw_cap, pos0, @@ -29563,12 +29746,14 @@ static bool metal_graph_encode_layer_attention_batch( const uint32_t raw_start = metal_graph_raw_start_for_span(g, pos0 + n_tokens - 1u, n_raw); - ok = ds4_gpu_store_raw_kv_batch_tensor(g->layer_raw_cache[il], - metal_graph_batch_kv(g), - g->raw_cap, - pos0, - n_tokens, - DS4_N_HEAD_DIM) != 0; + if (!batch_qkv_finalized) { + ok = ds4_gpu_store_raw_kv_batch_tensor(g->layer_raw_cache[il], + metal_graph_batch_kv(g), + g->raw_cap, + pos0, + n_tokens, + DS4_N_HEAD_DIM) != 0; + } if (ok) { metal_graph_debug_dump_tensor("raw_cache", g->layer_raw_cache[il], @@ -29922,6 +30107,26 @@ static bool metal_graph_encode_layer_attention_batch( } DS4_METAL_PROFILE_ATTN_STAGE("compressor"); + const bool topk_prefill_needed = + ratio == 4 && n_comp > DS4_N_INDEXER_TOP_K; +#if defined(__APPLE__) + /* A zero-prefix ratio-4 batch does not consume the indexer query or + * its per-head weights until the compressed cache grows beyond + * top-k. Keep building the indexer compressor/cache below, but skip + * these four ephemeral query dispatches while every attention row + * still uses the exact static-mixed path. */ + const bool prune_unused_indexer_query = + ratio == 4 && zero_prefix && n_tokens >= 32u && + !topk_prefill_needed && + !g->quality && !g->ssd_streaming && !g->ssd_streaming_cold && + g->placement == NULL && g->tp_world < 2u && + metal_graph_ported_m5_decode_feature_enabled( + "DS4_METAL_DISABLE_PRE_M5_BATCH_INDEXER_QUERY_PRUNE", + "DS4_METAL_DISABLE_M5_BATCH_INDEXER_QUERY_PRUNE"); +#else + const bool prune_unused_indexer_query = false; +#endif + if (ok && ratio == 4) { const uint32_t index_width = coff * DS4_N_INDEXER_HEAD_DIM; if (!layer->indexer_compressor_kv || !layer->indexer_compressor_gate || @@ -29958,14 +30163,14 @@ static bool metal_graph_encode_layer_attention_batch( (uint64_t)index_width * n_tokens, il, pos0); - if (ok) ok = metal_graph_matmul_plain_tensor(metal_graph_batch_indexer_q(g), + if (ok && !prune_unused_indexer_query) ok = metal_graph_matmul_plain_tensor(metal_graph_batch_indexer_q(g), model, layer->indexer_attn_q_b, q_rank, (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM, metal_graph_batch_qr_norm(g), n_tokens); - if (ok) ok = ds4_gpu_rope_tail_tensor(metal_graph_batch_indexer_q(g), + if (ok && !prune_unused_indexer_query) ok = ds4_gpu_rope_tail_tensor(metal_graph_batch_indexer_q(g), n_tokens, DS4_N_INDEXER_HEAD, DS4_N_INDEXER_HEAD_DIM, @@ -29979,10 +30184,10 @@ static bool metal_graph_encode_layer_attention_batch( attn_factor, DS4_ROPE_YARN_BETA_FAST, DS4_ROPE_YARN_BETA_SLOW) != 0; - if (ok) ok = ds4_gpu_dsv4_indexer_qat_tensor(metal_graph_batch_indexer_q(g), + if (ok && !prune_unused_indexer_query) ok = ds4_gpu_dsv4_indexer_qat_tensor(metal_graph_batch_indexer_q(g), n_tokens * DS4_N_INDEXER_HEAD, DS4_N_INDEXER_HEAD_DIM) != 0; - if (ok) ok = ds4_gpu_matmul_f16_tensor(metal_graph_batch_indexer_weights(g), + if (ok && !prune_unused_indexer_query) ok = ds4_gpu_matmul_f16_tensor(metal_graph_batch_indexer_weights(g), model->map, model->size, layer->indexer_proj->abs_offset, @@ -30297,12 +30502,14 @@ static bool metal_graph_encode_layer_attention_batch( bool use_indexed_comp = false; double index_stage_t0 = 0.0; - ok = ds4_gpu_store_raw_kv_batch_tensor(g->layer_raw_cache[il], - metal_graph_batch_kv(g), - g->raw_cap, - pos0, - n_tokens, - DS4_N_HEAD_DIM) != 0; + if (!batch_qkv_finalized) { + ok = ds4_gpu_store_raw_kv_batch_tensor(g->layer_raw_cache[il], + metal_graph_batch_kv(g), + g->raw_cap, + pos0, + n_tokens, + DS4_N_HEAD_DIM) != 0; + } if (ok && ratio == 4 && n_comp > DS4_N_INDEXER_TOP_K) { const float index_scale = 1.0f / sqrtf((float)(DS4_N_INDEXER_HEAD_DIM * DS4_N_INDEXER_HEAD)); if (index_stage_profile) { @@ -30422,7 +30629,6 @@ static bool metal_graph_encode_layer_attention_batch( if (ok) batch_attention_done = true; } - const bool topk_prefill_needed = ratio == 4 && n_comp > DS4_N_INDEXER_TOP_K; if (ok && !batch_attention_done && zero_prefix && topk_prefill_needed && n_comp != 0) { const float index_scale = 1.0f / sqrtf((float)(DS4_N_INDEXER_HEAD_DIM * DS4_N_INDEXER_HEAD)); @@ -30777,6 +30983,42 @@ static bool metal_graph_encode_layer_attention_batch( metal_graph_batch_heads(g), n_tokens) != 0; } + bool attn_out_hc_fused = false; +#if defined(__APPLE__) + if (ok && !attn_out_f16 && !attn_out_debug && !tp_row_split_attn && + n_tokens >= 512u && n_tokens <= 4096u && + (n_tokens % 32u) == 0u && !g->quality && !g->placement && + g->tp_world < 2u && !g->ssd_streaming && !g->ssd_streaming_cold && + !layer_stage_profile && !q_stage_profile && + metal_graph_debug_get_config()->prefix == NULL && + !metal_graph_directional_steering_attn_enabled(g) && + layer->attn_output_a->type == DS4_TENSOR_Q8_0 && + layer->attn_output_b->type == DS4_TENSOR_Q8_0 && + group_dim == 4096u && rank == 1024u && n_groups == 8u && + DS4_N_EMBD == 4096u && DS4_N_HC == 4u) { + const int fused = ds4_gpu_attention_output_q8_batch_hc_tensor( + metal_graph_batch_attn_out(g), + after_attn_hc_view, + metal_graph_batch_cur_hc(g), + hc_split_view, + metal_graph_batch_attn_low(g), + metal_graph_batch_group_tmp(g), + metal_graph_batch_low_tmp(g), + model->map, + model->size, + layer->attn_output_a->abs_offset, + layer->attn_output_b->abs_offset, + group_dim, + rank, + n_groups, + DS4_N_EMBD, + metal_graph_batch_heads(g), + n_tokens, + DS4_N_HC); + if (fused < 0) ok = false; + attn_out_hc_fused = fused > 0; + } +#endif uint64_t tp_attn_gate_seq = 0; /* Opt-in sub-chunk gate pipelining (see metal_graph_tp_subgate_pipeline; * measured net-negative on the M5 Max pair, kept for slower wires). @@ -30791,7 +31033,7 @@ static bool metal_graph_encode_layer_attention_batch( const bool tp_attn_pipeline = tp_row_split_attn && (n_tokens % 256u) == 0u && metal_graph_tp_subgate_pipeline(); - if (!attn_out_f16) { + if (!attn_out_f16 && !attn_out_hc_fused) { if (ok && tp_attn_pipeline) { /* Sub-chunk pipelined swap: the output projection runs in two * sub-halves of this rank's rows and each sub-half's row swap @@ -30899,7 +31141,9 @@ static bool metal_graph_encode_layer_attention_batch( if (ok && !attn_out_f16 && metal_graph_directional_steering_attn_enabled(g)) { ok = metal_graph_apply_directional_steering_attn(g, metal_graph_batch_attn_out(g), il, n_tokens); } - if (ok && attn_out_f16) { + if (ok && attn_out_hc_fused) { + /* The Q8 output-B specialization already wrote after_attn_hc_view. */ + } else if (ok && attn_out_f16) { ok = ds4_gpu_hc_expand_split_half_tensor(after_attn_hc_view, g->batch_q_half, metal_graph_batch_cur_hc(g), @@ -31209,10 +31453,40 @@ static bool metal_graph_encode_layer_ffn_batch( } const bool keep_ffn_out = metal_graph_needs_ffn_out(g, il, pos0); +#if defined(__APPLE__) + const bool fuse_moe_sum6_hc = + ds4_gpu_device_is_pre_m5_apple_silicon() && + ds4_gpu_moe_sum6_hc_expand_available() != 0 && + n_tokens >= 32u && n_tokens <= 4096u && decode_count == 0 && + !g->quality && !g->placement && g->tp_world < 2 && + !g->ssd_streaming && !g->ssd_streaming_cold && + !layer_stage_profile && !keep_ffn_out && + !metal_graph_directional_steering_ffn_enabled(g) && + metal_graph_debug_get_config()->prefix == NULL && + getenv("DS4_METAL_MOE_STAGE_PROFILE") == NULL && + getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL && + getenv("DS4_METAL_DISABLE_MOE_MM_ID_PAIR_SWIGLU") == NULL && + getenv("DS4_METAL_DISABLE_PRE_M5_BATCH_MOE_SUM6_HC_FUSION") == NULL && + layer->ffn_gate_exps->type == DS4_TENSOR_MXFP4 && + layer->ffn_up_exps->type == DS4_TENSOR_MXFP4 && + layer->ffn_down_exps->type == DS4_TENSOR_MXFP4 && + layer->ffn_gate_shexp->type == DS4_TENSOR_Q8_0 && + layer->ffn_up_shexp->type == DS4_TENSOR_Q8_0 && + layer->ffn_down_shexp->type == DS4_TENSOR_Q8_0 && + DS4_N_HC == 4u && DS4_N_EXPERT == 256u && + DS4_N_EXPERT_USED == 6u && DS4_N_EMBD == 4096u && + shared_dim == 2048u && expert_in_dim == 4096u && + expert_mid_dim == 2048u && down_in_dim == 2048u && + routed_out_dim == 4096u && gate_row_bytes == 2176u && + gate_expert_bytes == 4456448u && down_row_bytes == 1088u && + down_expert_bytes == 4456448u; +#else + const bool fuse_moe_sum6_hc = false; +#endif bool shared_down_f16 = false; #define DS4_METAL_TRY_SHARED_DOWN_F16() do { \ - if (ok && !tp_row_split_ffn && !keep_ffn_out && \ + if (ok && !fuse_moe_sum6_hc && !tp_row_split_ffn && !keep_ffn_out && \ !metal_graph_debug_wants("ffn_shexp", il, pos0)) { \ shared_down_f16 = ds4_gpu_matmul_q8_0_f16_out_tensor(g->batch_q_half, \ model->map, \ @@ -31421,7 +31695,8 @@ static bool metal_graph_encode_layer_ffn_batch( x_row, NULL, il, - false) != 0; + false, + NULL, NULL, NULL, NULL, NULL) != 0; ds4_gpu_tensor_free(w_row); ds4_gpu_tensor_free(sel_row); ds4_gpu_tensor_free(x_row); @@ -31463,6 +31738,7 @@ static bool metal_graph_encode_layer_ffn_batch( il, n_tokens, &g->batch_routed_mid_is_f16, + fuse_moe_sum6_hc, false) != 0; } if (ok) { @@ -31555,6 +31831,16 @@ static bool metal_graph_encode_layer_ffn_batch( DS4_N_EMBD, DS4_N_HC) != 0; } + else if (ok && fuse_moe_sum6_hc) { + ok = ds4_gpu_moe_sum6_hc_expand_split_tensor( + next_hc_view, + metal_graph_batch_routed_down(g), + metal_graph_batch_shared_out(g), + metal_graph_batch_after_attn_hc(g), + hc_split_view, + DS4_N_EMBD, + DS4_N_HC) != 0; + } else if (ok && shared_down_f16) { ok = ds4_gpu_hc_expand_add_split_half_add_tensor(next_hc_view, metal_graph_batch_routed_out(g), @@ -31831,12 +32117,15 @@ static bool metal_graph_eval_token_raw_swa( "DS4_METAL_GRAPH_TOKEN_PROFILE"); const bool throttle = graph_power_throttle_enabled(g); const double t0 = (profile || throttle) ? now_sec() : 0.0; + const bool stage_counters = ds4_gpu_stage_counters_enabled(); + if (stage_counters) ds4_gpu_stage_counter_reset(); bool ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = metal_graph_encode_token_raw_swa(g, model, weights, token, pos, logits != NULL, true); + if (ok) ok = metal_graph_encode_token_raw_swa(g, model, weights, token, pos, logits != NULL, true, NULL); const double t_encoded = (profile || throttle) ? now_sec() : 0.0; if (ok) ok = ds4_gpu_end_commands() != 0; const double t_done = (profile || throttle) ? now_sec() : 0.0; + if (ok && stage_counters) ds4_gpu_stage_counter_report(pos); if (ok && logits && g->tp_world == 2 && g->tp_logits_half) { const uint64_t tp_vhalf = (uint64_t)DS4_N_VOCAB / 2u; @@ -31866,6 +32155,171 @@ static bool metal_graph_eval_token_raw_swa( return ok; } +/* ------------------------------------------------------------------------ + * Greedy chain decode. + * + * The classic generate loop serializes every token through the host: + * waitUntilCompleted, a 517 KiB logits readback, a CPU argmax, and only then + * the next token's encode; the GPU idles in that window (~0.5 ms/token on + * M3 Ultra). The chain removes the host from the per-token critical path: + * each token's graph ends with the GPU argmax writing the next token id into + * a device ring, and the next token's embedding reads the id from the ring. + * Encoding runs DS4_GREEDY_CHAIN_AHEAD tokens ahead of the host's confirm + * cursor, so a token's command buffers are always committed before the GPU + * drains the previous one; the host only lags to fetch confirmed ids (one + * shared-event wait per token) for printing and stop checks. + * + * Bit-exactness: every kernel and its inputs are unchanged; the id travels + * through a device buffer instead of an encode-time constant, and the GPU + * argmax (already the split-kv verifier's greedy oracle) reproduces the CPU + * argmax including lowest-index tie-breaking, so the transcript is identical. + * --------------------------------------------------------------------- */ +#define DS4_GREEDY_CHAIN_AHEAD 2 + +/* Chains n_evals evals from pos0 with ring[0] = seed_token, calling on_token + * for the seed and then for every confirmed id in order; a false return stops + * the chain early (up to AHEAD-1 already-committed evals are discarded). + * Returns the number of approved tokens (>= 0, seed included) or -1 on error. + * When logits_out is non-NULL the final token's logits are read back after + * the chain drains. */ +static int metal_graph_greedy_chain( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + int seed_token, + uint32_t pos0, + int n_evals, + bool (*on_token)(void *ctx, int token), + void *on_token_ctx, + float *logits_out) { + ds4_gpu_tensor *ring = + ds4_gpu_tensor_alloc(((uint64_t)n_evals + 2u) * sizeof(int32_t)); + if (!ring) { + fprintf(stderr, "ds4: greedy chain ring allocation failed\n"); + return -1; + } + const int32_t seed32 = (int32_t)seed_token; + if (ds4_gpu_tensor_write(ring, 0, &seed32, sizeof(seed32)) == 0) { + ds4_gpu_tensor_free(ring); + fprintf(stderr, "ds4: greedy chain ring seed failed\n"); + return -1; + } + + int approved = 0; + if (getenv("DS4_GREEDY_CHAIN_DUMP_IDS")) { + fprintf(stderr, "ds4: chain id[0]=%d\n", seed_token); + } + if (!on_token(on_token_ctx, seed_token)) { + ds4_gpu_tensor_free(ring); + return 0; + } + approved = 1; + if (n_evals < 1) { + ds4_gpu_tensor_free(ring); + return approved; + } + if (getenv("DS4_GREEDY_CHAIN_DEBUG")) { + fprintf(stderr, "ds4: greedy chain engaged: pos0=%u n_evals=%d\n", + pos0, n_evals); + } + + bool ok = ds4_gpu_begin_commands() != 0; + int encoded = 0; /* evals committed: eval(j) for j in [0, encoded) */ + bool stop = false; + uint64_t ev_ring[DS4_GREEDY_CHAIN_AHEAD] = {0}; + for (int j = 0; ok && j < n_evals && !stop; j++) { + ds4_gpu_tensor *vin = ds4_gpu_tensor_view(ring, + (uint64_t)j * sizeof(int32_t), + sizeof(int32_t)); + g->chain_token_view = vin; + ok = vin != NULL && + metal_graph_encode_token_raw_swa(g, model, weights, + 0, pos0 + (uint32_t)j, + true, true, vin); + g->chain_token_view = NULL; + ds4_gpu_tensor_free(vin); + if (!ok) break; + ds4_gpu_tensor *vout = ds4_gpu_tensor_view(ring, + (uint64_t)(j + 1) * sizeof(int32_t), + sizeof(int32_t)); + ok = vout != NULL && + ds4_gpu_argmax_tensor(vout, metal_graph_logits(g), DS4_N_VOCAB) != 0; + ds4_gpu_tensor_free(vout); + if (!ok) break; + uint64_t ev = 0; + ok = ds4_gpu_signal_selected_readback_ready(&ev) != 0; + if (!ok) break; + ok = ds4_gpu_flush_commands() != 0; + if (!ok) break; + encoded = j + 1; + ev_ring[j % DS4_GREEDY_CHAIN_AHEAD] = ev; + /* Confirm the id that falls out of the encode-ahead window. */ + const int k = j - (DS4_GREEDY_CHAIN_AHEAD - 1); + if (k >= 0) { + const uint64_t ev_k = ev_ring[k % DS4_GREEDY_CHAIN_AHEAD]; + int32_t id = -1; + ok = ds4_gpu_wait_selected_readback_ready(ev_k, "greedy chain id") != 0 && + ds4_gpu_tensor_read(ring, + (uint64_t)(k + 1) * sizeof(int32_t), + &id, + sizeof(id)) != 0; + if (!ok) break; + if (getenv("DS4_GREEDY_CHAIN_DUMP_IDS")) { + fprintf(stderr, "ds4: chain id[%d]=%d\n", k + 1, (int)id); + } + if (getenv("DS4_GREEDY_CHAIN_VERIFY")) { + float *lbuf = xmalloc((size_t)DS4_N_VOCAB * sizeof(float)); + if (ds4_gpu_tensor_read(metal_graph_logits(g), 0, lbuf, + (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0) { + const int cpu_id = sample_argmax(lbuf, DS4_N_VOCAB); + if (cpu_id != (int)id) { + fprintf(stderr, "ds4: chain verify mismatch at id[%d]: gpu=%d cpu=%d\n", + k + 1, (int)id, cpu_id); + } + } + free(lbuf); + } + if (!on_token(on_token_ctx, (int)id)) { + stop = true; + } else { + approved++; + } + } + } + /* Drain the encode-ahead window: confirm the remaining ids. */ + for (int k = encoded - (DS4_GREEDY_CHAIN_AHEAD - 1); ok && !stop && k < encoded; k++) { + if (k < 0) continue; + const uint64_t ev_k = ev_ring[k % DS4_GREEDY_CHAIN_AHEAD]; + int32_t id = -1; + ok = ds4_gpu_wait_selected_readback_ready(ev_k, "greedy chain id") != 0 && + ds4_gpu_tensor_read(ring, + (uint64_t)(k + 1) * sizeof(int32_t), + &id, + sizeof(id)) != 0; + if (!ok) break; + if (getenv("DS4_GREEDY_CHAIN_DUMP_IDS")) { + fprintf(stderr, "ds4: chain id[%d]=%d\n", k + 1, (int)id); + } + if (!on_token(on_token_ctx, (int)id)) stop = true; + else approved++; + } + if (ds4_gpu_end_commands() == 0) ok = false; + if (ok && logits_out) { + ok = ds4_gpu_tensor_read(metal_graph_logits(g), + 0, + logits_out, + (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + } + ds4_gpu_tensor_free(ring); + if (!ok) { + if (ds4_gpu_synchronize() == 0) { + fprintf(stderr, "ds4: Metal synchronize after greedy chain failure also failed\n"); + } + return -1; + } + return approved; +} + static bool metal_graph_streaming_decode_prefill_wide_default( const ds4_weights *weights) { if (DS4_MODEL_VARIANT != DS4_VARIANT_FLASH || !weights || DS4_N_LAYER == 0) { @@ -32321,7 +32775,7 @@ static bool metal_graph_eval_token_raw_swa_top( uint32_t output_ways = 0; bool ok = ds4_gpu_begin_commands() != 0; if (ok) ok = metal_graph_encode_token_raw_swa(g, model, weights, - token, pos, false, true); + token, pos, false, true, NULL); if (ok) ok = metal_graph_encode_output_head_split_top1(g, model, weights, @@ -32382,7 +32836,7 @@ static bool metal_graph_eval_token_raw_swa_top( bool ok = ds4_gpu_begin_commands() != 0; if (ok) ok = metal_graph_encode_token_raw_swa(g, model, weights, - token, pos, true, true); + token, pos, true, true, NULL); if (ok) { ok = ds4_gpu_argmax_tensor(metal_graph_comp_selected(g), metal_graph_logits(g), @@ -44225,7 +44679,8 @@ static int glm_graph_routed_moe_one_dispatch( x, NULL, il, - force_resident); + force_resident, + NULL, NULL, NULL, NULL, NULL); } return ds4_gpu_glm_routed_moe_one_tensor(out, @@ -44363,6 +44818,7 @@ static int glm_graph_routed_moe_batch_dispatch( il, n_tokens, &g->batch_routed_mid_is_f16, + false, force_resident); } @@ -52957,6 +53413,19 @@ static int generate_glm_metal_argmax( /* Metal generation entry point. The model runs as one local whole-graph * pipeline: graph prefill followed by graph decode steps. Streaming PRO may * use decode-style prefill for short prompts. */ +typedef struct { + const ds4_vocab *vocab; + ds4_token_emit_fn emit; + void *emit_ud; +} metal_graph_chain_emit_ctx; + +static bool metal_graph_chain_on_token(void *vctx, int token) { + metal_graph_chain_emit_ctx *ctx = vctx; + if (vocab_token_is_generation_stop(ctx->vocab, token)) return false; + if (ctx->emit) ctx->emit(ctx->emit_ud, token); + return true; +} + static int generate_metal_graph_raw_swa( const ds4_model * model, const ds4_vocab * vocab, @@ -53091,7 +53560,51 @@ static int generate_metal_graph_raw_swa( int n_generated = 0; int n_decode_eval = 0; const double t_decode0 = now_sec(); - for (int i = 0; i < n_predict && pos < ctx_size; i++) { + /* Chained greedy decode keeps the token id on-device and encodes ahead, + * removing the per-token host sync from the GPU critical path. The + * classic loop below remains the fallback and the reference. */ + const int chain_evals = n_predict < ctx_size - pos ? n_predict - 1 + : ctx_size - pos - 1; + bool chain_cpu_router = false; + for (uint32_t il = 0; il < DS4_N_LAYER && !chain_cpu_router; il++) { + chain_cpu_router = + metal_graph_decode_cpu_router_applicable(&g, &weights->layer[il]); + } + const bool chain_ok = + !ssd_streaming && + !quality && + !chain_cpu_router && + chain_evals >= 1 && + getenv("DS4_DISABLE_GREEDY_CHAIN") == NULL && + !trace_top && + !token_timing && + getenv("DS4_METAL_GRAPH_TOKEN_PROFILE") == NULL && + getenv("DS4_METAL_DECODE_STAGE_PROFILE") == NULL && + !graph_power_throttle_enabled(&g); + if (chain_ok) { + const int seed = sample_argmax(logits, DS4_N_VOCAB); + metal_graph_chain_emit_ctx chain_ctx = { + .vocab = vocab, + .emit = emit, + .emit_ud = emit_ud, + }; + const int n = metal_graph_greedy_chain(&g, + model, + weights, + seed, + (uint32_t)pos, + chain_evals, + metal_graph_chain_on_token, + &chain_ctx, + logits); + if (n < 0) { + ok = false; + } else { + n_generated = n; + pos += n_generated; + } + } + for (int i = 0; !chain_ok && i < n_predict && pos < ctx_size; i++) { if (trace_top) { char label[64]; snprintf(label, sizeof(label), "step %d", i); @@ -53099,6 +53612,9 @@ static int generate_metal_graph_raw_swa( } int token = sample_argmax(logits, DS4_N_VOCAB); + if (getenv("DS4_GREEDY_CHAIN_DUMP_IDS")) { + fprintf(stderr, "ds4: classic id[%d]=%d\n", i, token); + } if (vocab_token_is_generation_stop(vocab, token)) break; if (emit) emit(emit_ud, token); @@ -64562,7 +65078,13 @@ int ds4_session_create(ds4_session **out, ds4_engine *e, int ctx_size) { free(s); return 1; } - if (s->graph.dspark_capture_enabled) { + /* Diagnostic, not news: the capture configuration is decided by the + * build and the weights, not by anything the caller did, and an + * embedder prints it in the middle of the user's screen on every + * session it creates. Opt in with DS4_DSPARK_VERBOSE, matching the + * other DS4_* diagnostics. */ + if (s->graph.dspark_capture_enabled && + getenv("DS4_DSPARK_VERBOSE") != NULL) { fprintf(stderr, "ds4: DSpark target-hidden capture enabled: layers="); for (uint32_t i = 0; i < s->graph.dspark_target_layer_count; i++) { fprintf(stderr, @@ -68830,6 +69352,122 @@ static bool glm53_graph_encode_native_session_batch( return ok; } +/* Session-level greedy chain decode: same device-resident token ring as the + * CLI path (metal_graph_greedy_chain), driven from an existing session. */ +bool ds4_session_chain_greedy_supported(const ds4_session *s) { + if (!s || !s->engine || s->engine->backend != DS4_BACKEND_METAL) return false; + if (!s->checkpoint_valid || s->checkpoint.len <= 0) return false; + if (ds4_session_is_cpu(s) || ds4_session_is_glm(s) || s->distributed) { + return false; + } + ds4_engine *e = s->engine; + /* The chain skips the per-token support-draft prep, so restrict it to + * sessions without an MTP/DSpark companion. */ + if (e->support_kind != DS4_SUPPORT_NONE) return false; + if (e->tp.active || s->graph.tp_world >= 2) return false; + if (s->graph.placement != NULL) return false; /* multi-tier: CLI never chains it */ + if (s->graph.ssd_streaming || s->graph.quality) return false; + if (metal_graph_directional_steering_attn_enabled(&s->graph) || + metal_graph_directional_steering_ffn_enabled(&s->graph)) return false; + for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + if (metal_graph_decode_cpu_router_applicable(&s->graph, + &e->weights.layer[il])) { + return false; + } + } + if (getenv("DS4_DISABLE_GREEDY_CHAIN") != NULL) return false; + if (getenv("DS4_METAL_GRAPH_TOKEN_PROFILE") != NULL) return false; + if (getenv("DS4_METAL_DECODE_STAGE_PROFILE") != NULL) return false; + if (getenv("DS4_METAL_GRAPH_DUMP_PREFIX") != NULL) return false; + if (graph_power_throttle_enabled(&s->graph)) return false; + return true; +} + +typedef struct { + ds4_session *s; + bool (*on_token)(void *ctx, int token); + void *on_token_ctx; + bool stopped_early; +} ds4_session_chain_tramp; + +static bool ds4_session_chain_tramp_on_token(void *vctx, int token) { + ds4_session_chain_tramp *t = (ds4_session_chain_tramp *)vctx; + /* Only tokens the caller approves enter the checkpoint; a false return + * stops the chain and leaves the already-encoded ahead-evals discarded + * (their KV rows are overwritten by any later eval at those positions). */ + if (!t->on_token(t->on_token_ctx, token)) { + t->stopped_early = true; + return false; + } + token_vec_push(&t->s->checkpoint, token); + return true; +} + +int ds4_session_eval_chain_greedy(ds4_session *s, int max_tokens, + bool (*on_token)(void *ctx, int token), + void *on_token_ctx, + bool *completed, + char *err, size_t errlen) { + if (!ds4_session_chain_greedy_supported(s)) { + if (errlen) snprintf(err, errlen, "greedy chain decode not supported for this session"); + return -1; + } + if (!on_token) { + if (errlen) snprintf(err, errlen, "greedy chain decode requires an on_token callback"); + return -1; + } + ds4_engine *e = s->engine; + int n = max_tokens; + const int ctx_left = s->ctx_size - s->checkpoint.len; + if (n > ctx_left) n = ctx_left; + if (n < 2) { + /* A one-token burst would never run a GPU eval, leaving s->logits + * stale; the caller uses the classic step for that. */ + if (errlen) snprintf(err, errlen, "greedy chain burst requires at least 2 tokens of headroom"); + return -1; + } + ds4_session_chain_tramp tramp = { s, on_token, on_token_ctx, false }; + const int seed = ds4_session_argmax(s); + const int approved = metal_graph_greedy_chain(&s->graph, + &e->model, + &e->weights, + seed, + (uint32_t)s->checkpoint.len, + n - 1, + ds4_session_chain_tramp_on_token, + &tramp, + s->logits); + if (approved < 0) { + if (errlen) snprintf(err, errlen, "%s greedy chain decode failed", + ds4_backend_name(e->backend)); + s->checkpoint_valid = false; + return -1; + } + s->checkpoint_valid = s->checkpoint.len > 0; + s->mtp_draft_valid = false; + s->dspark_draft_valid = false; + if (completed) *completed = !tramp.stopped_early; + return approved; +} +#else +bool ds4_session_chain_greedy_supported(const ds4_session *s) { + (void)s; + return false; +} + +int ds4_session_eval_chain_greedy(ds4_session *s, int max_tokens, + bool (*on_token)(void *ctx, int token), + void *on_token_ctx, + bool *completed, + char *err, size_t errlen) { + (void)s; (void)max_tokens; (void)on_token; (void)on_token_ctx; + (void)completed; + if (errlen) snprintf(err, errlen, "GPU support is not compiled in"); + return -1; +} +#endif + +#ifndef DS4_NO_GPU static bool ds4_sessions_eval_batch_metal_supported( ds4_decode_item *items, int count, @@ -69280,7 +69918,8 @@ static int ds4_sessions_eval_batch_metal( items[i].token, pos, true, - false); + false, + NULL); } } } @@ -73382,7 +74021,8 @@ static int ds4_sessions_eval_batch_cuda(ds4_decode_item *items, int count, items[i].token, (uint32_t)s->checkpoint.len, true, - false); + false, + NULL); } } if (ok) ok = ds4_gpu_end_commands() != 0; diff --git a/ds4.h b/ds4.h index e6dae1b9f..e8a09e651 100644 --- a/ds4.h +++ b/ds4.h @@ -488,6 +488,20 @@ int ds4_session_set_logits(ds4_session *s, const float *logits, int n); void ds4_session_gpu_warmup(ds4_session *s); int ds4_session_eval(ds4_session *s, int token, char *err, size_t errlen); +/* Greedy chain decode for sessions (Metal, greedy only): keeps the next token + * id on-device and encodes ahead, removing the per-token host sync. + * on_token receives the seed first, then every confirmed id in order; a false + * return stops the burst early without committing that token. Returns the + * number of approved tokens (>= 0) or -1 with err set. When *completed is + * false the burst stopped early and s->logits is stale; when true, s->logits + * holds the logits after the last approved token, as after ds4_session_eval. */ +bool ds4_session_chain_greedy_supported(const ds4_session *s); +int ds4_session_eval_chain_greedy(ds4_session *s, int max_tokens, + bool (*on_token)(void *ctx, int token), + void *on_token_ctx, + bool *completed, + char *err, size_t errlen); + typedef struct { ds4_session *session; int token; diff --git a/ds4_cuda.cu b/ds4_cuda.cu index fd36d7014..3d3746ba9 100644 --- a/ds4_cuda.cu +++ b/ds4_cuda.cu @@ -26000,7 +26000,20 @@ extern "C" int ds4_gpu_routed_moe_owned_packed_combine_tensor( extern "C" int ds4_gpu_routed_moe_one_tensor(ds4_gpu_tensor *out, ds4_gpu_tensor *gate, ds4_gpu_tensor *up, ds4_gpu_tensor *mid, ds4_gpu_tensor *down, const void *model_map, uint64_t model_size, uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset, uint32_t gate_type, uint32_t down_type, uint64_t gate_expert_bytes, uint64_t gate_row_bytes, uint64_t down_expert_bytes, uint64_t down_row_bytes, uint32_t expert_in_dim, uint32_t expert_mid_dim, uint32_t out_dim, const ds4_gpu_tensor *selected, const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, float clamp, const ds4_gpu_tensor *x, const ds4_gpu_tensor *add_in, uint32_t layer_index, - bool force_resident) { + bool force_resident, + const ds4_gpu_tensor *hc_shared_out, + const ds4_gpu_tensor *hc_residual, + const ds4_gpu_tensor *hc_split, + ds4_gpu_tensor *hc_out, + int *hc_fused_out) { + /* The decode MoE+HC tail fusion is Metal-only; CUDA callers always pass + * NULL hc tensors. */ + (void)hc_shared_out; (void)hc_residual; (void)hc_split; + if (hc_fused_out) *hc_fused_out = 0; + if (hc_out) { + fprintf(stderr, "ds4: routed MoE HC tail fusion is Metal-only\n"); + return 0; + } if (add_in) { if (!ds4_gpu_add_tensor(out, out, add_in, (uint32_t)(out->bytes / sizeof(float)))) return 0; @@ -26014,8 +26027,9 @@ extern "C" int ds4_gpu_routed_moe_one_tensor(ds4_gpu_tensor *out, ds4_gpu_tensor selected, weights, n_total_expert, n_expert, clamp, x, layer_index, 1, force_resident ? 0 : 1, 0); } -extern "C" int ds4_gpu_routed_moe_batch_tensor(ds4_gpu_tensor *out, ds4_gpu_tensor *gate, ds4_gpu_tensor *up, ds4_gpu_tensor *mid, ds4_gpu_tensor *down, const void *model_map, uint64_t model_size, uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset, uint32_t gate_type, uint32_t down_type, uint64_t gate_expert_bytes, uint64_t gate_row_bytes, uint64_t down_expert_bytes, uint64_t down_row_bytes, uint32_t expert_in_dim, uint32_t expert_mid_dim, uint32_t out_dim, const ds4_gpu_tensor *selected, const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, float clamp, const ds4_gpu_tensor *x, uint32_t layer_index, uint32_t n_tokens, bool *mid_is_f16, bool force_resident) { +extern "C" int ds4_gpu_routed_moe_batch_tensor(ds4_gpu_tensor *out, ds4_gpu_tensor *gate, ds4_gpu_tensor *up, ds4_gpu_tensor *mid, ds4_gpu_tensor *down, const void *model_map, uint64_t model_size, uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset, uint32_t gate_type, uint32_t down_type, uint64_t gate_expert_bytes, uint64_t gate_row_bytes, uint64_t down_expert_bytes, uint64_t down_row_bytes, uint32_t expert_in_dim, uint32_t expert_mid_dim, uint32_t out_dim, const ds4_gpu_tensor *selected, const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, float clamp, const ds4_gpu_tensor *x, uint32_t layer_index, uint32_t n_tokens, bool *mid_is_f16, bool defer_sum6, bool force_resident) { (void)force_resident; + if (defer_sum6) return 0; if (mid_is_f16) *mid_is_f16 = false; return routed_moe_launch(out, gate, up, mid, down, model_map, model_size, gate_offset, up_offset, down_offset, diff --git a/ds4_eval.c b/ds4_eval.c index bcb351453..bd7fda42b 100644 --- a/ds4_eval.c +++ b/ds4_eval.c @@ -3634,6 +3634,77 @@ static void eval_prefill_progress(void *ud, const char *event, int current, int if (paused_sec > 0.0) ui->phase_start_sec += paused_sec; } +/* Greedy-chain burst decode: carries the per-token bookkeeping of the classic + * decode loop so a chained burst emits exactly the same stream, stop checks, + * and think-close records. */ +typedef struct { + ds4_engine *engine; + eval_ui *ui; + int idx; + int generation_limit; + byte_buf *raw; + bool *generation_in_think; + bool *plain_in_think; + eval_think_close_info *think_close; + double *t0; + bool tty; + bool use_plain_color; + bool stop; + bool quit; + bool switch_case; +} eval_chain_ctx; + +static bool eval_chain_on_token(void *vctx, int token) { + eval_chain_ctx *c = (eval_chain_ctx *)vctx; + eval_ui *ui = c->ui; + if (c->tty) { + tui_consume_input(ui); + if (tui_has_quit_request(ui)) { c->quit = true; return false; } + if (tui_has_switch_request(ui, c->idx)) { c->switch_case = true; return false; } + double paused_sec = tui_wait_if_paused(ui, ui->in_think ? "thinking" : "answer"); + if (paused_sec > 0.0) { + ui->phase_start_sec += paused_sec; + *c->t0 += paused_sec; + } + if (tui_has_quit_request(ui)) { c->quit = true; return false; } + if (tui_has_switch_request(ui, c->idx)) { c->switch_case = true; return false; } + } + + if (ds4_token_is_stop(c->engine, token)) { c->stop = true; return false; } + size_t len = 0; + char *text = ds4_token_text(c->engine, token, &len); + buf_append(c->raw, text, len); + ui->generated++; + ui->generated_tokens[c->idx] = ui->generated; + tui_run_clock_tick(ui); + if (*c->generation_in_think && c->raw->v && strstr(c->raw->v, "")) { + *c->generation_in_think = false; + if (c->think_close->kind == EVAL_THINK_CLOSE_NONE) { + c->think_close->kind = EVAL_THINK_CLOSE_NATURAL; + c->think_close->token_index = ui->generated; + c->think_close->remaining_budget = + c->generation_limit - ui->generated + 1; + c->think_close->rank = 0; + } + } + double elapsed = now_sec() - ui->phase_start_sec; + ui->speed_tps = elapsed > 0.001 ? (double)ui->generated / elapsed : 0.0; + + if (c->tty) { + stream_append_token_text(ui, text, len, false); + tui_refresh(ui, ui->in_think ? "thinking" : "answer"); + } else { + if (*c->plain_in_think && strstr(c->raw->v ? c->raw->v : "", "")) { + *c->plain_in_think = false; + plain_reset_color(c->use_plain_color); + } + fwrite(text, 1, len, stdout); + fflush(stdout); + } + free(text); + return true; +} + static eval_run_result run_one_case(ds4_engine *engine, ds4_session *session, const eval_config *cfg, eval_ui *ui, FILE *trace, int idx, uint64_t *rng) { @@ -3774,6 +3845,12 @@ static eval_run_result run_one_case(ds4_engine *engine, ds4_session *session, double t0 = ui->phase_start_sec; int forced_close_pos = -1; + /* Greedy chain bursts keep the token id on-device between evals (same + * machinery as the CLI chain decode). Checked once per case; the burst + * branch below re-verifies per iteration that the think-close controller + * cannot intervene, so chained and classic decode pick identical tokens. */ + const bool chain_burst_ok = + cfg->temperature <= 0.0f && ds4_session_chain_greedy_supported(session); for (int i = 0; i < generation_limit; i++) { if (tty) { tui_consume_input(ui); @@ -3870,6 +3947,58 @@ static eval_run_result run_one_case(ds4_engine *engine, ds4_session *session, } } } + /* Chained greedy burst: while the think-close controller cannot + * intervene (outside its reply-budget window), decode a run of argmax + * tokens with the id kept on-device. The window is where the classic + * path may force a non-argmax token, so bursts stop ahead of it. */ + if (token < 0 && chain_burst_ok && remaining_budget >= 2) { + int burst = remaining_budget; + if (generation_in_think && think_close_tokens.len > 0) { + const int window = think_close_tokens.len == 1 + ? cfg->soft_limit_reply_budget + : cfg->hard_limit_reply_budget; + burst = remaining_budget - window; + } + if (burst >= 2) { + eval_chain_ctx cctx = { + .engine = engine, + .ui = ui, + .idx = idx, + .generation_limit = generation_limit, + .raw = &raw, + .generation_in_think = &generation_in_think, + .plain_in_think = &plain_in_think, + .think_close = &think_close, + .t0 = &t0, + .tty = tty, + .use_plain_color = use_plain_color, + }; + const int n = ds4_session_eval_chain_greedy(session, burst, + eval_chain_on_token, + &cctx, NULL, + err, sizeof(err)); + if (n < 0) { + plain_reset_color(use_plain_color); + ui->generated_tokens[idx] = ui->generated; + tui_run_clock_stop(ui); + fprintf(stderr, "ds4-eval: decode failed for %s: %s\n", + tc->id, err); + trace_write_case(trace, cfg, tc, idx, ui->ncases, "ERROR", + err, system, question, + raw.v ? raw.v : "", think_mode, + prompt_tokens, ui->generated, + now_sec() - t0, "?", &think_close); + free(question); + ds4_tokens_free(&think_close_tokens); + buf_free(&raw); + return EVAL_RUN_ERROR; + } + if (cctx.stop) break; + if (cctx.quit || cctx.switch_case) continue; /* handled above */ + i += n - 1; /* keep i in lockstep with ui->generated */ + continue; + } + } if (token < 0) token = ds4_session_sample(session, cfg->temperature, 0, cfg->top_p, cfg->min_p, rng); diff --git a/ds4_gpu.h b/ds4_gpu.h index 21d016019..6f2a311ea 100644 --- a/ds4_gpu.h +++ b/ds4_gpu.h @@ -110,6 +110,13 @@ int ds4_gpu_tensor_read_after_selected_event(const ds4_gpu_tensor *tensor, int ds4_gpu_end_commands(void); int ds4_gpu_synchronize(void); +/* Diagnostic GPU stage-counter profiler: timestamp samples taken at decode + * stage boundaries without ending the batch command buffer. */ +int ds4_gpu_stage_counters_enabled(void); +int ds4_gpu_stage_counter_sample(const char *label); +void ds4_gpu_stage_counter_reset(void); +void ds4_gpu_stage_counter_report(uint32_t pos); + int ds4_gpu_set_model_map(const void *model_map, uint64_t model_size); int ds4_gpu_set_model_fd(int fd); int ds4_gpu_set_model_fd_for_map(int fd, const void *model_map); @@ -186,7 +193,11 @@ enum { DS4_GPU_TEST_MXFP4_DOWN_HALF_LUT = 1u << 4, DS4_GPU_TEST_OUTPUT_HC_WEIGHTS4 = 1u << 5, DS4_GPU_TEST_HC_RMS_SCALE_PROJ = 1u << 6, + DS4_GPU_TEST_ATTN_OUT_LOW_Q8_STATIC = 1u << 7, + DS4_GPU_TEST_INDEXED_ATTN_PREFILL_RB4 = 1u << 8, + DS4_GPU_TEST_BATCH_ATTN_OUT_HC_FUSION = 1u << 9, }; +int ds4_gpu_test_get_quality(void); void ds4_gpu_test_set_flags(uint32_t flags); void ds4_gpu_release_zero_prefix_prefill_mask_cache(void); #else @@ -1119,6 +1130,28 @@ int ds4_gpu_head_rms_norm_rope_tail_tensor( float beta_slow, float eps); +int ds4_gpu_dsv4_batch_qnorm_rope_kv_finalize_available(void); + +int ds4_gpu_dsv4_batch_qnorm_rope_kv_finalize_tensor( + ds4_gpu_tensor *q, + ds4_gpu_tensor *kv, + ds4_gpu_tensor *raw_cache, + uint32_t raw_cap, + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + uint32_t pos0, + uint32_t n_ctx_orig, + bool inverse, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float eps); + int ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor( ds4_gpu_tensor *out, ds4_gpu_tensor *q_half, @@ -2250,6 +2283,31 @@ int ds4_gpu_attention_output_q8_batch_tensor( uint64_t out_dim, const ds4_gpu_tensor *heads, uint32_t n_tokens); + +#ifdef __APPLE__ +/* Optional pre-M5 batch Q8 attention-output B + HC4 epilogue. Returns 1 + * when fused work was encoded, 0 when unsupported without encoding work, + * and -1 after an attempted-path failure. */ +int ds4_gpu_attention_output_q8_batch_hc_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + ds4_gpu_tensor *low, + ds4_gpu_tensor *group_tmp, + ds4_gpu_tensor *low_tmp, + const void *model_map, + uint64_t model_size, + uint64_t out_a_offset, + uint64_t out_b_offset, + uint64_t group_dim, + uint64_t rank, + uint32_t n_groups, + uint64_t out_dim, + const ds4_gpu_tensor *heads, + uint32_t n_tokens, + uint32_t n_hc); +#endif int ds4_gpu_attention_output_q4_K_batch_tensor( ds4_gpu_tensor *out, ds4_gpu_tensor *low, @@ -2385,6 +2443,27 @@ int ds4_gpu_router_select_tensor( bool hash_mode, const ds4_gpu_tensor *logits); +/* One-token router select with the token id read from a device buffer + * (greedy chain decode); identical kernels to ds4_gpu_router_select_tensor. */ +int ds4_gpu_router_select_tensor_devtoken( + ds4_gpu_tensor *selected, + ds4_gpu_tensor *weights, + ds4_gpu_tensor *probs, + const void *model_map, + uint64_t model_size, + uint64_t bias_offset, + uint64_t hash_offset, + uint32_t hash_rows, + const ds4_gpu_tensor *token_dev, + uint32_t n_expert, + uint32_t n_expert_used, + float expert_weight_scale, + uint32_t n_expert_groups, + uint32_t n_group_used, + bool has_bias, + bool hash_mode, + const ds4_gpu_tensor *logits); + int ds4_gpu_router_select_batch_tensor( ds4_gpu_tensor *selected, ds4_gpu_tensor *weights, @@ -2687,7 +2766,12 @@ int ds4_gpu_routed_moe_one_tensor( const ds4_gpu_tensor *x, const ds4_gpu_tensor *add_in, uint32_t layer_index, - bool force_resident); + bool force_resident, + const ds4_gpu_tensor *hc_shared_out, + const ds4_gpu_tensor *hc_residual, + const ds4_gpu_tensor *hc_split, + ds4_gpu_tensor *hc_out, + int *hc_fused_out); int ds4_gpu_routed_moe_batch_tensor( ds4_gpu_tensor *out, @@ -2718,6 +2802,7 @@ int ds4_gpu_routed_moe_batch_tensor( uint32_t layer_index, uint32_t n_tokens, bool *mid_is_f16, + bool defer_sum6, bool force_resident); /* ========================================================================= @@ -2899,6 +2984,23 @@ int ds4_gpu_hc_expand_add_split_tensor( uint32_t n_embd, uint32_t n_hc); +#ifdef __APPLE__ +int ds4_gpu_moe_sum6_hc_expand_available(void); +int ds4_gpu_moe_sum6_hc_expand_split_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *expert_down, + const ds4_gpu_tensor *shared_out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc); +int ds4_gpu_test_moe_sum6_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *expert_down, + uint32_t n_embd, + uint32_t n_tokens); +#endif + int ds4_gpu_hc_expand_add_split_half_add_tensor( ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, diff --git a/ds4_metal.m b/ds4_metal.m index 3363d7df5..9f32ebff5 100644 --- a/ds4_metal.m +++ b/ds4_metal.m @@ -144,6 +144,7 @@ static id g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_pipeline_nsg1; static id g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_static_pipeline_nsg1; static id g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_pipeline_nsg1; +static id g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_hc4_pipeline_nsg1; static id g_moe_mul_mv_slots6_mxfp4_pair_swiglu_pipeline; static id g_moe_mul_mv_slots6_mxfp4_sum6_pipeline; static id g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline; @@ -168,6 +169,8 @@ static id g_dsv4_indexer_qat_pipeline; static id g_dsv4_kv_fp8_store_pipeline; static id g_dsv4_kv_rope_fp8_store_pipeline; +static id g_dsv4_batch_qkv_finalize_pipeline; +static bool g_dsv4_batch_qkv_finalize_unavailable; static id g_dsv4_ratio4_shift_pipeline; static id g_dsv4_compressor_pack_ratio4_pipeline; static id g_dsv4_compressor_pack_ratio4_decode_ggml_pipeline; @@ -189,8 +192,11 @@ static id g_dsv4_indexed_attention_heads8_pipeline; static id g_dsv4_indexed_attention_heads8_rb16_pipeline; static id g_dsv4_indexed_attention_heads16_dual_pipeline; +static id g_dsv4_indexed_attention_heads8_rb4_pipeline; +static id g_dsv4_indexed_attention_heads16_dual_rb4_pipeline; static id g_dsv4_indexed_attention_heads8_split_pipeline; static id g_dsv4_indexed_attention_heads8_split_reduce_pipeline; +static bool g_attn_out_low_q8_static_unavailable; static id g_dsv4_softplus_sqrt_pipeline; static id g_dsv4_router_finalize_one_pipeline; static id g_dsv4_router_finalize_one_simd_pipeline; @@ -265,6 +271,8 @@ static id g_glm_q6_k_down_f32_pipeline; static id g_dsv4_router_weights_batch_pipeline; static id g_dsv4_hc_expand4_pipeline; +static id g_dsv4_attn_out_q8_mm_hc_expand4_pipeline; +static id g_dsv4_moe_sum6_hc_expand4_pipeline; static NSMutableDictionary> *g_pipeline_cache; enum { @@ -992,6 +1000,62 @@ static void ds4_gpu_close_batch_encoder(void) { static double g_gpu_busy_accum; static uint64_t g_gpu_busy_cbs; +/* --- GPU stage-counter profiler ----------------------------------------- + * Diagnostic only: at decode stage boundaries the open batch command buffer + * is committed without waiting, so the GPU queue stays fed and each stage's + * GPU busy span (GPUEndTime-GPUStartTime) can be read after the token. + * compute-encoder counter sampling is unsupported on Apple Silicon GPUs. */ +enum { DS4_STAGE_COUNTER_CAP = 2048u }; +static id g_stage_cbs[DS4_STAGE_COUNTER_CAP]; +static char g_stage_names[DS4_STAGE_COUNTER_CAP][40]; +static uint32_t g_stage_cb_idx; + +int ds4_gpu_stage_counters_enabled(void) { + static int enabled = -1; + if (enabled < 0) { + enabled = getenv("DS4_METAL_STAGE_COUNTERS") != NULL; + } + return enabled; +} + +int ds4_gpu_stage_counter_sample(const char *label) { + if (!label || !g_batch_cb) return 1; + if (g_stage_cb_idx >= DS4_STAGE_COUNTER_CAP) return 1; + if (g_batch_has_work) { + ds4_gpu_close_batch_encoder(); + id cb = g_batch_cb; + [cb commit]; + g_stage_cbs[g_stage_cb_idx] = cb; + snprintf(g_stage_names[g_stage_cb_idx], + sizeof(g_stage_names[0]), "%s", label); + g_stage_cb_idx++; + g_batch_cb = nil; + g_batch_has_work = NO; + return ds4_gpu_begin_commands() != 0; + } + /* No dispatches since the last boundary: nothing to attribute. */ + return 1; +} + +void ds4_gpu_stage_counter_reset(void) { + for (uint32_t i = 0; i < g_stage_cb_idx; i++) g_stage_cbs[i] = nil; + g_stage_cb_idx = 0; +} + +void ds4_gpu_stage_counter_report(uint32_t pos) { + double total_busy = 0.0; + for (uint32_t i = 0; i < g_stage_cb_idx; i++) { + const double busy = + g_stage_cbs[i].GPUEndTime - g_stage_cbs[i].GPUStartTime; + total_busy += busy; + fprintf(stderr, "ds4: stage-counter pos=%u seq=%u %s=%.4f ms\n", + pos, i, g_stage_names[i], busy * 1000.0); + } + fprintf(stderr, "ds4: stage-counter pos=%u total-cb-busy=%.4f ms n=%u\n", + pos, total_busy * 1000.0, g_stage_cb_idx); + ds4_gpu_stage_counter_reset(); +} + /* A failed command buffer can leave a cross-threadgroup arrival counter at an * arbitrary partial value. Drop cached ownership instead of CPU-resetting * buffers that another in-flight command buffer might still reference; bound @@ -1003,7 +1067,7 @@ static void ds4_gpu_invalidate_completion_counters(void) { g_dsv4_hc_producer_last_completion = nil; } -static int ds4_gpu_wait_command_buffer(id cb, const char *label) { +int ds4_gpu_wait_command_buffer(id cb, const char *label) { [cb waitUntilCompleted]; if (getenv("DS4_METAL_GPU_BUSY_PROFILE")) { const double busy = cb.GPUEndTime - cb.GPUStartTime; @@ -5141,6 +5205,34 @@ static int16_t ds4_gpu_mv_ext_r1ptg(uint64_t n_tok) { float eps; } ds4_gpu_qkv_rms_norm_args; +typedef struct { + int32_t n_head; + int32_t head_dim; + int32_t head_dim4; + int32_t n_dims; + int32_t n_ctx_orig; + int32_t pos0; + int32_t inverse; + float eps; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; +} ds4_gpu_dsv4_head_norm_rope_args; + +_Static_assert(sizeof(ds4_gpu_dsv4_head_norm_rope_args) == 56, + "Metal DSV4 head norm/RoPE argument ABI changed"); + +typedef struct { + uint32_t raw_cap; + uint32_t raw_pos0; +} ds4_gpu_dsv4_batch_qkv_finalize_args; + +_Static_assert(sizeof(ds4_gpu_dsv4_batch_qkv_finalize_args) == 8, + "Metal batch Q/KV finalizer argument ABI changed"); + static ds4_gpu_rms_norm_args ds4_gpu_make_rms_norm_args(uint32_t n, uint32_t rows, float eps) { const uint64_t row_bytes = (uint64_t)n * sizeof(float); return (ds4_gpu_rms_norm_args) { @@ -5512,6 +5604,14 @@ static int ds4_gpu_encode_mul_mm_id_addr_mapped_tile( uint32_t shared_pad; } ds4_gpu_flash_kv_stage_f16_args; +/* Matches ds4_metal_args_flash_kv_direct in metal/flash_attn.metal. */ +typedef struct { + uint32_t raw_cap; + uint32_t raw_start; + uint32_t n_raw; + uint32_t pad0; +} ds4_gpu_flash_kv_direct_args; + typedef struct { int32_t ne01; int32_t ne30; @@ -5751,7 +5851,9 @@ static int ds4_gpu_encode_rope_tail_inplace( pos[t] = (int32_t)(pos0 + t * pos_step); } - if (pos_bytes > 4096u) { + const bool force_pos_buffer = + getenv("DS4_GPU_TEST_BATCH_QKV_FINALIZE_POS_BUFFER") != NULL; + if (pos_bytes > 4096u || force_pos_buffer) { /* * Metal inline setBytes data is meant for small constants. Long * prefill RoPE calls need thousands of positions; passing that much @@ -7796,6 +7898,9 @@ int ds4_gpu_init(void) { g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_pipeline_nsg1 = ds4_gpu_new_mul_mv_tg_multiple_pipeline( "kernel_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_f32", 1); + g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_hc4_pipeline_nsg1 = + ds4_gpu_new_mul_mv_tg_multiple_pipeline( + "kernel_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_hc4_f32", 1); if (!g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline_nsg1 || !g_moe_mul_mv_id_mxfp4_sum6_pipeline_nsg1 || !g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline_nsg1_tg_multiple || @@ -7804,7 +7909,8 @@ int ds4_gpu_init(void) { !g_moe_mul_mv_id_mxfp4_sum6_fixed_route_pipeline_nsg1 || !g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_pipeline_nsg1 || !g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_static_pipeline_nsg1 || - !g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_pipeline_nsg1) { + !g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_pipeline_nsg1 || + !g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_hc4_pipeline_nsg1) { g_queue = nil; g_device = nil; return 0; @@ -8463,6 +8569,10 @@ int ds4_gpu_init(void) { ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads8_rb16"); g_dsv4_indexed_attention_heads16_dual_pipeline = ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads16_dual"); + g_dsv4_indexed_attention_heads8_rb4_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads8_rb4"); + g_dsv4_indexed_attention_heads16_dual_rb4_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads16_dual_rb4"); g_dsv4_indexed_attention_heads8_split_pipeline = ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads8_split"); g_dsv4_indexed_attention_heads8_split_reduce_pipeline = @@ -8608,6 +8718,10 @@ int ds4_gpu_init(void) { ds4_gpu_get_pipeline("kernel_dsv4_router_weights_batch"); g_dsv4_hc_expand4_pipeline = ds4_gpu_get_pipeline("kernel_dsv4_hc_expand4"); + g_dsv4_attn_out_q8_mm_hc_expand4_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_attn_out_q8_mm_hc_expand4_batch"); + g_dsv4_moe_sum6_hc_expand4_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_moe_sum6_hc_expand4"); if (!g_dsv4_indexer_score_one_direct_pipeline || !g_dsv4_compressor_store_one_pipeline || !g_dsv4_sort_i32_rows_asc_pipeline || @@ -8735,6 +8849,10 @@ void ds4_gpu_test_set_flags(uint32_t flags) { g_test_flags = flags; } +int ds4_gpu_test_get_quality(void) { + return g_quality_mode != 0; +} + ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes) { if (!g_initialized && !ds4_gpu_init()) return NULL; if (bytes == 0 || bytes > (uint64_t)NSUIntegerMax) return NULL; @@ -10307,6 +10425,7 @@ void ds4_gpu_cleanup(void) { g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_pipeline_nsg1 = nil; g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_static_pipeline_nsg1 = nil; g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_pipeline_nsg1 = nil; + g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_hc4_pipeline_nsg1 = nil; g_moe_mul_mv_slots6_mxfp4_pair_swiglu_pipeline = nil; g_moe_mul_mv_slots6_mxfp4_sum6_pipeline = nil; g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline = nil; @@ -10331,6 +10450,8 @@ void ds4_gpu_cleanup(void) { g_dsv4_indexer_qat_pipeline = nil; g_dsv4_kv_fp8_store_pipeline = nil; g_dsv4_kv_rope_fp8_store_pipeline = nil; + g_dsv4_batch_qkv_finalize_pipeline = nil; + g_dsv4_batch_qkv_finalize_unavailable = false; g_dsv4_ratio4_shift_pipeline = nil; g_dsv4_compressor_pack_ratio4_pipeline = nil; g_dsv4_compressor_pack_ratio4_decode_ggml_pipeline = nil; @@ -10352,8 +10473,11 @@ void ds4_gpu_cleanup(void) { g_dsv4_indexed_attention_heads8_pipeline = nil; g_dsv4_indexed_attention_heads8_rb16_pipeline = nil; g_dsv4_indexed_attention_heads16_dual_pipeline = nil; + g_dsv4_indexed_attention_heads8_rb4_pipeline = nil; + g_dsv4_indexed_attention_heads16_dual_rb4_pipeline = nil; g_dsv4_indexed_attention_heads8_split_pipeline = nil; g_dsv4_indexed_attention_heads8_split_reduce_pipeline = nil; + g_attn_out_low_q8_static_unavailable = false; g_dsv4_softplus_sqrt_pipeline = nil; g_dsv4_router_finalize_one_pipeline = nil; g_dsv4_router_finalize_one_simd_pipeline = nil; @@ -10429,6 +10553,8 @@ void ds4_gpu_cleanup(void) { g_glm_q6_k_down_f32_pipeline = nil; g_dsv4_router_weights_batch_pipeline = nil; g_dsv4_hc_expand4_pipeline = nil; + g_dsv4_attn_out_q8_mm_hc_expand4_pipeline = nil; + g_dsv4_moe_sum6_hc_expand4_pipeline = nil; g_flash_attn_mask_buffer = nil; g_flash_attn_zero_mask_buffer = nil; g_flash_attn_pad_buffer = nil; @@ -21305,22 +21431,7 @@ int ds4_gpu_head_rms_norm_rope_tail_tensor( "kernel_dsv4_head_rms_norm_rope_tail_f32"); if (!pipeline) return 0; - struct { - int32_t n_head; - int32_t head_dim; - int32_t head_dim4; - int32_t n_dims; - int32_t n_ctx_orig; - int32_t pos0; - int32_t inverse; - float eps; - float freq_base; - float freq_scale; - float ext_factor; - float attn_factor; - float beta_fast; - float beta_slow; - } args = { + ds4_gpu_dsv4_head_norm_rope_args args = { .n_head = (int32_t)n_head, .head_dim = (int32_t)head_dim, .head_dim4 = (int32_t)(head_dim / 4u), @@ -21362,6 +21473,209 @@ int ds4_gpu_head_rms_norm_rope_tail_tensor( return 1; } +int ds4_gpu_dsv4_batch_qnorm_rope_kv_finalize_available(void) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (g_dsv4_batch_qkv_finalize_unavailable || + !ds4_gpu_device_is_pre_m5_apple_silicon() || + !g_use_dsv4_head_rms_norm_rope_tail_pipeline || + !ds4_gpu_device_name_contains("M3") || + !g_rope_tail_inplace_pair_pipeline || + getenv("DS4_METAL_DISABLE_INPLACE_ROPE_PAIR") != NULL) { + return 0; + } + if (!g_dsv4_batch_qkv_finalize_pipeline) { + g_dsv4_batch_qkv_finalize_pipeline = ds4_gpu_get_pipeline( + "kernel_dsv4_batch_kv_rope_fp8_store_f32"); + } + if (!g_dsv4_batch_qkv_finalize_pipeline) { + g_dsv4_batch_qkv_finalize_unavailable = true; + return 0; + } + return 1; +} + +int ds4_gpu_dsv4_batch_qnorm_rope_kv_finalize_tensor( + ds4_gpu_tensor *q, + ds4_gpu_tensor *kv, + ds4_gpu_tensor *raw_cache, + uint32_t raw_cap, + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + uint32_t pos0, + uint32_t n_ctx_orig, + bool inverse, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!q || !kv || !raw_cache || raw_cap == 0u || n_tok <= 1u || + n_tok > raw_cap || n_head != 64u || head_dim != 512u || + n_rot != 64u || pos0 > (uint32_t)INT32_MAX - n_tok || + !ds4_gpu_dsv4_batch_qnorm_rope_kv_finalize_available()) { + return 0; + } + + @autoreleasepool { + id qbuf = ds4_gpu_tensor_buffer(q); + id kvbuf = ds4_gpu_tensor_buffer(kv); + id rawbuf = ds4_gpu_tensor_buffer(raw_cache); + const uint64_t q_bytes = + (uint64_t)n_tok * n_head * head_dim * sizeof(float); + const uint64_t kv_bytes = + (uint64_t)n_tok * head_dim * sizeof(float); + const uint64_t raw_bytes = + (uint64_t)raw_cap * head_dim * sizeof(float); + if (!qbuf || !kvbuf || !rawbuf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(kv) < kv_bytes || + ds4_gpu_tensor_bytes(raw_cache) < raw_bytes) { + fprintf(stderr, + "ds4: Metal batch Q/KV finalizer received undersized buffers\n"); + return 0; + } + + id kv_pipeline = + g_dsv4_batch_qkv_finalize_pipeline; + id q_pipeline = + g_dsv4_head_rms_norm_rope_tail_pipeline; + if (!q_pipeline || !kv_pipeline) { + g_dsv4_batch_qkv_finalize_unavailable = true; + return 0; + } + if (g_batch_encoder_concurrent) return 0; + + ds4_gpu_dsv4_head_norm_rope_args args = { + .n_head = (int32_t)n_head, + .head_dim = (int32_t)head_dim, + .head_dim4 = (int32_t)(head_dim / 4u), + .n_dims = (int32_t)n_rot, + .n_ctx_orig = (int32_t)n_ctx_orig, + .pos0 = (int32_t)pos0, + .inverse = inverse ? 1 : 0, + .eps = eps, + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + }; + ds4_gpu_dsv4_batch_qkv_finalize_args store = { + .raw_cap = raw_cap, + .raw_pos0 = pos0, + }; + + int32_t pos_stack[256]; + int32_t *positions = pos_stack; + if (n_tok > (uint32_t)(sizeof(pos_stack) / sizeof(pos_stack[0]))) { + positions = malloc((size_t)n_tok * sizeof(*positions)); + if (!positions) { + fprintf(stderr, + "ds4: failed to allocate batch Q/KV finalizer positions\n"); + return 0; + } + } + for (uint32_t t = 0; t < n_tok; t++) { + positions[t] = (int32_t)(pos0 + t); + } + + const NSUInteger pos_bytes = + (NSUInteger)n_tok * sizeof(*positions); + id posbuf = nil; + if (pos_bytes > 4096u) { + posbuf = ds4_gpu_new_transient_buffer( + pos_bytes, "ds4_batch_qkv_finalize_positions"); + if (!posbuf) { + if (positions != pos_stack) free(positions); + return 0; + } + memcpy([posbuf contents], positions, pos_bytes); + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) { + if (positions != pos_stack) free(positions); + return 0; + } + + const BOOL batch_owned_encoder = + g_batch_cb && cb == g_batch_cb; + if (batch_owned_encoder) { + ds4_gpu_close_batch_encoder(); + g_batch_encoder_concurrent = YES; + } + id enc = batch_owned_encoder + ? ds4_gpu_compute_encoder(cb) + : [cb computeCommandEncoderWithDispatchType:MTLDispatchTypeConcurrent]; + if (!enc || enc.dispatchType != MTLDispatchTypeConcurrent) { + if (batch_owned_encoder) { + ds4_gpu_close_batch_encoder(); + g_batch_encoder_concurrent = NO; + } + if (positions != pos_stack) free(positions); + return 0; + } + + /* Keep the established Q PSO intact: only its dispatch is moved into + * the concurrent encoder. This avoids the fast-math drift seen when + * the same source body was embedded in a compound Q/KV kernel. */ + [enc setComputePipelineState:q_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_head, n_tok, 1) + threadsPerThreadgroup:MTLSizeMake( + ds4_gpu_rms_norm_pipeline_threads(head_dim, q_pipeline), + 1, + 1)]; + + /* The KV PSO is independent of Q, and preserves the batch pair-RoPE, + * FP8 round trip, and F16-rounded raw-ring write in one threadgroup per + * token. DispatchTypeConcurrent lets the GPU hide this short tail + * under the 64 Q-head threadgroups for each token. */ + ds4_gpu_rope_tail_batch_args rope_args = + ds4_gpu_make_rope_tail_args( + n_tok, 1u, head_dim, n_rot, n_ctx_orig, inverse, + freq_base, freq_scale, ext_factor, attn_factor, + beta_fast, beta_slow); + [enc setComputePipelineState:kv_pipeline]; + [enc setBytes:&rope_args length:sizeof(rope_args) atIndex:0]; + [enc setBytes:&store length:sizeof(store) atIndex:1]; + if (posbuf) { + [enc setBuffer:posbuf offset:0 atIndex:2]; + } else { + [enc setBytes:positions length:pos_bytes atIndex:2]; + } + [enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv) atIndex:3]; + [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_cache) atIndex:4]; + [enc setThreadgroupMemoryLength:64u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1u, n_tok, 1) + threadsPerThreadgroup:MTLSizeMake(64u, 1, 1)]; + + if (batch_owned_encoder) { + ds4_gpu_close_batch_encoder(); + g_batch_encoder_concurrent = NO; + } else { + [enc endEncoding]; + } + + if (positions != pos_stack) free(positions); + if (!ds4_gpu_finish_command_buffer( + cb, owned, "batch Q norm/RoPE + KV finalizer")) { + return 0; + } + } + + return 1; +} + int ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor( ds4_gpu_tensor *out, ds4_gpu_tensor *q_half, @@ -24226,7 +24540,85 @@ static int ds4_gpu_encode_fill_f32_rows( return 1; } -int ds4_gpu_attention_output_q8_batch_tensor( +typedef struct { + ds4_gpu_tensor *out_hc; + const ds4_gpu_tensor *residual_hc; + const ds4_gpu_tensor *split; + id __strong weight_buffer; + NSUInteger weight_offset; +} ds4_gpu_attn_out_q8_hc_target; + +static int ds4_gpu_encode_attn_out_q8_mm_hc( + id cb, + const ds4_gpu_attn_out_q8_hc_target *target, + const ds4_gpu_tensor *low, + uint32_t n_tokens) { + if (!cb || !target || !target->out_hc || !target->residual_hc || + !target->split || !target->weight_buffer || !low || n_tokens == 0u) { + return 0; + } + + id lowbuf = ds4_gpu_tensor_buffer(low); + id resbuf = ds4_gpu_tensor_buffer(target->residual_hc); + id splitbuf = ds4_gpu_tensor_buffer(target->split); + id outbuf = ds4_gpu_tensor_buffer(target->out_hc); + if (!lowbuf || !resbuf || !splitbuf || !outbuf) return 0; + + const uint32_t in_dim = 8192u; + const uint32_t out_dim = 4096u; + const uint32_t n_hc = 4u; + const uint64_t row_bytes = (uint64_t)(in_dim / 32u) * 34u; + ds4_gpu_mul_mm_args mm = + ds4_gpu_make_mm_args(in_dim, out_dim, n_tokens, row_bytes); + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + ds4_gpu_hc_expand_args hc = { + .n_embd = out_dim, + .n_hc = n_hc, + .n_tokens = (int64_t)n_tokens, + .nb_block0 = sizeof(float), + .nb_block1 = (uint64_t)out_dim * sizeof(float), + .nb_add0 = sizeof(float), + .nb_add1 = (uint64_t)out_dim * sizeof(float), + .nb_res0 = sizeof(float), + .nb_res1 = (uint64_t)out_dim * sizeof(float), + .nb_res2 = (uint64_t)n_hc * out_dim * sizeof(float), + .nb_post0 = sizeof(float), + .nb_post1 = mix_hc * sizeof(float), + .nb_comb0 = sizeof(float), + .nb_comb1 = (uint64_t)n_hc * sizeof(float), + .nb_comb2 = mix_hc * sizeof(float), + .nb0 = sizeof(float), + .nb1 = (uint64_t)out_dim * sizeof(float), + .nb2 = (uint64_t)n_hc * out_dim * sizeof(float), + .has_add = 0, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_attn_out_q8_mm_hc_expand4_pipeline]; + [enc setBytes:&mm length:sizeof(mm) atIndex:0]; + [enc setBuffer:target->weight_buffer offset:target->weight_offset atIndex:1]; + [enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(low) atIndex:2]; + [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(target->residual_hc) atIndex:3]; + [enc setBuffer:splitbuf + offset:ds4_gpu_tensor_offset(target->split) + + (NSUInteger)n_hc * sizeof(float) + atIndex:4]; + [enc setBuffer:splitbuf + offset:ds4_gpu_tensor_offset(target->split) + + (NSUInteger)(2u * n_hc) * sizeof(float) + atIndex:5]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(target->out_hc) atIndex:6]; + [enc setBytes:&hc length:sizeof(hc) atIndex:7]; + [enc setThreadgroupMemoryLength:8192u atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens / 32u, + (NSUInteger)out_dim / 64u, + 1u) + threadsPerThreadgroup:MTLSizeMake(128u, 1u, 1u)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_attention_output_q8_batch_impl( ds4_gpu_tensor *out, ds4_gpu_tensor *low, ds4_gpu_tensor *group_tmp, @@ -24240,7 +24632,8 @@ int ds4_gpu_attention_output_q8_batch_tensor( uint32_t n_groups, uint64_t out_dim, const ds4_gpu_tensor *heads, - uint32_t n_tokens) { + uint32_t n_tokens, + const ds4_gpu_attn_out_q8_hc_target *hc_target) { if (!g_initialized && !ds4_gpu_init()) return 0; if (!out || !low || !group_tmp || !low_tmp || !heads || !model_map || group_dim == 0 || rank == 0 || n_groups == 0 || out_dim == 0 || n_tokens == 0 || @@ -24554,9 +24947,14 @@ int ds4_gpu_attention_output_q8_batch_tensor( DS4_METAL_PROFILE_ATTN_OUT_STAGE("low_proj"); if (ok) { - ok = ds4_gpu_matmul_q8_0_tensor(out, model_map, model_size, - out_b_offset, - low_dim, out_dim, low, n_tokens) != 0; + if (hc_target) { + ok = ds4_gpu_encode_attn_out_q8_mm_hc( + cb, hc_target, low, n_tokens) != 0; + } else { + ok = ds4_gpu_matmul_q8_0_tensor(out, model_map, model_size, + out_b_offset, + low_dim, out_dim, low, n_tokens) != 0; + } } DS4_METAL_PROFILE_ATTN_OUT_STAGE("out_proj"); @@ -24568,6 +24966,135 @@ int ds4_gpu_attention_output_q8_batch_tensor( } } +int ds4_gpu_attention_output_q8_batch_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *low, + ds4_gpu_tensor *group_tmp, + ds4_gpu_tensor *low_tmp, + const void *model_map, + uint64_t model_size, + uint64_t out_a_offset, + uint64_t out_b_offset, + uint64_t group_dim, + uint64_t rank, + uint32_t n_groups, + uint64_t out_dim, + const ds4_gpu_tensor *heads, + uint32_t n_tokens) { + return ds4_gpu_attention_output_q8_batch_impl( + out, low, group_tmp, low_tmp, model_map, model_size, + out_a_offset, out_b_offset, group_dim, rank, n_groups, out_dim, + heads, n_tokens, NULL); +} + +int ds4_gpu_attention_output_q8_batch_hc_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + ds4_gpu_tensor *low, + ds4_gpu_tensor *group_tmp, + ds4_gpu_tensor *low_tmp, + const void *model_map, + uint64_t model_size, + uint64_t out_a_offset, + uint64_t out_b_offset, + uint64_t group_dim, + uint64_t rank, + uint32_t n_groups, + uint64_t out_dim, + const ds4_gpu_tensor *heads, + uint32_t n_tokens, + uint32_t n_hc) { + if (!g_initialized && !ds4_gpu_init()) return -1; + + const bool force = + (g_test_flags & DS4_GPU_TEST_BATCH_ATTN_OUT_HC_FUSION) != 0u; + const NSUInteger dynamic_bytes = 8192u; + const NSUInteger static_bytes = g_dsv4_attn_out_q8_mm_hc_expand4_pipeline ? + g_dsv4_attn_out_q8_mm_hc_expand4_pipeline.staticThreadgroupMemoryLength : 0u; + const NSUInteger max_tg_bytes = g_device ? + g_device.maxThreadgroupMemoryLength : 0u; + const bool pipeline_ok = + g_dsv4_attn_out_q8_mm_hc_expand4_pipeline != nil && + g_dsv4_attn_out_q8_mm_hc_expand4_pipeline.threadExecutionWidth == 32u && + g_dsv4_attn_out_q8_mm_hc_expand4_pipeline.maxTotalThreadsPerThreadgroup >= 128u && + static_bytes <= max_tg_bytes && + dynamic_bytes <= max_tg_bytes - static_bytes; + const uint64_t low_dim = (uint64_t)n_groups * rank; + const uint64_t row_a_bytes = group_dim / 32u * 34u; + const uint64_t row_b_bytes = low_dim / 32u * 34u; + const uint64_t out_a_bytes = (uint64_t)n_groups * rank * row_a_bytes; + const uint64_t out_b_bytes = out_dim * row_b_bytes; + const uint64_t heads_bytes = + (uint64_t)n_tokens * n_groups * group_dim * sizeof(float); + const uint64_t low_bytes = + (uint64_t)n_tokens * low_dim * sizeof(float); + const uint64_t out_bytes = + (uint64_t)n_tokens * out_dim * sizeof(float); + const uint64_t hc_bytes = + (uint64_t)n_tokens * n_hc * out_dim * sizeof(float); + const uint64_t split_bytes = + (uint64_t)n_tokens * (2ull * n_hc + (uint64_t)n_hc * n_hc) * + sizeof(float); + bool eligible = + out && out_hc && residual_hc && split && low && group_tmp && low_tmp && + heads && model_map && + (ds4_gpu_device_is_pre_m5_apple_silicon() || force) && + (!g_quality_mode || force) && + !g_ssd_streaming_mode && !ds4_gpu_tp_world_is_two() && + getenv("DS4_METAL_DISABLE_PRE_M5_BATCH_ATTN_OUT_HC_FUSION") == NULL && + getenv("DS4_METAL_ATTN_OUT_STAGE_PROFILE") == NULL && + getenv("DS4_METAL_Q8_PREFILL_PROFILE") == NULL && + pipeline_ok && group_dim == 4096u && rank == 1024u && + n_groups == 8u && low_dim == 8192u && out_dim == 4096u && + n_hc == 4u && (n_tokens >= 512u || force) && + n_tokens <= 4096u && + (n_tokens % 32u) == 0u && + out_a_offset <= model_size && out_a_bytes <= model_size - out_a_offset && + out_b_offset <= model_size && out_b_bytes <= model_size - out_b_offset && + ds4_gpu_tensor_buffer(out) != nil && + ds4_gpu_tensor_buffer(out_hc) != nil && + ds4_gpu_tensor_buffer(residual_hc) != nil && + ds4_gpu_tensor_buffer(split) != nil && + ds4_gpu_tensor_buffer(low) != nil && + ds4_gpu_tensor_buffer(heads) != nil && + ds4_gpu_tensor_bytes(out) >= out_bytes && + ds4_gpu_tensor_bytes(out_hc) >= hc_bytes && + ds4_gpu_tensor_bytes(residual_hc) >= hc_bytes && + ds4_gpu_tensor_bytes(split) >= split_bytes && + ds4_gpu_tensor_bytes(low) >= low_bytes && + ds4_gpu_tensor_bytes(heads) >= heads_bytes; + + uint64_t out_b_inner = 0u; + id out_b_buf = nil; + if (eligible) { + out_b_buf = ds4_gpu_wrap_model_range( + model_map, model_size, out_b_offset, out_b_bytes, &out_b_inner); + eligible = out_b_buf != nil && out_b_inner <= NSUIntegerMax; + } + if (!eligible) { + if (force && + getenv("DS4_METAL_REQUIRE_PRE_M5_BATCH_ATTN_OUT_HC_FUSION") != NULL) { + fprintf(stderr, + "ds4: required Metal batch attention-output HC fusion was not selected\n"); + } + return 0; + } + + ds4_gpu_attn_out_q8_hc_target target = { + .out_hc = out_hc, + .residual_hc = residual_hc, + .split = split, + .weight_buffer = out_b_buf, + .weight_offset = (NSUInteger)out_b_inner, + }; + return ds4_gpu_attention_output_q8_batch_impl( + out, low, group_tmp, low_tmp, model_map, model_size, + out_a_offset, out_b_offset, group_dim, rank, n_groups, out_dim, + heads, n_tokens, &target) ? 1 : -1; +} + int ds4_gpu_attention_output_q4_K_batch_tensor( ds4_gpu_tensor *out, ds4_gpu_tensor *low, @@ -25268,8 +25795,37 @@ int ds4_gpu_attention_output_low_q8_tensor( .nb1 = (uint64_t)rank * sizeof(float), .nr0 = 2, }; - id pipeline = - ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_attn_out_low_q8_0_f32", 4); + /* The one-token Flash decode shape is invariant: eight independent + * 4096 -> 1024 Q8_0 projections. Its static-trip PSO retains the + * generic NR0=2 / NSG=4 arithmetic and reduction tree; only group + * offsets and loop bounds become compile-time constants. */ + const bool force_flash_decode_static_for_test = + (g_test_flags & DS4_GPU_TEST_ATTN_OUT_LOW_Q8_STATIC) != 0u; + const bool use_pre_m5_flash_decode_static = + (ds4_gpu_device_is_pre_m5_apple_silicon() || + force_flash_decode_static_for_test) && + (!g_attn_out_low_q8_static_unavailable || + force_flash_decode_static_for_test) && + getenv("DS4_METAL_DISABLE_PRE_M5_DECODE_PORTS") == NULL && + getenv("DS4_METAL_DISABLE_PRE_M5_ATTN_OUT_LOW_Q8_STATIC") == NULL && + group_dim == 4096u && rank == 1024u && n_groups == 8u && + low_dim == 8192u && row_a_bytes == 4352u && + out_a_bytes == 35651584u; + id pipeline = ds4_gpu_get_mul_mv_pipeline( + use_pre_m5_flash_decode_static ? + "kernel_dsv4_attn_out_low_q8_0_flash_decode_static_f32" : + "kernel_dsv4_attn_out_low_q8_0_f32", + 4); + if (!pipeline && use_pre_m5_flash_decode_static) { + /* A custom/older Metal source may not contain this optional + * kernel. Latch the miss until cleanup so decode does not + * retry PSO creation and log once per layer and token. */ + g_attn_out_low_q8_static_unavailable = true; + if (!force_flash_decode_static_for_test) { + pipeline = ds4_gpu_get_mul_mv_pipeline( + "kernel_dsv4_attn_out_low_q8_0_f32", 4); + } + } ok = ds4_gpu_encode_attn_out_low_q8_direct(cb, pipeline, &args, @@ -25876,9 +26432,9 @@ static int ds4_gpu_encode_flash_kv_stage_f16( bool shared_pad, bool *did_fuse_pad) { if (did_fuse_pad) *did_fuse_pad = false; - if (!cb || !raw || !comp || !dst || raw_cap == 0 || + if (!cb || !raw || !dst || raw_cap == 0 || raw_start >= raw_cap || n_raw == 0 || n_raw > raw_cap || - n_comp == 0 || head_dim == 0) { + (n_comp != 0 && !comp) || head_dim == 0) { return 0; } @@ -25965,6 +26521,9 @@ static int ds4_gpu_encode_flash_kv_stage_f16( dst_offset)) { return 0; } + if (n_comp == 0) { + return 1; + } return ds4_gpu_encode_copy_to_f16_1d( cb, comp, @@ -27716,7 +28275,9 @@ static int ds4_gpu_encode_flash_attention_gathered_heads( ds4_gpu_ported_m5_decode_feature_enabled( "DS4_METAL_DISABLE_PRE_M5_FLASH_ATTN_PACKED32_REDUCE", NULL) && - !g_quality_mode && use_mask == 0u && comp_kv_f16 != 0u && n_comp != 0u && + !g_quality_mode && use_mask == 0u && comp_kv_f16 != 0u && + (n_comp != 0u || + getenv("DS4_METAL_DISABLE_DECODE_RAW_PACKED32") == NULL) && n_head == 64u && head_dim == 512u && nsg == 1u && nwg == 32u && n_keys <= 1024u && g_decode_attn_rope_fuse != 0 && g_decode_attn_rope_args.head_dim == 512 && @@ -27786,6 +28347,36 @@ static int ds4_gpu_encode_flash_attention_gathered_heads( packed_pipeline.threadExecutionWidth == 32u && packed_pipeline.maxTotalThreadsPerThreadgroup >= packed_threads && (max_tgmem == 0u || max_tgmem >= packed_shared_bytes); + /* Direct-KV read (OPT-IN, default off): the packed32 consumer computes + * each row's raw-ring / compressed-cache source inline (bit-exact with + * the staged layout), skipping the staging dispatch. On M3 Ultra this + * measured ~2x slower in the attention segment: the F32 raw ring rows + * are re-read and re-converted (float4 -> half4 -> float4) by all 64 + * head threadgroups on every token, doubling raw-region traffic versus + * the one-time staging conversion. Keep for A/B and future devices: + * enable with DS4_METAL_ENABLE_DECODE_RAW_DIRECT_KV. */ + const bool direct_kv_requested = + packed_requested && + getenv("DS4_METAL_ENABLE_DECODE_RAW_DIRECT_KV") != NULL && + getenv("DS4_METAL_DISABLE_DECODE_RAW_DIRECT_KV") == NULL; + id direct_kv_pipeline = nil; + if (direct_kv_requested) { + direct_kv_pipeline = ds4_gpu_get_flash_attn_vec_pipeline( + "kernel_dsv4_flash_attn_vec_packed32_reduce_rope_f16_dk512_dv512_direct_kv", + true, true, false, false, false, false, + (int32_t)head_dim, (int32_t)head_dim, 1, 32); + } + const bool use_direct_kv = + direct_kv_pipeline != nil && + direct_kv_pipeline.threadExecutionWidth == 32u && + direct_kv_pipeline.maxTotalThreadsPerThreadgroup >= packed_threads && + (max_tgmem == 0u || max_tgmem >= packed_shared_bytes); + if (getenv("DS4_METAL_REQUIRE_DECODE_RAW_DIRECT_KV") != NULL && + direct_kv_requested && !use_direct_kv) { + fprintf(stderr, + "ds4: required Metal decode direct-KV attention kernel was not selected\n"); + return 0; + } if (packed_requested && getenv("DS4_METAL_TRACE_M5_FLASH_ATTN_PACKED32_REDUCE") != NULL) { fprintf(stderr, @@ -27828,7 +28419,8 @@ static int ds4_gpu_encode_flash_attention_gathered_heads( } bool pad_fused = false; - if (!ds4_gpu_encode_flash_kv_stage_f16( + if (!use_direct_kv && + !ds4_gpu_encode_flash_kv_stage_f16( cb, rawbuf, ds4_gpu_tensor_offset(raw_kv), @@ -27853,12 +28445,12 @@ static int ds4_gpu_encode_flash_attention_gathered_heads( } id pad_pipeline = nil; - if (has_kvpad && !pad_fused) { + if (has_kvpad && !pad_fused && !use_direct_kv) { pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); if (!pad_pipeline) return 0; } - if (has_kvpad && !pad_fused) { + if (has_kvpad && !pad_fused && !use_direct_kv) { ds4_gpu_flash_attn_pad_args pad_args = { .ne11 = (int32_t)n_keys, .ne_12_2 = 1, @@ -27929,6 +28521,41 @@ static int ds4_gpu_encode_flash_attention_gathered_heads( 2u * ds4_gpu_align_up_ns(head_dim, 128u)) * nsg; const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); + if (use_direct_kv) { + ds4_gpu_flash_kv_direct_args kv_args = { + .raw_cap = raw_cap, + .raw_start = raw_start, + .n_raw = n_raw, + .pad0 = 0, + }; + id direct_enc = ds4_gpu_compute_encoder(cb); + [direct_enc setComputePipelineState:direct_kv_pipeline]; + [direct_enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0]; + [direct_enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [direct_enc setBuffer:rawbuf + offset:ds4_gpu_tensor_offset(raw_kv) atIndex:2]; + /* n_comp == 0 leaves compbuf nil; bind a dummy (the kernel never + * reads it: every row is a raw row when n_keys == n_raw). */ + [direct_enc setBuffer:(compbuf ? compbuf : rawbuf) + offset:(compbuf ? ds4_gpu_tensor_offset(comp_kv) + : ds4_gpu_tensor_offset(raw_kv)) + atIndex:3]; + [direct_enc setBuffer:flash_mask_buffer offset:0 atIndex:4]; + [direct_enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; + [direct_enc setBytes:&kv_args length:sizeof(kv_args) atIndex:6]; + [direct_enc setBuffer:headsbuf + offset:ds4_gpu_tensor_offset(heads) atIndex:7]; + [direct_enc setBytes:&g_decode_attn_rope_args + length:sizeof(g_decode_attn_rope_args) atIndex:8]; + [direct_enc setThreadgroupMemoryLength:packed_shared_bytes atIndex:0]; + [direct_enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(packed_threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, direct_enc); + g_decode_attn_rope_fuse = 0; + g_decode_attn_rope_fuse_used = 1; + return 1; + } + if (use_packed) { id packed_enc = ds4_gpu_compute_encoder(cb); [packed_enc setComputePipelineState:packed_pipeline]; @@ -28955,6 +29582,49 @@ int ds4_gpu_attention_indexed_mixed_batch_heads_tensor( !decode_one_token && !g_quality_mode && ds4_gpu_mpp_available() && n_head == 64u && top_k == 512u && window == 128u && head_dim == 512u; + const bool force_prefill_heads8_rb4_for_test = + (g_test_flags & DS4_GPU_TEST_INDEXED_ATTN_PREFILL_RB4) != 0u; + const bool prefill_heads8_rb4 = + !decode_one_token && !g_quality_mode && n_tokens >= 32u && + n_head == 64u && top_k == 512u && window == 128u && + head_dim == 512u && ratio == 4u && + n_comp > top_k && comp_kv_f16 != 0u && !g_ssd_streaming_mode && + !ds4_gpu_tp_world_is_two() && + (ds4_gpu_device_is_pre_m5_apple_silicon() || + force_prefill_heads8_rb4_for_test) && + getenv("DS4_METAL_DISABLE_PRE_M5_INDEXED_ATTN_PREFILL_RB4") == NULL && + g_dsv4_indexed_attention_heads8_rb4_pipeline != nil && + g_dsv4_indexed_attention_heads8_rb4_pipeline.threadExecutionWidth == 32u && + g_dsv4_indexed_attention_heads8_rb4_pipeline.maxTotalThreadsPerThreadgroup >= 256u; + const NSUInteger dual_rb4_static_bytes = + g_dsv4_indexed_attention_heads16_dual_rb4_pipeline != nil ? + g_dsv4_indexed_attention_heads16_dual_rb4_pipeline.staticThreadgroupMemoryLength : 0u; + const NSUInteger max_tg_bytes = g_device ? g_device.maxThreadgroupMemoryLength : 0u; + const bool prefill_heads16_dual_rb4 = + prefill_heads8_rb4 && + getenv("DS4_METAL_DISABLE_PRE_M5_INDEXED_ATTN_PREFILL_HEADS16_DUAL_RB4") == NULL && + g_dsv4_indexed_attention_heads16_dual_rb4_pipeline != nil && + g_dsv4_indexed_attention_heads16_dual_rb4_pipeline.threadExecutionWidth == 32u && + g_dsv4_indexed_attention_heads16_dual_rb4_pipeline.maxTotalThreadsPerThreadgroup >= 256u && + (max_tg_bytes == 0u || + (dual_rb4_static_bytes <= max_tg_bytes && + 4096u <= max_tg_bytes - dual_rb4_static_bytes)); + if (force_prefill_heads8_rb4_for_test && + getenv("DS4_METAL_REQUIRE_PRE_M5_INDEXED_ATTN_PREFILL_HEADS16_DUAL_RB4") != NULL && + !prefill_heads16_dual_rb4) { + fprintf(stderr, + "ds4: required Metal indexed-attention prefill heads16 dual RB4 kernel was not selected\n"); + return 0; + } + if (force_prefill_heads8_rb4_for_test && + getenv("DS4_METAL_REQUIRE_PRE_M5_INDEXED_ATTN_PREFILL_RB4") != NULL && + !prefill_heads8_rb4) { + fprintf(stderr, + "ds4: required Metal indexed-attention prefill RB4 kernel was not selected\n"); + return 0; + } + const bool use_prefill_dual_heads = + prefill_dual_heads && !prefill_heads8_rb4; const uint32_t decode_splits = decode_one_token && !g_quality_mode ? 12u : 1u; const bool split_decode = decode_splits > 1u; @@ -28966,7 +29636,13 @@ int ds4_gpu_attention_indexed_mixed_batch_heads_tensor( decode_one_token ? ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads8_rb16_pipeline, "kernel_dsv4_indexed_mixed_attention_heads8_rb16") : - prefill_dual_heads ? + prefill_heads16_dual_rb4 ? + ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads16_dual_rb4_pipeline, + "kernel_dsv4_indexed_mixed_attention_heads16_dual_rb4") : + prefill_heads8_rb4 ? + ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads8_rb4_pipeline, + "kernel_dsv4_indexed_mixed_attention_heads8_rb4") : + use_prefill_dual_heads ? ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads16_dual_pipeline, "kernel_dsv4_indexed_mixed_attention_heads16_dual") : ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads8_pipeline, @@ -29103,13 +29779,15 @@ int ds4_gpu_attention_indexed_mixed_batch_heads_tensor( atIndex:4]; [enc setBuffer:sinks_buf offset:(NSUInteger)sinks_inner atIndex:5]; [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:6]; - [enc setThreadgroupMemoryLength:(decode_one_token ? 16u : 1u) * + [enc setThreadgroupMemoryLength:(decode_one_token ? 16u : + prefill_heads8_rb4 ? 4u : 1u) * 128u * 4u * sizeof(uint16_t) atIndex:0]; [enc dispatchThreadgroups: MTLSizeMake((NSUInteger)n_tokens, - ((NSUInteger)n_head + (prefill_dual_heads ? 15u : 7u)) / - (prefill_dual_heads ? 16u : 8u), + ((NSUInteger)n_head + + (prefill_heads16_dual_rb4 || use_prefill_dual_heads ? 15u : 7u)) / + (prefill_heads16_dual_rb4 || use_prefill_dual_heads ? 16u : 8u), 1) threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; ds4_gpu_end_compute_encoder(cb, enc); @@ -29345,7 +30023,16 @@ int ds4_gpu_attention_decode_heads_tensor( id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, sinks_offset, sink_bytes, &sinks_inner); if (!sinks_buf) return 0; - if (n_comp == 0) { + /* Raw-only layers (ratio 0/128 with n_comp == 0) historically used a + * five-dispatch raw path (ring copy, standalone pad, vec, plain + * reduce, standalone inverse-RoPE tail). The gathered path below + * handles n_comp == 0 identically (raw-rows-first staging with pad + * fusion, rope-fused reduce) in two to three dispatches with the + * same kernels and reduction topology, so decode routes there by + * default; the env rollback restores the raw path. */ + if (n_comp == 0 && + (use_mask != 0 || + getenv("DS4_METAL_DISABLE_DECODE_RAW_GATHERED_ATTN") != NULL)) { int owned = 0; id cb = ds4_gpu_command_buffer(&owned); if (!cb) return 0; @@ -30618,6 +31305,68 @@ static int ds4_gpu_encode_mul_mv_id_sum6( return 1; } +/* HC tail-fused variant of ds4_gpu_encode_mul_mv_id_sum6 for + * kernel_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_hc4_f32: same + * grid and matvec arguments, plus the pre-materialized shared_out addend + * and the HC=4 residual/post/comb/dst streams. */ +static int ds4_gpu_encode_mul_mv_id_sum6_hc4( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_hc_expand_args *hc, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + id shared_out, + NSUInteger shared_out_off, + id residual, + NSUInteger residual_off, + id post, + NSUInteger post_off, + id comb, + NSUInteger comb_off, + id hc_dst, + NSUInteger hc_dst_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg) { + if (!cb || !pipeline || !args || !hc || !src0 || !src1 || !dst || !ids || + !shared_out || !residual || !post || !comb || !hc_dst || + args->ne00 <= 0 || args->ne01 <= 0 || + args->nei0 <= 0 || args->nei0 > 8 || args->nei1 <= 0) { + return 0; + } + + const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:hc length:sizeof(*hc) atIndex:1]; + [enc setBuffer:src0 offset:src0_off atIndex:2]; + [enc setBuffer:src1 offset:src1_off atIndex:3]; + [enc setBuffer:dst offset:dst_off atIndex:4]; + [enc setBuffer:ids offset:ids_off atIndex:5]; + [enc setBuffer:dst offset:dst_off atIndex:6]; + [enc setBuffer:shared_out offset:shared_out_off atIndex:7]; + [enc setBuffer:residual offset:residual_off atIndex:8]; + [enc setBuffer:post offset:post_off atIndex:9]; + [enc setBuffer:comb offset:comb_off atIndex:10]; + [enc setBuffer:hc_dst offset:hc_dst_off atIndex:11]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + static int ds4_gpu_encode_q4_gather_slots6( id cb, id pipeline, @@ -31847,6 +32596,46 @@ static int ds4_gpu_encode_moe_sum6( return 1; } +int ds4_gpu_test_moe_sum6_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *expert_down, + uint32_t n_embd, + uint32_t n_tokens) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !expert_down || n_embd == 0 || n_tokens == 0) return 0; + + @autoreleasepool { + const uint64_t row_bytes = (uint64_t)n_embd * sizeof(float); + if ((uint64_t)n_tokens > UINT64_MAX / row_bytes || + (uint64_t)n_tokens > UINT64_MAX / (6u * row_bytes)) { + return 0; + } + const uint64_t out_bytes = (uint64_t)n_tokens * row_bytes; + const uint64_t expert_bytes = (uint64_t)n_tokens * 6u * row_bytes; + id outbuf = ds4_gpu_tensor_buffer(out); + id expertsbuf = ds4_gpu_tensor_buffer(expert_down); + if (!outbuf || !expertsbuf || + ds4_gpu_tensor_bytes(out) < out_bytes || + ds4_gpu_tensor_bytes(expert_down) < expert_bytes) { + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb || + !ds4_gpu_encode_moe_sum6(cb, + expertsbuf, + ds4_gpu_tensor_offset(expert_down), + outbuf, + ds4_gpu_tensor_offset(out), + n_embd, + n_tokens)) { + return 0; + } + return ds4_gpu_finish_command_buffer(cb, owned, "test MoE sum6"); + } +} + static int ds4_gpu_encode_moe_sum8( id cb, id experts, @@ -38283,6 +39072,105 @@ int ds4_gpu_router_select_tensor( return 1; } +int ds4_gpu_router_select_tensor_devtoken( + ds4_gpu_tensor *selected, + ds4_gpu_tensor *weights, + ds4_gpu_tensor *probs, + const void *model_map, + uint64_t model_size, + uint64_t bias_offset, + uint64_t hash_offset, + uint32_t hash_rows, + const ds4_gpu_tensor *token_dev, + uint32_t n_expert, + uint32_t n_expert_used, + float expert_weight_scale, + uint32_t n_expert_groups, + uint32_t n_group_used, + bool has_bias, + bool hash_mode, + const ds4_gpu_tensor *logits) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!selected || !weights || !probs || !logits || !model_map || !token_dev || + n_expert == 0 || n_expert_used == 0) return 0; + if (n_expert_groups > 1u || n_group_used > 0u) { + fprintf(stderr, "ds4: Metal router group gating is not part of this DeepSeek V4 path\n"); + return 0; + } + + @autoreleasepool { + id logitsbuf = ds4_gpu_tensor_buffer(logits); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + id probsbuf = ds4_gpu_tensor_buffer(probs); + id tokenbuf = ds4_gpu_tensor_buffer(token_dev); + if (!logitsbuf || !selectedbuf || !weightsbuf || !probsbuf || !tokenbuf || + ds4_gpu_tensor_bytes(logits) < (uint64_t)n_expert * sizeof(float) || + ds4_gpu_tensor_bytes(selected) < (uint64_t)n_expert_used * sizeof(int) || + ds4_gpu_tensor_bytes(weights) < (uint64_t)n_expert_used * sizeof(float) || + ds4_gpu_tensor_bytes(probs) < (uint64_t)n_expert * sizeof(float) || + ds4_gpu_tensor_bytes(token_dev) < sizeof(int32_t)) { + fprintf(stderr, "ds4: Metal router select received undersized buffers\n"); + return 0; + } + + uint64_t bias_inner = 0; + uint64_t hash_inner = 0; + id biasbuf = nil; + id hashbuf = nil; + NSUInteger bias_set_offset = 0; + NSUInteger hash_set_offset = 0; + if (has_bias && !hash_mode) { + const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float); + biasbuf = ds4_gpu_wrap_model_range(model_map, model_size, bias_offset, bias_bytes, &bias_inner); + if (!biasbuf) return 0; + bias_set_offset = (NSUInteger)bias_inner; + } + if (hash_mode) { + const uint64_t hash_bytes = (uint64_t)hash_rows * n_expert_used * sizeof(int32_t); + hashbuf = ds4_gpu_wrap_model_range(model_map, model_size, hash_offset, hash_bytes, &hash_inner); + if (!hashbuf) return 0; + hash_set_offset = (NSUInteger)hash_inner; + } + + const bool had_batch = g_batch_cb != nil; + if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + int ok = cb && + ds4_gpu_encode_router_select(cb, + selected, + weights, + probs, + logitsbuf, + ds4_gpu_tensor_offset(logits), + biasbuf, + bias_set_offset, + hashbuf, + hash_set_offset, + tokenbuf, + ds4_gpu_tensor_offset(token_dev), + nil, + 0, + NULL, + hash_rows, + 0, + 1, + n_expert, + n_expert_used, + expert_weight_scale, + has_bias && !hash_mode, + hash_mode, + false); + if (!had_batch) { + ok = ds4_gpu_end_commands() != 0 && ok; + } + if (!ok) return 0; + } + + return 1; +} + int ds4_gpu_router_select_batch_tensor( ds4_gpu_tensor *selected, ds4_gpu_tensor *weights, @@ -38518,10 +39406,16 @@ int ds4_gpu_routed_moe_one_tensor( const ds4_gpu_tensor *x, const ds4_gpu_tensor *add_in, uint32_t layer_index, - bool force_resident) { + bool force_resident, + const ds4_gpu_tensor *hc_shared_out, + const ds4_gpu_tensor *hc_residual, + const ds4_gpu_tensor *hc_split, + ds4_gpu_tensor *hc_out, + int *hc_fused_out) { BOOL parallel_ffn_scope __attribute__((cleanup(ds4_gpu_parallel_ffn_scope_cleanup))) = g_parallel_q8_pending; + if (hc_fused_out) *hc_fused_out = 0; if (!g_initialized && !ds4_gpu_init()) return 0; /* TP sharding: only the owned contiguous expert range is mapped, * so bind from the owned base, validate only its bytes, and tell the @@ -38576,6 +39470,29 @@ int ds4_gpu_routed_moe_one_tensor( fprintf(stderr, "ds4: Metal routed tensor MoE received undersized expert output buffer\n"); return 0; } + /* Decode HC tail fusion: the caller pre-encoded the plain shared-down + * matvec into hc_shared_out and expects the down-sum6 dispatch to add + * it and expand the four HC streams. mix_hc is 2*n_hc + n_hc*n_hc. */ + const uint64_t hc_bytes = 4ull * out_dim * sizeof(float); + const uint64_t hc_split_bytes = 24ull * sizeof(float); + id hc_sharedbuf = nil; + id hc_resbuf = nil; + id hc_splitbuf = nil; + id hc_outbuf = nil; + if (hc_out) { + hc_sharedbuf = ds4_gpu_tensor_buffer(hc_shared_out); + hc_resbuf = ds4_gpu_tensor_buffer(hc_residual); + hc_splitbuf = ds4_gpu_tensor_buffer(hc_split); + hc_outbuf = ds4_gpu_tensor_buffer(hc_out); + if (!hc_sharedbuf || !hc_resbuf || !hc_splitbuf || !hc_outbuf || + ds4_gpu_tensor_bytes(hc_shared_out) < out_bytes || + ds4_gpu_tensor_bytes(hc_residual) < hc_bytes || + ds4_gpu_tensor_bytes(hc_split) < hc_split_bytes || + ds4_gpu_tensor_bytes(hc_out) < hc_bytes) { + fprintf(stderr, "ds4: Metal routed MoE HC fusion received undersized buffers\n"); + return 0; + } + } if ((uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes || (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes) { @@ -38751,6 +39668,18 @@ int ds4_gpu_routed_moe_one_tensor( down_args.nb01 == 1088 && down_args.nei0 == 6 && g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_pipeline_nsg1 != nil; + /* Decode HC tail fusion on the exact static sum6 shape: the caller + * pre-encoded the plain shared-down matvec into hc_shared_out, and + * the fused down kernel adds it and expands the four HC streams. + * Bit-exact with the unfused pair (see the kernel comment). When + * the caller passed hc_out but this gate fails, the plain sum6 runs + * instead and *hc_fused_out stays 0 so the caller can re-run the + * legacy shared-down+HC dispatch. */ + const bool use_mxfp4_moe_decode_sum6_hc4 = + use_mxfp4_moe_decode_static_trip_down && + hc_out != NULL && + out_dim == 4096 && + g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_hc4_pipeline_nsg1 != nil; id pair_swiglu_pipeline = nil; if (gate_type == DS4_METAL_TENSOR_IQ2_XXS) { pair_swiglu_pipeline = g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline; @@ -38796,6 +39725,8 @@ int ds4_gpu_routed_moe_one_tensor( down_sum6_pipeline = g_moe_mul_mv_id_mxfp4_sum6_pipeline; if (ds4_gpu_mxfp4_moe_decode_nsg1_enabled(n_tokens)) { down_sum6_pipeline = + use_mxfp4_moe_decode_sum6_hc4 ? + g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_hc4_pipeline_nsg1 : use_mxfp4_moe_decode_static_trip_down ? g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_pipeline_nsg1 : use_mxfp4_moe_decode_sum6_full_rows ? @@ -41007,6 +41938,57 @@ int ds4_gpu_routed_moe_one_tensor( 2); } } else if (ok && direct_down_sum) { + if (use_mxfp4_moe_decode_sum6_hc4) { + const uint64_t mix_hc = 24ull; /* 2*n_hc + n_hc*n_hc, n_hc = 4 */ + ds4_gpu_hc_expand_args hc_args = { + .n_embd = (int64_t)out_dim, + .n_hc = 4, + .n_tokens = 1, + .nb_block0 = sizeof(float), + .nb_block1 = (uint64_t)out_dim * sizeof(float), + .nb_add0 = sizeof(float), + .nb_add1 = (uint64_t)out_dim * sizeof(float), + .nb_res0 = sizeof(float), + .nb_res1 = (uint64_t)out_dim * sizeof(float), + .nb_res2 = (uint64_t)4 * out_dim * sizeof(float), + .nb_post0 = sizeof(float), + .nb_post1 = mix_hc * sizeof(float), + .nb_comb0 = sizeof(float), + .nb_comb1 = (uint64_t)4 * sizeof(float), + .nb_comb2 = mix_hc * sizeof(float), + .nb0 = sizeof(float), + .nb1 = (uint64_t)out_dim * sizeof(float), + .nb2 = (uint64_t)4 * out_dim * sizeof(float), + .has_add = 1, + }; + ok = ds4_gpu_encode_mul_mv_id_sum6_hc4(cb, + down_sum6_pipeline, + &down_args, + &hc_args, + down_buf, + (NSUInteger)down_inner, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selectedbuf, + ds4_gpu_tensor_offset(selected), + hc_sharedbuf, + ds4_gpu_tensor_offset(hc_shared_out), + hc_resbuf, + ds4_gpu_tensor_offset(hc_residual), + hc_splitbuf, + ds4_gpu_tensor_offset(hc_split) + + (NSUInteger)4 * sizeof(float), + hc_splitbuf, + ds4_gpu_tensor_offset(hc_split) + + (NSUInteger)8 * sizeof(float), + hc_outbuf, + ds4_gpu_tensor_offset(hc_out), + down_smem, + down_sum6_nsg); + if (ok && hc_fused_out) *hc_fused_out = 1; + } else { ok = ds4_gpu_encode_mul_mv_id_sum6(cb, down_sum6_pipeline, &down_args, @@ -41022,6 +42004,7 @@ int ds4_gpu_routed_moe_one_tensor( add_in ? ds4_gpu_tensor_offset(add_in) : 0, down_smem, down_sum6_nsg); + } } else if (ok) { ok = ds4_gpu_encode_mul_mv_id(cb, down_mv_pipeline, @@ -41098,6 +42081,7 @@ int ds4_gpu_routed_moe_batch_tensor( uint32_t layer_index, uint32_t n_tokens, bool *mid_is_f16, + bool defer_sum6, bool force_resident) { (void)force_resident; if (!g_initialized && !ds4_gpu_init()) return 0; @@ -41127,6 +42111,7 @@ int ds4_gpu_routed_moe_batch_tensor( } const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes; const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes; + if (defer_sum6 && n_tokens < 32u) return 0; /* * PRO Q4 routed expert tensors are multi-GiB per layer. A one-token @@ -41256,7 +42241,8 @@ int ds4_gpu_routed_moe_batch_tensor( x, NULL, layer_index, - false); + false, + NULL, NULL, NULL, NULL, NULL); } @autoreleasepool { @@ -41461,6 +42447,17 @@ int ds4_gpu_routed_moe_batch_tensor( 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; + /* The established resident MXFP4 specializations were originally + * promoted only for >=2K-token prefills. Extend the same exact kernels + * to the 32..2047 range behind one aggregate rollback so the sparse + * expert map and tail tiles do not fall back to the scan-heavy path. */ + const bool use_pre_m5_mxfp4_small_prefill_defaults = + ds4_gpu_device_is_pre_m5_apple_silicon() && + n_tokens >= 32u && n_tokens < 2048u && + getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_SMALL_PREFILL") == NULL; + const bool use_pre_m5_mxfp4_prefill_defaults = + ds4_gpu_device_is_pre_m5_apple_silicon() && + (n_tokens >= 2048u || use_pre_m5_mxfp4_small_prefill_defaults); /* * The MXFP4 32x32 specialization uses two SIMDgroups and 8 KiB of * threadgroup memory, and exactly culls SIMDgroup 1 on at-most-16-row @@ -41468,9 +42465,8 @@ int ds4_gpu_routed_moe_batch_tensor( * it the resident pre-M5 default for large prefill. */ const bool use_pre_m5_mxfp4_mm_id_pair_swiglu_compact_tile_default = - ds4_gpu_device_is_pre_m5_apple_silicon() && + use_pre_m5_mxfp4_prefill_defaults && !g_ssd_streaming_mode && - n_tokens >= 2048u && getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_MM_ID_PAIR_SWIGLU_COMPACT_TILE") == NULL; const bool use_mxfp4_mm_id_pair_swiglu_compact_tile = use_mm_id_pair_swiglu && @@ -41485,8 +42481,7 @@ int ds4_gpu_routed_moe_batch_tensor( * by the existing padded direct launches and changes no arithmetic. */ const bool use_pre_m5_mxfp4_mm_id_map_scatter_default = - ds4_gpu_device_is_pre_m5_apple_silicon() && - n_tokens >= 2048u && + use_pre_m5_mxfp4_prefill_defaults && getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_MM_ID_MAP_SCATTER") == NULL; const bool use_mxfp4_mm_id_map_scatter = use_mxfp4_mm_id_pair_swiglu_compact_tile && @@ -41505,9 +42500,8 @@ int ds4_gpu_routed_moe_batch_tensor( * shape and prefixes covered by the full-model A/B gate. */ const bool use_pre_m5_mxfp4_mm_id_pair_tail_simdgroup_cull_default = - ds4_gpu_device_is_pre_m5_apple_silicon() && + use_pre_m5_mxfp4_prefill_defaults && !g_ssd_streaming_mode && - n_tokens >= 2048u && getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_MM_ID_PAIR_TAIL_SIMDGROUP_CULL") == NULL; const bool use_mxfp4_mm_id_pair_tail_simdgroup_cull = use_mm_id_pair_swiglu && @@ -41517,9 +42511,8 @@ int ds4_gpu_routed_moe_batch_tensor( (use_pre_m5_mxfp4_mm_id_pair_tail_simdgroup_cull_default || (g_test_flags & DS4_GPU_TEST_MXFP4_PAIR_TAIL_CULL) != 0u); const bool use_pre_m5_mxfp4_mm_id_down_tail_simdgroup_cull_default = - ds4_gpu_device_is_pre_m5_apple_silicon() && + use_pre_m5_mxfp4_prefill_defaults && !g_ssd_streaming_mode && - n_tokens >= 2048u && getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_MM_ID_DOWN_TAIL_SIMDGROUP_CULL") == NULL; const bool use_mxfp4_mm_id_down_tail_simdgroup_cull = use_mm_id && @@ -41535,8 +42528,7 @@ int ds4_gpu_routed_moe_batch_tensor( * pre-M5 Apple-Silicon default for large prefill. */ const bool use_pre_m5_mxfp4_mm_id_down_half_lut_default = - ds4_gpu_device_is_pre_m5_apple_silicon() && - n_tokens >= 2048u && + use_pre_m5_mxfp4_prefill_defaults && getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_MM_ID_DOWN_HALF_LUT") == NULL; const bool use_mxfp4_mm_id_down_half_lut = use_mm_id && @@ -41641,6 +42633,16 @@ int ds4_gpu_routed_moe_batch_tensor( } } + /* The graph may defer the fixed top-6 reduction into its HC consumer, + * but only when the normal grouped batch path writes all expert rows + * into the caller-owned tensor. Reject before opening/encoding a + * command buffer so a failed preflight can never leave partial work. */ + if (defer_sum6 && + (!use_mm_id || n_expert != 6u || !expertsbuf || + ds4_gpu_tensor_bytes(experts) < down_scratch_bytes)) { + return 0; + } + if (use_iq2_batch_selected_addr) { const int had_batch = g_batch_cb != nil; if (had_batch && ds4_gpu_end_commands() == 0) { @@ -42320,6 +43322,7 @@ int ds4_gpu_routed_moe_batch_tensor( DS4_METAL_PROFILE_MOE_STAGE("down"); if (ok && n_expert > 1 && + !defer_sum6 && !direct_down_sum && !use_q4_batch_expert_table && !use_iq2_batch_selected_addr) { @@ -43747,6 +44750,128 @@ int ds4_gpu_hc_expand_add_split_tensor( return 1; } +int ds4_gpu_moe_sum6_hc_expand_available(void) { + if (!g_initialized && !ds4_gpu_init()) return 0; + return g_dsv4_moe_sum6_hc_expand4_pipeline != nil && + g_dsv4_moe_sum6_hc_expand4_pipeline.maxTotalThreadsPerThreadgroup > 0u; +} + +int ds4_gpu_moe_sum6_hc_expand_split_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *expert_down, + const ds4_gpu_tensor *shared_out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out_hc || !expert_down || !shared_out || !residual_hc || !split || + n_embd == 0 || n_hc != 4 || + !ds4_gpu_moe_sum6_hc_expand_available()) { + return 0; + } + + @autoreleasepool { + id outbuf = ds4_gpu_tensor_buffer(out_hc); + id expertsbuf = ds4_gpu_tensor_buffer(expert_down); + id sharedbuf = ds4_gpu_tensor_buffer(shared_out); + id resbuf = ds4_gpu_tensor_buffer(residual_hc); + id splitbuf = ds4_gpu_tensor_buffer(split); + const uint64_t hc_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); + const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out_hc); + if (hc_row_bytes == 0 || out_tensor_bytes < hc_row_bytes || + out_tensor_bytes % hc_row_bytes != 0) { + fprintf(stderr, + "ds4: Metal MoE sum6 HC output size is not a whole HC token row\n"); + return 0; + } + + const uint64_t n_tokens64 = out_tensor_bytes / hc_row_bytes; + if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) { + fprintf(stderr, + "ds4: Metal MoE sum6 HC token count is outside supported range\n"); + return 0; + } + + const uint64_t block_values = (uint64_t)n_embd; + const uint64_t hc_values = (uint64_t)n_hc * n_embd; + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + if (block_values > UINT64_MAX / (6u * sizeof(float)) || + n_tokens64 > UINT64_MAX / (6u * block_values * sizeof(float)) || + n_tokens64 > UINT64_MAX / (block_values * sizeof(float)) || + n_tokens64 > UINT64_MAX / (hc_values * sizeof(float)) || + n_tokens64 > UINT64_MAX / (mix_hc * sizeof(float))) { + fprintf(stderr, "ds4: Metal MoE sum6 HC activation size overflow\n"); + return 0; + } + + const uint64_t expert_bytes = + n_tokens64 * 6u * block_values * sizeof(float); + const uint64_t block_bytes = + n_tokens64 * block_values * sizeof(float); + const uint64_t hc_bytes = n_tokens64 * hc_values * sizeof(float); + const uint64_t split_bytes = n_tokens64 * mix_hc * sizeof(float); + if (!outbuf || !expertsbuf || !sharedbuf || !resbuf || !splitbuf || + ds4_gpu_tensor_bytes(expert_down) < expert_bytes || + ds4_gpu_tensor_bytes(shared_out) < block_bytes || + ds4_gpu_tensor_bytes(residual_hc) < hc_bytes || + ds4_gpu_tensor_bytes(split) < split_bytes) { + fprintf(stderr, + "ds4: Metal MoE sum6 HC received undersized activation buffers\n"); + return 0; + } + + ds4_gpu_hc_expand_args args = { + .n_embd = n_embd, + .n_hc = n_hc, + .n_tokens = (int64_t)n_tokens64, + .nb_block0 = sizeof(float), + .nb_block1 = (uint64_t)n_embd * sizeof(float), + .nb_add0 = sizeof(float), + .nb_add1 = (uint64_t)n_embd * sizeof(float), + .nb_res0 = sizeof(float), + .nb_res1 = (uint64_t)n_embd * sizeof(float), + .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float), + .nb_post0 = sizeof(float), + .nb_post1 = mix_hc * sizeof(float), + .nb_comb0 = sizeof(float), + .nb_comb1 = (uint64_t)n_hc * sizeof(float), + .nb_comb2 = mix_hc * sizeof(float), + .nb0 = sizeof(float), + .nb1 = (uint64_t)n_embd * sizeof(float), + .nb2 = (uint64_t)n_hc * n_embd * sizeof(float), + .has_add = 1, + }; + const uint64_t n_elem = (uint64_t)n_embd * n_tokens64; + NSUInteger nth = g_dsv4_moe_sum6_hc_expand4_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth > n_elem) nth = (NSUInteger)n_elem; + if (nth == 0u) return 0; + const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_moe_sum6_hc_expand4_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:expertsbuf offset:ds4_gpu_tensor_offset(expert_down) atIndex:1]; + [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2]; + [enc setBuffer:splitbuf + offset:ds4_gpu_tensor_offset(split) + (NSUInteger)n_hc * sizeof(float) + atIndex:3]; + [enc setBuffer:splitbuf + offset:ds4_gpu_tensor_offset(split) + (NSUInteger)(2u * n_hc) * sizeof(float) + atIndex:4]; + [enc setBuffer:sharedbuf offset:ds4_gpu_tensor_offset(shared_out) atIndex:5]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:6]; + [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "MoE sum6 HC expand"); + } +} + int ds4_gpu_hc_expand_add_split_half_add_tensor( ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, diff --git a/metal/dsv4_hc.metal b/metal/dsv4_hc.metal index c46794675..fab24b396 100644 --- a/metal/dsv4_hc.metal +++ b/metal/dsv4_hc.metal @@ -692,6 +692,233 @@ kernel void kernel_dsv4_hc_expand4( } } +// Batch attention-output tail specialization for the exact pre-M5 legacy +// Q8_0 matmul tile: +// +// attn_out = attn_low @ Wob +// after_attn_hc = HCPost(attn_out, residual_hc, split) +// +// The matrix body is the aligned Q8_0/F32 specialization of kernel_mul_mm, +// kept in the same load/dequantize/MMA order. Its 64x32 F32 result is staged +// through threadgroup memory before the scalar HC=4 epilogue, preserving the +// materialized F32 boundary and kernel_dsv4_hc_expand4's statement order while +// removing the global attn_out write/read and the standalone HC dispatch. +kernel void kernel_dsv4_attn_out_q8_mm_hc_expand4_batch( + constant ds4_metal_args_mul_mm & mm [[buffer(0)]], + device const char * weight [[buffer(1)]], + device const char * input [[buffer(2)]], + device const char * residual [[buffer(3)]], + device const char * post [[buffer(4)]], + device const char * comb [[buffer(5)]], + device char * dst [[buffer(6)]], + constant ds4_metal_args_dsv4_hc_expand & hc [[buffer(7)]], + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + constexpr int NR0 = 64; + constexpr int NR1 = 32; + constexpr int NK = 32; + constexpr int NL0 = NK / 16; + constexpr int NL1 = NK / 8; + + if (hc.n_hc != 4 || mm.ne0 != hc.n_embd || mm.ne1 != hc.n_tokens || + mm.ne00 != 8192 || mm.ne0 != 4096 || (mm.ne1 & 31) != 0) { + return; + } + + threadgroup half * sa = (threadgroup half *)shmem; + threadgroup half * sb = (threadgroup half *)(shmem + 4096); + + const int im = tgpig.z; + const int r0 = tgpig.y * NR0; + const int r1 = tgpig.x * NR1; + const short lr0 = (short)tiitg / NL0; + const short lr1 = (short)tiitg / NL1; + const short il0 = tiitg % NL0; + short il = il0; + + const int i12 = im % mm.ne12; + const int i13 = im / mm.ne12; + const uint64_t offset0 = + (i12 / mm.r2) * mm.nb02 + (i13 / mm.r3) * mm.nb03; + const short offset1 = il0 / 2; + device const block_q8_0 * x = + (device const block_q8_0 *)(weight + mm.nb01 * (r0 + lr0) + offset0) + + offset1; + + const short iy = 8 * (tiitg % NL1); + device const float * y = (device const float *)(input + + mm.nb13 * i13 + mm.nb12 * i12 + mm.nb11 * (r1 + lr1) + mm.nb10 * iy); + + simdgroup_half8x8 ma[4]; + simdgroup_half8x8 mb[2]; + simdgroup_float8x8 mc[8]; + FOR_UNROLL (short i = 0; i < 8; ++i) { + mc[i] = make_filled_simdgroup_matrix(0.0f); + } + + for (int loop_k = 0; loop_k < mm.ne00; loop_k += NK) { + half4x4 temp_a; + dequantize_q8_0(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; + *(sa + 64 * ib + 8 * ly + lx) = temp_a[i / 4][i % 4]; + } + + const short sx = tiitg % NL1; + const short sy = (tiitg / NL1) / 8; + const short ly = (tiitg / NL1) % 8; + const short ib = 4 * sx + sy; + *(threadgroup half2x4 *)(sb + 64 * ib + 8 * ly) = + (half2x4)(*((device float2x4 *)y)); + + il = (il + 2 < 2) ? il + 2 : il % 2; + x = (il < 2) ? x + 1 : x; + y += NK; + + threadgroup_barrier(mem_flags::mem_threadgroup); + + 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; + } + } + + // Every simdgroup must finish its final sa/sb reads before the same 8 KiB + // allocation is reused as the row-major F32 output tile. + threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup float * tile = (threadgroup float *)shmem; + threadgroup float * tile_sg = + tile + 32 * (sgitg & 1) + 16 * (sgitg >> 1) * NR0; + FOR_UNROLL (short i = 0; i < 8; ++i) { + simdgroup_store(mc[i], + tile_sg + 8 * (i % 4) + 8 * NR0 * (i / 4), + NR0, 0, false); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint e = tiitg; e < (uint)(NR0 * NR1); e += 128u) { + const int64_t t = r1 + e / NR0; + const int64_t d = r0 + e % NR0; + const float block_v = tile[e]; + + const float rv0 = *((device const float *)(residual + + d * hc.nb_res0 + 0 * hc.nb_res1 + t * hc.nb_res2)); + const float rv1 = *((device const float *)(residual + + d * hc.nb_res0 + 1 * hc.nb_res1 + t * hc.nb_res2)); + const float rv2 = *((device const float *)(residual + + d * hc.nb_res0 + 2 * hc.nb_res1 + t * hc.nb_res2)); + const float rv3 = *((device const float *)(residual + + d * hc.nb_res0 + 3 * hc.nb_res1 + t * hc.nb_res2)); + + for (int64_t dst_hc = 0; dst_hc < 4; ++dst_hc) { + float acc = block_v * *((device const float *)(post + + dst_hc * hc.nb_post0 + t * hc.nb_post1)); + acc += *((device const float *)(comb + + dst_hc * hc.nb_comb0 + 0 * hc.nb_comb1 + t * hc.nb_comb2)) * rv0; + acc += *((device const float *)(comb + + dst_hc * hc.nb_comb0 + 1 * hc.nb_comb1 + t * hc.nb_comb2)) * rv1; + acc += *((device const float *)(comb + + dst_hc * hc.nb_comb0 + 2 * hc.nb_comb1 + t * hc.nb_comb2)) * rv2; + acc += *((device const float *)(comb + + dst_hc * hc.nb_comb0 + 3 * hc.nb_comb1 + t * hc.nb_comb2)) * rv3; + *((device float *)(dst + d * hc.nb0 + dst_hc * hc.nb1 + + t * hc.nb2)) = acc; + } + } +} + +// Batch routed-FFN tail specialization: +// +// routed_out = sum(expert_down[0..5]) +// next_hc = HCPost(routed_out + shared_out, residual_hc, split) +// +// The scalar statements intentionally mirror kernel_dsv4_moe_sum6_f32 and +// kernel_dsv4_hc_expand4 in their original order. The intermediate routed +// row is dead on the admitted graph path, so consuming the six expert rows +// here removes its global write/read and the standalone sum dispatch. +kernel void kernel_dsv4_moe_sum6_hc_expand4( + constant ds4_metal_args_dsv4_hc_expand & args, + device const char * expert_down, + device const char * residual, + device const char * post, + device const char * comb, + device const char * shared_out, + device char * dst, + uint gid [[thread_position_in_grid]]) { + if (args.n_hc != 4) { + return; + } + + const int64_t n_elem = args.n_embd * args.n_tokens; + if ((int64_t) gid >= n_elem) { + return; + } + + const int64_t d = ((int64_t) gid) % args.n_embd; + const int64_t t = ((int64_t) gid) / args.n_embd; + device const float *s = (device const float *)(expert_down + + (uint64_t)t * 6u * (uint64_t)args.n_embd * sizeof(float)); + + float block_v = s[d]; + block_v += s[(uint64_t)args.n_embd + d]; + block_v += s[2u * (uint64_t)args.n_embd + d]; + block_v += s[3u * (uint64_t)args.n_embd + d]; + block_v += s[4u * (uint64_t)args.n_embd + d]; + block_v += s[5u * (uint64_t)args.n_embd + d]; + block_v += *((device const float *)(shared_out + + d*args.nb_add0 + t*args.nb_add1)); + + const float r0 = *((device const float *)(residual + + d*args.nb_res0 + 0*args.nb_res1 + t*args.nb_res2)); + const float r1 = *((device const float *)(residual + + d*args.nb_res0 + 1*args.nb_res1 + t*args.nb_res2)); + const float r2 = *((device const float *)(residual + + d*args.nb_res0 + 2*args.nb_res1 + t*args.nb_res2)); + const float r3 = *((device const float *)(residual + + d*args.nb_res0 + 3*args.nb_res1 + t*args.nb_res2)); + + for (int64_t dst_hc = 0; dst_hc < 4; ++dst_hc) { + float acc = block_v * *((device const float *)(post + + dst_hc*args.nb_post0 + t*args.nb_post1)); + + acc += *((device const float *)(comb + dst_hc*args.nb_comb0 + + 0*args.nb_comb1 + t*args.nb_comb2)) * r0; + acc += *((device const float *)(comb + dst_hc*args.nb_comb0 + + 1*args.nb_comb1 + t*args.nb_comb2)) * r1; + acc += *((device const float *)(comb + dst_hc*args.nb_comb0 + + 2*args.nb_comb1 + t*args.nb_comb2)) * r2; + acc += *((device const float *)(comb + dst_hc*args.nb_comb0 + + 3*args.nb_comb1 + t*args.nb_comb2)) * r3; + + *((device float *)(dst + d*args.nb0 + dst_hc*args.nb1 + + t*args.nb2)) = acc; + } +} + // Decode-time FFN tail fusion: // // shared_out = shared_mid @ Wshared_down diff --git a/metal/dsv4_misc.metal b/metal/dsv4_misc.metal index 5620b2443..c949838df 100644 --- a/metal/dsv4_misc.metal +++ b/metal/dsv4_misc.metal @@ -5936,6 +5936,270 @@ kernel void kernel_dsv4_indexed_mixed_attention_heads16_dual( } } +// Prefill specialization of the eight-head indexed-attention kernel. Stage +// four K/V rows per threadgroup barrier, then consume those rows in their +// original order. The online-softmax update sequence for each head is +// unchanged; only the row loads and barriers are batched. +kernel void kernel_dsv4_indexed_mixed_attention_heads8_rb4( + constant ds4_metal_args_dsv4_indexed_attention &args, + device const char *q, + device const char *raw_kv, + device const char *comp_kv, + device const char *topk, + device const char *sinks, + device char *dst, + threadgroup half4 *kv_shared [[threadgroup(0)]], + uint2 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + constexpr uint rows_per_block = 4u; + constexpr uint vecs_per_row = 128u; + + const uint token = tgpig.x; + const uint head = tgpig.y*8u + (uint)sg; + if (token >= args.n_tokens || head >= args.n_head) return; + + device const float4 *q4 = (device const float4 *)(q + + (uint64_t)token*args.q_token_stride + + (uint64_t)head*args.q_head_stride); + const half4 q0 = (half4)q4[lane + 0]; + const half4 q1 = (half4)q4[lane + 32]; + const half4 q2 = (half4)q4[lane + 64]; + const half4 q3 = (half4)q4[lane + 96]; + + float M = -FLT_MAX/2.0f; + float S = 0.0f; + float4 o0 = 0.0f; + float4 o1 = 0.0f; + float4 o2 = 0.0f; + float4 o3 = 0.0f; + + const uint qpos = args.pos0 + token; + const uint last_pos = args.pos0 + args.n_tokens - 1u; + const uint first_raw_pos = last_pos + 1u - args.n_raw; + const uint raw_last_pos = first_raw_pos + args.n_raw - 1u; + const uint window_first = (args.window != 0u && qpos + 1u > args.window) ? + qpos + 1u - args.window : 0u; + const uint first = max(first_raw_pos, window_first); + const uint last = min(qpos, raw_last_pos); + if (first <= last) { + for (uint pos0 = first; pos0 <= last; pos0 += rows_per_block) { + const uint n_rows = min(rows_per_block, last - pos0 + 1u); + for (uint off = (uint)tid; + off < n_rows * vecs_per_row; + off += 256u) { + const uint r = off / vecs_per_row; + const uint c = off - r * vecs_per_row; + const uint logical = pos0 + r - first_raw_pos; + const uint row = (args.raw_start + logical)%args.raw_cap; + device const float4 *src = (device const float4 *)(raw_kv + + (uint64_t)row*args.raw_row_stride); + kv_shared[off] = (half4)src[c]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint r = 0; r < n_rows; r++) { + dsv4_attend_shared_h4_row_at(kv_shared, r, + q0, q1, q2, q3, + args.scale, lane, M, S, o0, o1, o2, o3); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + + const uint visible = min((qpos + 1u)/args.ratio, args.n_comp); + device const int32_t *row_topk = (device const int32_t *)(topk + + (uint64_t)token*args.topk_token_stride); + bool stop = false; + for (uint i = 0; i < args.top_k && !stop; i += rows_per_block) { + uint rows[rows_per_block]; + uint n_rows = 0; + for (uint j = 0; j < rows_per_block && i + j < args.top_k; j++) { + const int32_t idx = row_topk[i + j]; + if (idx < 0) continue; + if ((uint)idx >= visible) { + stop = true; + break; + } + rows[n_rows++] = (uint)idx; + } + if (n_rows == 0) continue; + for (uint off = (uint)tid; + off < n_rows * vecs_per_row; + off += 256u) { + const uint r = off / vecs_per_row; + const uint c = off - r * vecs_per_row; + kv_shared[off] = dsv4_load_cache_h4(comp_kv, + args.comp_row_stride, + rows[r], + c, + args.comp_kv_f16 != 0u); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint r = 0; r < n_rows; r++) { + dsv4_attend_shared_h4_row_at(kv_shared, r, + q0, q1, q2, q3, + args.scale, lane, M, S, o0, o1, o2, o3); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + dsv4_attend_sink(((device const float *)sinks)[head], + M, S, o0, o1, o2, o3); + const float inv_s = S == 0.0f ? 0.0f : 1.0f/S; + device float4 *dst4 = (device float4 *)(dst + + (uint64_t)token*args.dst_token_stride + + (uint64_t)head*args.dst_head_stride); + dst4[lane + 0] = o0*inv_s; dst4[lane + 32] = o1*inv_s; + dst4[lane + 64] = o2*inv_s; dst4[lane + 96] = o3*inv_s; +} + +// Combine the existing two-head-per-simdgroup layout with four-row staging. +// Each head retains the RB4 row order and online-softmax update sequence. +kernel void kernel_dsv4_indexed_mixed_attention_heads16_dual_rb4( + constant ds4_metal_args_dsv4_indexed_attention &args, + device const char *q, + device const char *raw_kv, + device const char *comp_kv, + device const char *topk, + device const char *sinks, + device char *dst, + threadgroup half4 *kv_shared [[threadgroup(0)]], + uint2 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + constexpr uint rows_per_block = 4u; + constexpr uint vecs_per_row = 128u; + + const uint token = tgpig.x; + const uint head0 = tgpig.y*16u + (uint)sg; + const uint head1 = head0 + 8u; + if (token >= args.n_tokens || head0 >= args.n_head) return; + + device const float4 *qa = (device const float4 *)(q + + (uint64_t)token*args.q_token_stride + + (uint64_t)head0*args.q_head_stride); + const half4 qa0 = (half4)qa[lane + 0]; + const half4 qa1 = (half4)qa[lane + 32]; + const half4 qa2 = (half4)qa[lane + 64]; + const half4 qa3 = (half4)qa[lane + 96]; + half4 qb0 = half4(0.0h), qb1 = half4(0.0h); + half4 qb2 = half4(0.0h), qb3 = half4(0.0h); + if (head1 < args.n_head) { + device const float4 *qb = (device const float4 *)(q + + (uint64_t)token*args.q_token_stride + + (uint64_t)head1*args.q_head_stride); + qb0 = (half4)qb[lane + 0]; + qb1 = (half4)qb[lane + 32]; + qb2 = (half4)qb[lane + 64]; + qb3 = (half4)qb[lane + 96]; + } + + float Ma = -FLT_MAX/2.0f, Sa = 0.0f; + float Mb = -FLT_MAX/2.0f, Sb = 0.0f; + float4 ao0 = 0.0f, ao1 = 0.0f, ao2 = 0.0f, ao3 = 0.0f; + float4 bo0 = 0.0f, bo1 = 0.0f, bo2 = 0.0f, bo3 = 0.0f; + + const uint qpos = args.pos0 + token; + const uint last_pos = args.pos0 + args.n_tokens - 1u; + const uint first_raw_pos = last_pos + 1u - args.n_raw; + const uint raw_last_pos = first_raw_pos + args.n_raw - 1u; + const uint window_first = (args.window != 0u && qpos + 1u > args.window) ? + qpos + 1u - args.window : 0u; + const uint first = max(first_raw_pos, window_first); + const uint last = min(qpos, raw_last_pos); + if (first <= last) { + for (uint pos0 = first; pos0 <= last; pos0 += rows_per_block) { + const uint n_rows = min(rows_per_block, last - pos0 + 1u); + for (uint off = (uint)tid; + off < n_rows * vecs_per_row; + off += 256u) { + const uint r = off / vecs_per_row; + const uint c = off - r * vecs_per_row; + const uint logical = pos0 + r - first_raw_pos; + const uint row = (args.raw_start + logical)%args.raw_cap; + device const float4 *src = (device const float4 *)(raw_kv + + (uint64_t)row*args.raw_row_stride); + kv_shared[off] = (half4)src[c]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint r = 0; r < n_rows; r++) { + dsv4_attend_shared_h4_row_at(kv_shared, r, + qa0, qa1, qa2, qa3, + args.scale, lane, Ma, Sa, ao0, ao1, ao2, ao3); + if (head1 < args.n_head) { + dsv4_attend_shared_h4_row_at(kv_shared, r, + qb0, qb1, qb2, qb3, + args.scale, lane, Mb, Sb, bo0, bo1, bo2, bo3); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + + const uint visible = min((qpos + 1u)/args.ratio, args.n_comp); + device const int32_t *row_topk = (device const int32_t *)(topk + + (uint64_t)token*args.topk_token_stride); + bool stop = false; + for (uint i = 0; i < args.top_k && !stop; i += rows_per_block) { + uint rows[rows_per_block]; + uint n_rows = 0; + for (uint j = 0; j < rows_per_block && i + j < args.top_k; j++) { + const int32_t idx = row_topk[i + j]; + if (idx < 0) continue; + if ((uint)idx >= visible) { + stop = true; + break; + } + rows[n_rows++] = (uint)idx; + } + if (n_rows == 0) continue; + for (uint off = (uint)tid; + off < n_rows * vecs_per_row; + off += 256u) { + const uint r = off / vecs_per_row; + const uint c = off - r * vecs_per_row; + kv_shared[off] = dsv4_load_cache_h4(comp_kv, + args.comp_row_stride, + rows[r], + c, + args.comp_kv_f16 != 0u); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint r = 0; r < n_rows; r++) { + dsv4_attend_shared_h4_row_at(kv_shared, r, + qa0, qa1, qa2, qa3, + args.scale, lane, Ma, Sa, ao0, ao1, ao2, ao3); + if (head1 < args.n_head) { + dsv4_attend_shared_h4_row_at(kv_shared, r, + qb0, qb1, qb2, qb3, + args.scale, lane, Mb, Sb, bo0, bo1, bo2, bo3); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + dsv4_attend_sink(((device const float *)sinks)[head0], + Ma, Sa, ao0, ao1, ao2, ao3); + const float ia = Sa == 0.0f ? 0.0f : 1.0f/Sa; + device float4 *da = (device float4 *)(dst + + (uint64_t)token*args.dst_token_stride + + (uint64_t)head0*args.dst_head_stride); + da[lane + 0] = ao0*ia; da[lane + 32] = ao1*ia; + da[lane + 64] = ao2*ia; da[lane + 96] = ao3*ia; + if (head1 < args.n_head) { + dsv4_attend_sink(((device const float *)sinks)[head1], + Mb, Sb, bo0, bo1, bo2, bo3); + const float ib = Sb == 0.0f ? 0.0f : 1.0f/Sb; + device float4 *db = (device float4 *)(dst + + (uint64_t)token*args.dst_token_stride + + (uint64_t)head1*args.dst_head_stride); + db[lane + 0] = bo0*ib; db[lane + 32] = bo1*ib; + db[lane + 64] = bo2*ib; db[lane + 96] = bo3*ib; + } +} + // Decode specialization of kernel_dsv4_indexed_mixed_attention_heads8. // Generation attends one token at a time, so the ratio-4 indexed path spends a // visible amount of time repeatedly staging the same K/V row for the eight diff --git a/metal/dsv4_rope.metal b/metal/dsv4_rope.metal index 279365c8f..81f57db0d 100644 --- a/metal/dsv4_rope.metal +++ b/metal/dsv4_rope.metal @@ -58,6 +58,11 @@ struct ds4_metal_args_dsv4_head_norm_rope { float beta_slow; }; +struct ds4_metal_args_dsv4_batch_qkv_finalize { + uint32_t raw_cap; + uint32_t raw_pos0; +}; + 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)); @@ -423,6 +428,101 @@ kernel void kernel_dsv4_head_rms_norm_rope_tail_f32( } } +// KV half of the batch finalizer as a separate PSO. The host can dispatch it +// concurrently with the unchanged Q-head norm/RoPE PSO, retaining that PSO's +// exact compiled arithmetic. The KV PSO collapses the four post-q_b KV +// dispatches; together the active path turns five serial dispatches into two +// concurrent dispatches. +kernel void kernel_dsv4_batch_kv_rope_fp8_store_f32( + constant ds4_metal_args_dsv4_rope_tail & args, + constant ds4_metal_args_dsv4_batch_qkv_finalize & store, + device const int32_t * positions, + device float * kvraw, + device float * raw_cache, + threadgroup float * scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg [[threads_per_threadgroup]]) { + const int i1 = tgpig.x; + const int i2 = tgpig.y; + const int n_nope = (int)args.ne00 - args.n_dims; + if (args.mode != 0 || n_nope < 0) { + return; + } + + device float *kv = (device float *)((device char *)kvraw + + (uint64_t)i2 * args.nb02 + (uint64_t)i1 * args.nb01); + + 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)positions[i2]; + const float inv_ndims = -1.f / args.n_dims; + + 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 + float cos_theta; + float sin_theta; + rope_yarn(theta, 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 = kv[j0]; + const float x1 = kv[j1]; + kv[j0] = x0 * cos_theta - x1 * sin_theta; + kv[j1] = x0 * sin_theta + x1 * cos_theta; + } + + threadgroup_barrier(mem_flags::mem_device_and_threadgroup); + + device float *raw = raw_cache + + (uint64_t)((store.raw_pos0 + (uint)i2) % store.raw_cap) * + (uint64_t)args.ne00; + for (int off = 0; off < n_nope; off += 64) { + float v = 0.0f; + if (off + (int)tid < n_nope) { + v = kv[off + tid]; + scratch[tid] = abs(v); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint stride = 32; stride > 0; stride >>= 1) { + if (tid < stride) { + scratch[tid] = max(scratch[tid], scratch[tid + stride]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + const float amax = max(scratch[0], 1.0e-4f); + const float fp8_scale = exp2(ceil(log2(amax / 448.0f))); + if (off + (int)tid < n_nope) { + const float q = dsv4_e4m3fn_dequant( + clamp(v / fp8_scale, -448.0f, 448.0f)) * fp8_scale; + kv[off + tid] = q; + raw[off + tid] = (float)((half)q); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + for (int i = n_nope + (int)tid; i < args.ne00; i += 64) { + raw[i] = (float)((half)kv[i]); + } +} + // 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 @@ -881,3 +981,64 @@ kernel void kernel_dsv4_flash_attn_vec_packed32_reduce_rope_f16_dk512_dv512( tiitg, 32u * 32u); } + +// Direct-KV sibling of the packed32 reduce+RoPE decode kernel above: reads +// the F32 raw ring and F16 compressed cache in place (see +// ds4_flash_attn_vec_packed8_reduce_f16_512_direct_kv in flash_attn.metal) +// so the gathered KV staging dispatch is skipped entirely. +kernel void kernel_dsv4_flash_attn_vec_packed32_reduce_rope_f16_dk512_dv512_direct_kv( + constant ds4_metal_args_flash_attn_ext_vec & args [[buffer(0)]], + device const char * q [[buffer(1)]], + device const char * raw_kv [[buffer(2)]], + device const char * comp_kv [[buffer(3)]], + device const char * mask [[buffer(4)]], + device const char * sinks [[buffer(5)]], + constant ds4_metal_args_flash_kv_direct & kv [[buffer(6)]], + device char * dst [[buffer(7)]], + constant ds4_metal_args_dsv4_rope_affine_pair & rope + [[buffer(8)]], + threadgroup char * shmem [[threadgroup(0)]], + uint head [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + /* Same uniform specialization guard as the staged packed32 kernel; nb11 + * / nb21 remain the F16 compressed row stride (raw rows are a fixed + * 2048-byte F32 stride inside the helper). */ + if (!FC_flash_attn_ext_vec_has_mask || + !FC_flash_attn_ext_vec_has_sinks || + FC_flash_attn_ext_vec_has_bias || + FC_flash_attn_ext_vec_has_scap || + FC_flash_attn_ext_vec_nsg != 1 || + FC_flash_attn_ext_vec_nwg != 32 || + FC_flash_attn_ext_vec_ns10 != 512 || + FC_flash_attn_ext_vec_ns20 != 512 || + args.ne01 != 1 || args.ne02 != 64 || args.ne03 != 1 || + args.ne_12_2 != 1 || args.ne_12_3 != 1 || + args.ne31 != 1 || args.ne32 != 1 || args.ne33 != 1 || + args.ne11 <= 0 || args.ne11 > 1024 || head >= (uint)args.ne02 || + args.nb02 != 2048 || args.nb11 != 1024 || args.nb21 != 1024 || + kv.n_raw == 0 || kv.n_raw > (uint)args.ne11 || + kv.raw_cap < kv.n_raw || kv.raw_start >= kv.raw_cap || + rope.head_dim != 512 || rope.n_dims != 64 || + rope.row_bytes != 2048 || rope.inverse == 0) { + return; + } + + ds4_flash_attn_vec_packed8_reduce_f16_512_direct_kv( + args, q, raw_kv, comp_kv, mask, sinks, kv, dst, shmem, + head, tiisg, sgitg); + + /* Same producer/consumer boundary as the current reduce+RoPE kernel. */ + threadgroup_barrier(mem_flags::mem_device); + + const int n_nope = rope.head_dim - rope.n_dims; + device char * row = dst + (uint64_t)head * rope.row_bytes; + ds4_rope_tail_pair_affine_row(rope, + (device const char *)row, + row, + n_nope, + rope.pos0, + tiitg, + 32u * 32u); +} diff --git a/metal/flash_attn.metal b/metal/flash_attn.metal index e3232a72d..cc0cfcc46 100644 --- a/metal/flash_attn.metal +++ b/metal/flash_attn.metal @@ -1690,3 +1690,320 @@ static inline void ds4_flash_attn_vec_packed8_reduce_f16_512( threadgroup_barrier(mem_flags::mem_threadgroup); } } + + +struct ds4_metal_args_flash_kv_direct { + uint raw_cap; + uint raw_start; + uint n_raw; + uint pad0; +}; + +/* Direct-read sibling of ds4_flash_attn_vec_packed8_reduce_f16_512 for + * gathered decode: skips the contiguous F16 K/V staging dispatch entirely + * and computes each row's source inline. Logical rows below n_raw read the + * F32 raw ring (single conditional wrap subtract, then float4 -> half4 -> + * float4 exactly like kernel_dsv4_flash_kv_stage_f16 followed by the F16 + * consumer, so the half bits are identical), rows at/above n_raw read the + * F16 compressed cache in place (same bits as the ushort4 transport), and + * rows past ne11 take an in-place zero row with a -MAXHALF mask, matching + * the fused pad writes bit-for-bit. Invalid rows still execute their dot + * and lo += 0.0h * weight accumulations (never skipped), the + * simd_max(sm[lane]) > -MAXHALF block-skip gate is unchanged, and the whole + * reduction topology (simd_sum trees, 8 simdgroups x 32 virtual splits, + * 33-column partial plane) is untouched. */ +static inline void ds4_flash_attn_vec_packed8_reduce_f16_512_direct_kv( + constant ds4_metal_args_flash_attn_ext_vec & args, + device const char * q, + device const char * raw_kv, + device const char * comp_kv, + device const char * mask, + device const char * sinks, + constant ds4_metal_args_flash_kv_direct & kv, + device char * dst, + threadgroup char * shmem, + uint head, + ushort tiisg, + ushort sgitg) { + constexpr short NW = 32; + constexpr short C = 32; + constexpr short NSG = 8; + constexpr short NWG = 32; + constexpr short DK4 = 128; + constexpr short DV4 = 128; + constexpr short SH = 128; + + /* Same 24,448-byte dynamic layout as the staged variant. */ + threadgroup half4 *q_shared = (threadgroup half4 *)shmem; + threadgroup half *score_banks = + (threadgroup half *)(q_shared + DK4); + threadgroup volatile float *weights = + (threadgroup volatile float *)(score_banks + NSG * SH); + threadgroup volatile float *stats = weights + NWG * C; + threadgroup volatile float *sink_scale = stats + 2 * NWG; + threadgroup volatile float4 *partial_plane = + (threadgroup volatile float4 *)(sink_scale + NWG); + + const short lane = (short)tiisg; + threadgroup half *bank = score_banks + (short)sgitg * SH; + threadgroup float *ss = (threadgroup float *)bank; + threadgroup half *sm = bank + 2 * C; + + device const float4 *q4 = + (device const float4 *)(q + (uint64_t)head * args.nb02); + if (sgitg == 0) { + for (short i = lane; i < DK4; i += NW) { + q_shared[i] = (half4)q4[i]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (short iwg = (short)sgitg; iwg < NWG; iwg += NSG) { + float S = 0.0f; + float M = -FLT_MAX / 2; + float out_scale = 1.0f; + const int ic_original = (int)iwg * C; + + weights[(uint)iwg * C + (uint)lane] = 0.0f; + ss[lane] = 0.0f; + sm[lane] = (half)0.0h; + simdgroup_barrier(mem_flags::mem_threadgroup); + + if (ic_original < args.ne11) { + device const half *pm = (device const half *)mask; + const int ic = ic_original; + + /* In-place tail: rows past ne11 take the pad mask bits 0xfbff + * (-MAXHALF), exactly as the fused pad dispatch wrote them. */ + const uint mrow = (uint)ic + (uint)lane; + sm[lane] = mrow < (uint)args.ne11 + ? pm[mrow] + : as_type((ushort)0xfbffu); + if (simd_max(sm[lane]) > -MAXHALF) { + threadgroup const half4 *pq4 = q_shared; + pq4 += lane; + + /* Per-block fast paths: a 32-row block almost never mixes + * sources, so only the raw/comp-straddling, ring-wrapping or + * tail block pays per-row addressing. The fast loops keep + * the staged kernel's exact load walk; the raw path inserts + * the same float4 -> half4 -> float4 conversion the staging + * dispatch + F16 consumer performed. */ + const uint blk_first = (uint)ic; + const uint blk_last = (uint)ic + (uint)(C - 1); + const bool all_valid = (uint)ic + (uint)C <= (uint)args.ne11; + const bool all_raw = blk_last < kv.n_raw; + const bool raw_no_wrap = + kv.raw_start + blk_last < kv.raw_cap; + + float lane_mqk = 0.0f; + if (all_valid && all_raw && raw_no_wrap) { + device const float4 *pk4f = + (device const float4 *)(raw_kv + + (uint64_t)(kv.raw_start + blk_first) * 2048u); + pk4f += lane; + FOR_UNROLL (short cc = 0; cc < C; ++cc) { + float mqk = 0.0f; + FOR_UNROLL (short ii = 0; ii < DK4 / NW; ++ii) { + mqk += dot((float4)(half4)pk4f[cc * DK4 + ii * NW], + (float4)pq4[ii * NW]); + } + mqk = simd_sum(mqk); + if (lane == cc) { + lane_mqk = mqk; + } + } + } else if (all_valid && blk_first >= kv.n_raw) { + device const half4 *pk4 = + (device const half4 *)(comp_kv + + (uint64_t)(blk_first - kv.n_raw) * 1024u); + pk4 += lane; + FOR_UNROLL (short cc = 0; cc < C; ++cc) { + float mqk = 0.0f; + FOR_UNROLL (short ii = 0; ii < DK4 / NW; ++ii) { + mqk += dot((float4)pk4[cc * DK4 + ii * NW], + (float4)pq4[ii * NW]); + } + mqk = simd_sum(mqk); + if (lane == cc) { + lane_mqk = mqk; + } + } + } else { + FOR_UNROLL (short cc = 0; cc < C; ++cc) { + /* Per-row base: a packed 32-row block can straddle + * the raw/comp boundary and the ring wrap. */ + const uint r = (uint)ic + (uint)cc; + float mqk = 0.0f; + if (r < (uint)args.ne11 && r < kv.n_raw) { + uint phys = kv.raw_start + r; + if (phys >= kv.raw_cap) phys -= kv.raw_cap; + device const float4 *row4 = + (device const float4 *)(raw_kv + + (uint64_t)phys * 2048u); + FOR_UNROLL (short ii = 0; ii < DK4 / NW; ++ii) { + mqk += dot((float4)(half4)row4[ii * NW + lane], + (float4)pq4[ii * NW]); + } + } else if (r < (uint)args.ne11) { + device const half4 *row4 = + (device const half4 *)(comp_kv + + (uint64_t)(r - kv.n_raw) * 1024u); + FOR_UNROLL (short ii = 0; ii < DK4 / NW; ++ii) { + mqk += dot((float4)row4[ii * NW + lane], + (float4)pq4[ii * NW]); + } + } else { + /* Pad rows were zero; execute the same dot so + * the signed-zero products match the staged + * path. */ + FOR_UNROLL (short ii = 0; ii < DK4 / NW; ++ii) { + mqk += dot(float4(0.0f), (float4)pq4[ii * NW]); + } + } + mqk = simd_sum(mqk); + if (lane == cc) { + lane_mqk = mqk; + } + } + } + + ss[lane] = fma(lane_mqk, args.scale, + (float)sm[lane]); + simdgroup_barrier(mem_flags::mem_threadgroup); + + const float old_m = M; + const float score = ss[lane]; + M = simd_max(max(M, score)); + const float ms = exp(old_m - M); + const float vs = exp(score - M); + S = S * ms + simd_sum(vs); + ss[lane] = vs; + simdgroup_barrier(mem_flags::mem_threadgroup); + + weights[(uint)iwg * C + (uint)lane] = ss[lane]; + } + + if (FC_flash_attn_ext_vec_has_sinks && iwg == 0) { + const float old_m = M; + const float sink = lane == 0 + ? ((device const float *)sinks)[head] + : -FLT_MAX / 2; + M = simd_max(max(M, sink)); + const float ms = exp(old_m - M); + const float vs = exp(sink - M); + S = S * ms + simd_sum(vs); + out_scale = ms; + } + } else if (FC_flash_attn_ext_vec_has_sinks && iwg == 0) { + const float old_m = M; + const float sink = lane == 0 + ? ((device const float *)sinks)[head] + : -FLT_MAX / 2; + M = simd_max(max(M, sink)); + const float ms = exp(old_m - M); + const float vs = exp(sink - M); + S = S * ms + simd_sum(vs); + out_scale = ms; + } + + if (lane == 0) { + stats[2 * (uint)iwg + 0] = S; + stats[2 * (uint)iwg + 1] = M; + sink_scale[(uint)iwg] = out_scale; + } + simdgroup_barrier(mem_flags::mem_threadgroup); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + const short split = lane; + float reduce_S = stats[2 * (uint)split + 0]; + float reduce_M = stats[2 * (uint)split + 1]; + const float reduce_max = simd_max(reduce_M); + const float reduce_ms = exp(reduce_M - reduce_max); + reduce_S = simd_sum(reduce_S * reduce_ms); + const float reduce_inv = + reduce_S == 0.0f ? 0.0f : 1.0f / reduce_S; + + device float4 *dst4 = + (device float4 *)(dst + + (uint64_t)head * 512u * sizeof(float)); + + for (short quadrant = 0; quadrant < 4; ++quadrant) { + for (short iwg = (short)sgitg; iwg < NWG; iwg += NSG) { + float4 lo = float4(0.0f); + const int ic_original = (int)iwg * C; + if (ic_original < args.ne11) { + const int ic = ic_original; + + threadgroup volatile float *split_weights = + weights + (uint)iwg * C; + const short oc = quadrant * NW + lane; + /* Same per-block fast paths as the QK pass. */ + const uint blk_first = (uint)ic; + const uint blk_last = (uint)ic + (uint)(C - 1); + const bool all_valid = (uint)ic + (uint)C <= (uint)args.ne11; + const bool all_raw = blk_last < kv.n_raw; + const bool raw_no_wrap = + kv.raw_start + blk_last < kv.raw_cap; + if (all_valid && all_raw && raw_no_wrap) { + device const float4 *pv4f = + (device const float4 *)(raw_kv + + (uint64_t)(kv.raw_start + blk_first) * 2048u); + FOR_UNROLL (short cc = 0; cc < C; ++cc) { + lo += (float4)(half4)pv4f[cc * DV4 + oc] * + float4(split_weights[cc]); + } + } else if (all_valid && blk_first >= kv.n_raw) { + device const half4 *pv4 = + (device const half4 *)(comp_kv + + (uint64_t)(blk_first - kv.n_raw) * 1024u); + FOR_UNROLL (short cc = 0; cc < C; ++cc) { + lo += float4(pv4[cc * DV4 + oc]) * + float4(split_weights[cc]); + } + } else { + FOR_UNROLL (short cc = 0; cc < C; ++cc) { + /* Same per-row source selection as the QK pass; invalid + * rows contribute +0.0f * weight, never a skipped + * iteration. */ + const uint r = (uint)ic + (uint)cc; + float4 vv = float4(0.0f); + if (r < (uint)args.ne11 && r < kv.n_raw) { + uint phys = kv.raw_start + r; + if (phys >= kv.raw_cap) phys -= kv.raw_cap; + vv = (float4)(half4)((device const float4 *)(raw_kv + + (uint64_t)phys * 2048u))[oc]; + } else if (r < (uint)args.ne11) { + vv = (float4)((device const half4 *)(comp_kv + + (uint64_t)(r - kv.n_raw) * 1024u))[oc]; + } + lo += vv * float4(split_weights[cc]); + } + } + + float4 acc = float4(0.0f); + acc += lo; + if (iwg == 0) { + acc *= sink_scale[0]; + } + lo = acc; + } + partial_plane[(uint)iwg * 33u + (uint)lane] = lo; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (short out_lane = (short)sgitg; out_lane < NW; out_lane += NSG) { + const float4 materialized = + (float4)partial_plane[(uint)lane * 33u + (uint)out_lane]; + const float4 reduced = simd_sum(materialized * reduce_ms); + if (lane == 0) { + dst4[quadrant * NW + out_lane] = reduced * reduce_inv; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} diff --git a/metal/moe.metal b/metal/moe.metal index 7aeb9d922..9ec961fa9 100644 --- a/metal/moe.metal +++ b/metal/moe.metal @@ -3444,6 +3444,98 @@ kernel void kernel_dsv4_attn_out_low_q8_0_f32( sgitg); } +/* Exact one-token DeepSeek V4 Flash decode shape for the attention-output low + * projection. The generic direct wrapper above has to recover the group and + * all tensor strides from runtime arguments, then walks a runtime 4096-wide K + * loop. Decode always uses eight 4096 -> 1024 Q8_0 projections here. Keep + * the established NR0=2 / NSG=4 lane mapping and reduction tree, but make the + * group offsets, row stride and four K trips literals so the compiler can + * eliminate the generic index arithmetic. */ +#define DS4_ATTN_OUT_LOW_Q8_STATIC_GROUPS 8 +#define DS4_ATTN_OUT_LOW_Q8_STATIC_K 4096 +#define DS4_ATTN_OUT_LOW_Q8_STATIC_ROWS 1024 +#define DS4_ATTN_OUT_LOW_Q8_STATIC_BLOCKS 128 +#define DS4_ATTN_OUT_LOW_Q8_STATIC_ROW_BYTES 4352 +#define DS4_ATTN_OUT_LOW_Q8_STATIC_GROUP_BYTES 4456448 + +static inline void ds4_attn_out_low_q8_static_impl( + device const char * src0s, + device const char * src1, + device char * dst, + threadgroup char * shmem, + uint3 tgpig, + ushort tiisg, + ushort sgitg) { + constexpr short NR0 = N_R0_Q8_0; + constexpr short NW = N_SIMDWIDTH; + constexpr short NQ = 8; + constexpr short NSG = 4; + + const uint group = tgpig.z; + const int r0 = (int)tgpig.x * NR0; + device const char *src0 = src0s + + (uint64_t)group * DS4_ATTN_OUT_LOW_Q8_STATIC_GROUP_BYTES; + device const float *y = (device const float *)(src1 + + (uint64_t)group * DS4_ATTN_OUT_LOW_Q8_STATIC_K * sizeof(float)); + device float *out = (device float *)(dst + + (uint64_t)group * DS4_ATTN_OUT_LOW_Q8_STATIC_ROWS * sizeof(float)); + + device const block_q8_0 *ax[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + ax[row] = (device const block_q8_0 *)(src0 + + (uint64_t)(r0 + row) * DS4_ATTN_OUT_LOW_Q8_STATIC_ROW_BYTES); + } + + float sumf[NR0] = { 0.0f }; + const short ix = tiisg / (NW / NQ); + const short il = tiisg % (NW / NQ); + const int ib0 = sgitg * NQ + ix; + device const float *yb = y + ib0 * QK8_0 + il * NQ; + float yl[NQ]; + + /* Every lane visits ib0 + {0, 32, 64, 96}, exactly as the generic loop. */ + for (int ib = ib0; ib < DS4_ATTN_OUT_LOW_Q8_STATIC_BLOCKS; + ib += NSG * NQ) { + for (short i = 0; i < NQ; ++i) { + yl[i] = yb[i]; + } + + for (short row = 0; row < NR0; ++row) { + device const int8_t *qs = ax[row][ib].qs + il * NQ; + float sumq = 0.0f; + FOR_UNROLL (short i = 0; i < NQ; ++i) { + sumq += qs[i] * yl[i]; + } + sumf[row] += sumq * ax[row][ib].d; + } + yb += NSG * NQ * QK8_0; + } + + helper_mv_reduce_and_write( + out, sumf, r0, DS4_ATTN_OUT_LOW_Q8_STATIC_ROWS, + tiisg, sgitg, shmem); +} + +kernel void kernel_dsv4_attn_out_low_q8_0_flash_decode_static_f32( + constant ds4_metal_args_mul_mv_id & args, + device const char * src0s, + device const char * src1, + device char * dst, + 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]]) { + /* This PSO is private to the exact host gate in + * ds4_gpu_attention_output_low_q8_tensor. Avoid repeating that proof in + * every shader lane: the point of this specialization is to remove the + * dynamic shape arithmetic from this very large dispatch. */ + ds4_attn_out_low_q8_static_impl( + src0s, src1, dst, shmem, tgpig, tiisg, sgitg); + (void)args; + (void)tiitg; +} + kernel void kernel_dsv4_attn_out_low_q4_K_f32( constant ds4_metal_args_mul_mv_id & args, device const char * src0s, @@ -6608,6 +6700,134 @@ kernel void kernel_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_f32( (void)tiitg; } +/* Layout mirror of ds4_metal_args_dsv4_hc_expand (metal/dsv4_hc.metal); the + * files are concatenated into one library with moe.metal first, so the + * struct cannot be shared. Only the decode t=0 fields are read here. */ +struct ds4_metal_args_moe_down_hc4 { + int64_t n_embd; + int64_t n_hc; + int64_t n_tokens; + uint64_t nb_block0; + uint64_t nb_block1; + uint64_t nb_add0; + uint64_t nb_add1; + uint64_t nb_res0; + uint64_t nb_res1; + uint64_t nb_res2; + uint64_t nb_post0; + uint64_t nb_post1; + uint64_t nb_comb0; + uint64_t nb_comb1; + uint64_t nb_comb2; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + int32_t has_add; +}; + +/* Decode FFN tail fusion on top of the fixed-route full-rows static sum6: + * + * after_ffn_hc = HCPost(down_sum6 + shared_out, residual_hc, split) + * + * The host reorders the layer tail so the plain shared-down Q8_0 matvec + * materializes shared_out before this dispatch. The matvec body, the + * routed_out store and the per-row simd_sum are byte-identical to + * kernel_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_f32; the epilogue + * then mirrors kernel_dsv4_shared_down_hc_expand4_q8_0 statement-for- + * statement (block_v = routed value; block_v += shared_out[d]; per dst_hc + * 0..3: acc = block_v*post, then comb_k*r_k in k order), so the result is + * bit-identical to the unfused dispatch pair. */ +kernel void kernel_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_hc4_f32( + constant ds4_metal_args_mul_mv_id &args, + constant ds4_metal_args_moe_down_hc4 &hc, + device const char *src0s, + device const char *src1, + device char *dst, + device const char *ids, + device const char *add_in, + device const char *shared_out, + device const char *residual, + device const char *post, + device const char *comb, + device char *hc_dst, + 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]]) { + if (hc.n_hc != 4 || hc.n_tokens != 1) { + return; + } + const short NSG = FC_mul_mv_nsg; + const uint32_t first_row = (uint32_t)((tgpig.x * NSG + sgitg) * N_R0_MXFP4); + device const int32_t *token_ids = (device const int32_t *)ids; + device const char *token_src1 = src1; + threadgroup float *lut = (threadgroup float *)shmem; + if (sgitg == 0) lut[tiisg] = ds4_metal_mxfp4_values[tiisg & 15]; + threadgroup_barrier(mem_flags::mem_threadgroup); + + /* Same defensive fallback as the base kernel: unreachable in production + * because the host proves the shape before selecting this pipeline. */ + const bool static_shape = + args.ne00 == DS4_MXFP4_DOWN_STATIC_NB * QK_MXFP4 && + args.nb01 == (uint64_t)DS4_MXFP4_DOWN_STATIC_ROW_BLOCKS * + sizeof(block_mxfp4) && + args.nei0 == DS4_MXFP4_DOWN_STATIC_SLOTS; + + float2 sumf = 0.0f; + if (static_shape) { + for (short slot = 0; slot < DS4_MXFP4_DOWN_STATIC_SLOTS; slot++) { + const int32_t expert = token_ids[slot]; + device const char *expert_base = + src0s + (int64_t)expert * args.nb02; + device const float *y = + (device const float *)(token_src1 + (uint64_t)slot * args.nb11); + sumf += ds4_mxfp4_accumulate_full_rows_static( + expert_base, y, first_row, lut, tiisg); + } + } else { + for (int slot = 0; slot < args.nei0; slot++) { + const int32_t expert = token_ids[slot]; + device const char *expert_base = + src0s + (int64_t)expert * args.nb02; + device const float *y = + (device const float *)(token_src1 + (uint64_t)slot * args.nb11); + sumf += ds4_mxfp4_accumulate_full_rows( + expert_base, args.nb01, y, args.ne00, first_row, lut, tiisg); + } + } + + device float *out = (device float *)dst; + FOR_UNROLL (short row = 0; row < N_R0_MXFP4; row++) { + const float value = simd_sum(sumf[row]); + if (tiisg == 0) { + const uint32_t d = first_row + row; + out[d] = value + + (args.tp_addend ? ((device const float *)add_in)[d] : 0.0f); + + float block_v = out[d]; + block_v += *((device const float *)(shared_out + (uint64_t)d * sizeof(float))); + + const float r0 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 0 * hc.nb_res1)); + const float r1 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 1 * hc.nb_res1)); + const float r2 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 2 * hc.nb_res1)); + const float r3 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 3 * hc.nb_res1)); + + for (int64_t dst_hc = 0; dst_hc < 4; ++dst_hc) { + float acc = block_v * *((device const float *)(post + dst_hc * hc.nb_post0)); + + acc += *((device const float *)(comb + dst_hc * hc.nb_comb0 + 0 * hc.nb_comb1)) * r0; + acc += *((device const float *)(comb + dst_hc * hc.nb_comb0 + 1 * hc.nb_comb1)) * r1; + acc += *((device const float *)(comb + dst_hc * hc.nb_comb0 + 2 * hc.nb_comb1)) * r2; + acc += *((device const float *)(comb + dst_hc * hc.nb_comb0 + 3 * hc.nb_comb1)) * r3; + + *((device float *)(hc_dst + (uint64_t)d * hc.nb0 + dst_hc * hc.nb1)) = acc; + } + } + } + (void)tiitg; +} + kernel void kernel_mul_mv_slots6_mxfp4_sum6_f32( constant ds4_metal_args_mul_mv_id &args, device const char *src00, device const char *src01, diff --git a/rocm/ds4_rocm_glm.cuh b/rocm/ds4_rocm_glm.cuh index e0bface99..f9692a0d1 100644 --- a/rocm/ds4_rocm_glm.cuh +++ b/rocm/ds4_rocm_glm.cuh @@ -5197,7 +5197,8 @@ static int glm_rocm_routed_moe_wrap( expert_mid_dim, out_dim, selected, weights, n_total_expert, n_expert, swiglu_clamp, x, NULL, layer_index, - force_resident); + force_resident, + NULL, NULL, NULL, NULL, NULL); } return ds4_gpu_routed_moe_batch_tensor(out, &gate_tmp, &up_tmp, mid, &down_tmp, model_map, model_size, gate_offset, diff --git a/rocm/ds4_rocm_moe_launch.cuh b/rocm/ds4_rocm_moe_launch.cuh index 16ad7bcb4..40b72ce94 100644 --- a/rocm/ds4_rocm_moe_launch.cuh +++ b/rocm/ds4_rocm_moe_launch.cuh @@ -2804,7 +2804,13 @@ static int routed_moe_launch( return ok; } -extern "C" int ds4_gpu_routed_moe_one_tensor(ds4_gpu_tensor *out, ds4_gpu_tensor *gate, ds4_gpu_tensor *up, ds4_gpu_tensor *mid, ds4_gpu_tensor *down, const void *model_map, uint64_t model_size, uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset, uint32_t gate_type, uint32_t down_type, uint64_t gate_expert_bytes, uint64_t gate_row_bytes, uint64_t down_expert_bytes, uint64_t down_row_bytes, uint32_t expert_in_dim, uint32_t expert_mid_dim, uint32_t out_dim, const ds4_gpu_tensor *selected, const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, float clamp, const ds4_gpu_tensor *x, const ds4_gpu_tensor *add_in, uint32_t layer_index, bool force_resident) { +extern "C" int ds4_gpu_routed_moe_one_tensor(ds4_gpu_tensor *out, ds4_gpu_tensor *gate, ds4_gpu_tensor *up, ds4_gpu_tensor *mid, ds4_gpu_tensor *down, const void *model_map, uint64_t model_size, uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset, uint32_t gate_type, uint32_t down_type, uint64_t gate_expert_bytes, uint64_t gate_row_bytes, uint64_t down_expert_bytes, uint64_t down_row_bytes, uint32_t expert_in_dim, uint32_t expert_mid_dim, uint32_t out_dim, const ds4_gpu_tensor *selected, const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, float clamp, const ds4_gpu_tensor *x, const ds4_gpu_tensor *add_in, uint32_t layer_index, bool force_resident, const ds4_gpu_tensor *hc_shared_out, const ds4_gpu_tensor *hc_residual, const ds4_gpu_tensor *hc_split, ds4_gpu_tensor *hc_out, int *hc_fused_out) { + (void)hc_shared_out; (void)hc_residual; (void)hc_split; + if (hc_fused_out) *hc_fused_out = 0; + if (hc_out) { + fprintf(stderr, "ds4: routed MoE HC tail fusion is Metal-only\n"); + return 0; + } if (add_in) { fprintf(stderr, "ds4: routed MoE addend fold is Metal-only\n"); return 0; diff --git a/speed-bench/DECODE-CAMPAIGN-HANDOFF.md b/speed-bench/DECODE-CAMPAIGN-HANDOFF.md new file mode 100644 index 000000000..88c040385 --- /dev/null +++ b/speed-bench/DECODE-CAMPAIGN-HANDOFF.md @@ -0,0 +1,117 @@ +# Decode campaign — handoff (session of Aug 21, M3 Ultra, round 2) + +Goal: push greedy decode past **45 tokens/s** on the MXFP4 `ds4flash.gguf` +model, bit-exact. **Reached: 45.5 t/s interleaved (45.47/45.55), all +transcripts md5 `db0c504c…`, `make test` 44/44, harness bit-identical.** +Round-1 state was 43.2–44.0 t/s; round 2 added +2.7% (attention parity) and ++1.75% (greedy chain decode), both validated per the campaign protocol. + +## Round-2 changes (working tree; commit state at bottom) + +1. **Raw-layer gathered attention** — `n_comp == 0` decode layers (ratio 0/128) + now use the gathered path (fused staging + packed32 reduce + fused inverse + RoPE) instead of the five-dispatch raw path. Found via commit-only stage + counters: the raw path cost ~65 µs vs ~32 µs per layer, invisible in the + averaged ledger ("attention core 44.3 µs flat"). ~0.7 ms/token at short + context. Rollback envs: `DS4_METAL_DISABLE_DECODE_RAW_GATHERED_ATTN`, + `DS4_METAL_DISABLE_DECODE_RAW_PACKED32`. Also fixed the staging wrapper's + `n_comp == 0` guard in `ds4_gpu_encode_flash_kv_stage_f16`. +2. **Greedy chain decode** — `metal_graph_greedy_chain` (ds4.c, engaged from + `generate_metal_graph_raw_swa` only): GPU argmax writes the next token id + into a device ring; the next token's embedding gathers it from the ring + (the by-value embed and the batched embed use the same get_rows/repeat + kernels — bit-identical). Encode runs two tokens ahead of the host's + confirm cursor; the host only lags (one MTLSharedEvent wait per token) to + print and stop-check. Removes the per-token `waitUntilCompleted` + 517 KiB + logits readback + CPU argmax boundary (~0.5 ms/token). Kill switch: + `DS4_DISABLE_GREEDY_CHAIN=1`. Diagnostics: `DS4_GREEDY_CHAIN_DEBUG`, + `DS4_GREEDY_CHAIN_DUMP_IDS`, `DS4_GREEDY_CHAIN_VERIFY`. + **Hash-layer gotcha that cost an hour**: the first `DS4_N_HASH_LAYER` + layers route experts by token id (`ffn_gate_tid2eid`); the select kernel + already supports a device-resident token (`use_token_buffer` in + `kernel_dsv4_router_finalize_one`), plumbed via + `ds4_gpu_router_select_tensor_devtoken` + `g->chain_token_view`. The + host-side `metal_graph_decode_set_hash_selected_override` is skipped in + chain mode — the resident fixed-route MoE never consumes it (the override/ + readback blocks live under `use_selected_slots`, all gated on + `g_ssd_streaming_mode`). Feeding token=0 instead was the one divergence: + deterministic drift from layer 0, tipping an argmax 33 tokens in. + Symptom signature if it regresses: transcripts match for N tokens then + diverge deterministically. + +## Verified numbers (interleaved CLI, `-c 8192 -n 128 --temp 0`, lighthouse prompt) + +| variant | t/s | md5 | +|---|---|---| +| round-1 baseline | 43.40–43.52 | `db0c504c…` | +| + attention parity only | 44.59–44.67 | `db0c504c…` | +| + greedy chain (current) | **45.47–45.55** | `db0c504c…` | + +Long-context (2K prompt): 44.12 chained vs 43.40 classic. 1024-token run and +an early-stop prompt: md5-identical. Harness (prefix 2048): 529 rows / +68,389,120 logits / 528 ids bit-identical, 43.18 vs 43.10 t/s (only layers +0–1 qualify at long prefix). SSD streaming decode smoke: OK (chain declines, +classic path taken). + +## Remaining headroom (re-estimated) + +GPU busy ≈ 21.85 ms/token now; wall ≈ 21.96 ms. The boundary is ~0.1–0.15 ms +(argmax→embed dependency + drain). Still open, expectations lowered by the +round-1 evidence: HC epilogue tail fusions (~0.3 ms), KV-staging direct read +(~0.25 ms; a prior attempt diverged — the addressing facts are in README's +stage-counter section and the code comments at cpy.metal:147), MXFP4 +concurrent shared-expert stream (machinery exists hard-gated to IQ2, +ds4_metal.m:9286; medium confidence, and the "more parallelism throttles this +GPU" lesson applies). Each could add ~0.2–0.6 t/s; none is needed for 45. + +## Watch item + +The balanced harness failed twice at `step=0 variant=control` ("metal decode +failed") with an early round-2 binary, then passed 9+ consecutive runs with +semantically identical code. Unexplained; both failures immediately followed +sustained benching (the documented thermal-transient window). If it recurs, +reproduce with `--warmup 1 --tokens 2` and capture full stderr. + +## Tools (unchanged from round 1) + +- Commit-only stage profiler: `DS4_METAL_DECODE_STAGE_PROFILE=1 + DS4_METAL_STAGE_COUNTERS=1 ./ds4 …` (trustworthy; end-and-wait inflates). +- Balanced A/B harness: `make metal-decode-schedule-bench && ./speed-bench/ + metal_decode_schedule_bench -m ds4flash.gguf --candidate-env NAME + --include-selection --tokens 512` (aborts unless bit-identical). +- Transcript md5 oracle: `./ds4 -m ds4flash.gguf -p "Write a short story + about a lighthouse keeper." -c 8192 -n 128 --temp 0 | md5` → `db0c504c…`. +- Tensor dump bisect: `DS4_METAL_GRAPH_DUMP_PREFIX=/tmp/x + DS4_METAL_GRAPH_DUMP_NAME= DS4_METAL_GRAPH_DUMP_LAYER=N + DS4_METAL_GRAPH_DUMP_POS=P ./ds4 …` (synchronizes and dumps; comparing + classic vs chain dumps located the hash-router divergence in one pass). +- One `ds4` instance at a time; idle ~60 s after heavy runs (transient + 2–10× slowdown recovers by itself). + +## Round-1 closed avenues still stand + +Eight bit-exact kernel variants (all ≤0.03%), DSpark strict/non-strict +(43.07 / peaks 40.2), three microbatch increments, split-schedule sweep +(default 4 optimal), readback-bubble premise (disproved). Details in +speed-bench/README.md. The round-1 verdict "bit-exact and >45 t/s are +jointly infeasible" was wrong: the ledger's averaged `attn_inv_rope` line +hid the 65/32 µs parity alternation, and the boundary's "wake/launch +latency" was recoverable by keeping the token id on-device. + +## Round-3 addendum (Aug 22): session chain + headroom list exhausted + +- **Session chain decode**: `ds4_session_eval_chain_greedy` (ds4.c) reuses + `metal_graph_greedy_chain` for session API callers; `ds4-eval` decodes in + bursts capped to stay out of the think-close controller window. Bit-exact + (traces identical), eval decode **44.5 → 45.2 t/s**. Kill switch + `DS4_DISABLE_GREEDY_CHAIN=1` covers both CLI and session paths. +- **MoE down-sum6+HC4 tail fusion**: landed bit-exact, speed-neutral, + default on (rollback `DS4_METAL_DISABLE_DECODE_MOE_HC_FUSION`). +- **KV-staging direct read**: implemented bit-exact (wrap-safe) but −7% + (per-head F32 re-read/re-convert amplification); landed gated OFF + (opt-in `DS4_METAL_ENABLE_DECODE_RAW_DIRECT_KV`). +- The round-2 "remaining headroom" list is now closed out: both quantified + items measured (neutral / negative), the concurrent shared-expert stream + stays contraindicated by the throttling evidence. Wall−GPU gap is ~0.1 + ms/token; the next real gain must come from the MoE or attention core + itself. diff --git a/speed-bench/README.md b/speed-bench/README.md index 6d4188e72..9433c2773 100644 --- a/speed-bench/README.md +++ b/speed-bench/README.md @@ -1,5 +1,10 @@ ## Benchmarking +Session handoff (Aug 21 decode campaign): see +`speed-bench/DECODE-CAMPAIGN-HANDOFF.md` for the cold-restart kit — verified +baselines, the closed-avenue list, tool usage, and the decision required +before resuming work toward >45 t/s. + Here we collect prefill and generation speed obtained with different hardware. Run `ds4-bench` as: @@ -27,6 +32,204 @@ python3 speed-bench/plot_speed.py speed-bench/m3_max.csv --title "M3 Max t/s" The script uses only the Python standard library. By default it writes a file next to the CSV using the `_ts.svg` suffix, such as `speed-bench/m3_max_ts.svg`. +### DSpark speculation on M3 Ultra (measured, not yet profitable) + +DSpark non-strict decode was measured end to end on M3 Ultra with the MXFP4 +model and the 0731 support GGUF. Three systemic costs keep it below plain +decode (~43.7 t/s) today; scheduler tuning alone cannot fix them: + +1. The verify pass runs the speculative suffix through the generic batch + prefill kernels (`metal_graph_encode_layer_batch`): ~46-60 ms per verify + vs the ~29-33 ms memory floor for two rows (the extra routed-expert reads + are inherent). A commit-only stage decomposition showed a uniform + 1.5-2.5x per-stage excess (routed MoE 14.3 ms vs 5.9 decode, HC pre + 7.2 vs 1.7, output projection 7.1 vs 4.8, attention 4.4 vs 2.1). + The fix is the N<=6 microbatch verifier on the decode-grade kernels that + the verifier's own header comment calls out as "not yet" written. + Caution for that build: three measured increments (per-row MoE, per-row + HC pre, and a strict-oracle-verified dual-row HC-pre kernel) each + recovered nothing, and the batch MoE already runs at its distinct-expert + floor. The honest reading is that the ~29-33 ms "perfect sharing" + verify floor is not deliverable on this GPU; the N=2 verify near 50 ms + is close to its real floor, so speculation is unlikely to beat the + 43.7 t/s plain decode on M3 Ultra with this model. +2. The draft propose chain costs ~3-8 ms/cycle and its confidence gate + (`sigmoid(confidence0) >= threshold`, Metal default 0.6) declines on + 45-75% of cycles; each such cycle still pays the propose before falling + back to one plain decode. Plain decode inside a DSpark session measures + a normal ~23.3 ms (the hidden-state capture is not a decode tax). +3. Per-cycle bookkeeping (checkpoint, snapshot, commit, propose-fail + waste) accounts for a further ~4-5 ms/cycle beyond the measured + propose+verify+decode components. + +Scheduler and confidence knobs were swept: `DS4_DSPARK_SCHEDULER_NO_DRAFT_SKIP=0` +(retry the draft every cycle) raised proposals from 91/179 to 125/147 cycles +and accepted drafts to 101 (80.8% accept); the best combination measured +39.5 t/s at `--dspark-confidence 0.75` on a code prompt, still below the +43.7 t/s plain-decode equilibrium. + +A genuine dual-row HC-pre producer kernel (one F16 mix-weight fetch serving +both rows, each row keeping the one-row kernel's exact reduction trees) was +also built for the N=2 verify. Correctness was proven with the strict-mode +oracle — `--dspark-strict` output matched plain decode byte-for-byte — after +fixing a buffer-index mismatch that silently misbound five kernel arguments +(the strict oracle is the right validation tool for any verify-kernel work). +Measured honestly, per-verify cost was unchanged (~52 ms, n=4, vs ~50 ms +rollback over 57 verifies) and the diverged non-strict token stream made the +draft decline on 99% of cycles: the serialized hc_pre share had been +overestimated by the stage decomposition, and the win is not there. + +Two dispatch-routing microbatch increments were also built, validated, and +measured negative before reverting: per-row routed MoE through the +single-token static-trip kernels (bit-identical output, verify_layer +unchanged at 1225 vs 1217 ms) and per-row HC pre through the decode fused +producer kernel (valid non-strict output, 39.3 vs 40.2 t/s). Both show the +verify excess lives in the batch kernels' execution, not in dispatch +routing; the microbatch build must write genuine small-N batched kernels. Strict mode (`--dspark-strict`) +measures 43.07 t/s, i.e. no gain, as accepted blocks are re-run through +one-token decode to stay byte-identical. + +### Metal decode stage GPU counters + +The end-and-wait stage profiler (`DS4_METAL_DECODE_STAGE_PROFILE=1`) adds a +synchronization per stage boundary and changes the schedule, so its numbers +are inflated by per-boundary waits. The stage-counter diagnostic keeps the +production token mostly intact: every stage boundary commits the open batch +command buffer without waiting, so the GPU queue stays fed, and each stage's +GPU busy span is printed after the token: + +``` +DS4_METAL_DECODE_STAGE_PROFILE=1 DS4_METAL_STAGE_COUNTERS=1 ./ds4 -m ds4flash.gguf \ + -p "Write a short story." -c 8192 -n 24 --temp 0 +``` + +Both env vars are required: the first arms the boundary macros, the second +switches them from end-and-wait to commit-only sampling. The concurrent +shared-expert/routed-MoE overlap stays armed under counters (only the +serializing profiler disables it), but the per-stage command buffers queue in +order, so overlapped stages report their serialized costs. The per-token +`total-cb-busy` line matches the production GPU-busy time (about 22.5 ms on +M3 Ultra at a short context), which is the check that the attribution is +faithful. M3 Ultra decode at a short context attributes the token roughly as: +routed MoE 5.9 ms, attention output projections 4.8 ms, Q lora path 5.1 ms +(Q-A/KV/compressor quad kernel 41 us + Q-B matvec 59 us per layer), attention +core plus inverse RoPE 2.1 ms, router/shared gate-up 1.9 ms, and about 3.4 ms +of per-layer HC pre/post bookkeeping, with the remaining dense Q8_0 matvecs +streaming at 590-650 GB/s, i.e. at the memory wall. + +### Metal decode raw-layer gathered attention A/B (45 t/s round, part 1) + +Decode attention has two schedules with a per-layer-parity split: ratio-4 +layers use the gathered path (fused KV staging with pad fusion + packed32 +attention with fused inverse RoPE, two dispatches), while ratio-0/128 layers +with `n_comp == 0` used a five-dispatch raw path (ring copy, standalone pad, +vec, plain reduce, standalone RoPE tail). Commit-only stage counters showed +the raw layers at ~65 µs versus ~32 µs — about 0.7 ms/token, hidden in the +averaged ledger. Decode now routes `n_comp == 0` layers through the gathered +path (the staging kernel already handles `n_comp == 0`; the packed32 gate's +`n_comp != 0` term is relaxed) whenever `use_mask == 0`. Same kernels, same +reduction topology; the routing change is packaging only. Rollbacks: +`DS4_METAL_DISABLE_DECODE_RAW_GATHERED_ATTN` (raw path) and +`DS4_METAL_DISABLE_DECODE_RAW_PACKED32` (staged vec+reduce instead of +packed32 on raw layers). + +``` +./speed-bench/metal_decode_schedule_bench \ + --candidate-env DS4_METAL_DISABLE_DECODE_RAW_GATHERED_ATTN \ + --include-selection --tokens 512 +``` + +Balanced M3 Ultra A/B at the harness's 2048-token prefix: 43.18/43.10 t/s +(+0.16%; only layers 0–1 qualify once odd layers hold compressed rows), all +529 frontier rows / 68,389,120 logits / 528 selected ids bit-identical. In +the campaign CLI regime (short prompt, `-n 128`), interleaved runs gave +44.59/44.67 t/s new versus 43.52/43.40 t/s rollback (+2.7%), transcripts +md5-identical (`db0c504c…`). + +### Metal greedy chain decode A/B (45 t/s round, part 2) + +The classic one-shot decode loop serialized every token through the host: +`waitUntilCompleted`, a 517 KiB logits readback, a CPU argmax, then the next +token's encode — about 0.5 ms/token of GPU idle at the boundary. Chained +greedy decode (`metal_graph_greedy_chain`, engaged by +`generate_metal_graph_raw_swa` when resident, non-quality, non-streaming, +greedy) keeps the token id on-device: each token's graph ends with the GPU +argmax writing the next id into a device ring, and the next token's embedding +gathers it from the ring. Encoding runs two tokens ahead of the host's +confirm cursor, so command buffers are always committed before the GPU drains +the previous token; the host lags only to print and check stop tokens (one +shared-event wait per token, hidden by the encode-ahead). The hash-layer +router select reads the id from the ring through the existing +`use_token_buffer` kernel argument; the host-side hash override is skipped +(the resident fixed-route MoE never consumes it). Bit-exactness: all kernels +and inputs are unchanged, and the GPU argsort top-1 reproduces the CPU argmax +including lowest-index ties, so the transcript is identical. +`DS4_DISABLE_GREEDY_CHAIN=1` restores the classic loop; +`DS4_GREEDY_CHAIN_DEBUG` / `DS4_GREEDY_CHAIN_DUMP_IDS` are diagnostics. +Interleaved M3 Ultra CLI runs: 45.47/45.55 t/s chained versus 44.73/44.73 +classic (+1.75%), transcripts md5-identical; at a 2K-token prefix 44.12/43.40 +(+1.7%); a 1024-token run and an early-stop prompt matched md5 exactly. + +Combined round state: 43.46 → 45.51 t/s (+4.7%), bit-exact, `make test` +44/44. + +### Session greedy chain (ds4-eval) + two headroom probes (Aug 22, round 3) + +**Session chain decode** brings the round-2 chain to the session API used by +`ds4-eval`. `ds4_session_chain_greedy_supported` / +`ds4_session_eval_chain_greedy` (ds4.c, next to `ds4_session_eval`) drive +`metal_graph_greedy_chain` on an existing session graph: seed = CPU argmax of +`s->logits` (identical to the classic temp-0 sample), pos0 = +`checkpoint.len`, approved tokens pushed into `s->checkpoint` by a trampoline +only after the caller's callback approves them; on full completion +`logits_out = s->logits` preserves the session logits invariant (on early +stop the logits are stale — eval only stops early on stop-token/quit/switch, +where they are never read). Guards mirror the CLI set plus session state: +no GLM/CPU/distributed/TP/multi-tier, `support_kind == DS4_SUPPORT_NONE`, no +ssd-streaming/quality/CPU-router/steering, same env kill switch +(`DS4_DISABLE_GREEDY_CHAIN=1` forces the classic loop everywhere). + +`ds4-eval` decodes in bursts (eval_chain_on_token carries the classic loop's +per-token bookkeeping verbatim). Bursts are capped to stay out of the +think-close controller window (`remaining - soft_limit` while thinking, len>1 +closes use the hard window) — inside the window the classic step runs, so +forced (non-argmax) closes behave identically. Verified bit-exact: 3-case +traces (`--questions 3 -n 4096`, think and `--nothink`) identical except +volatile timestamp/seed/elapsed fields; per-case answers, token counts, and +think_close records match. Speed: **45.22–45.24 t/s chained vs +44.44–44.59 classic** (+1.5%) on the eval cases; CLI oracle unchanged +(`db0c504c…`, 45.90 t/s), `make test` 44/44. + +**Probe 1 — MoE down-sum6+HC4 tail fusion** (landed, default ON; rollback +`DS4_METAL_DISABLE_DECODE_MOE_HC_FUSION`): layer tail reordered to +pair-SwiGLU → plain shared-down matvec → new +`kernel_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_hc4_f32` +(metal/moe.metal) doing the down-sum6 + `shared_out` add + HC4 expand in one +dispatch, eliminating the `routed_out` f32 materialization and one kernel +boundary per layer. Bit-exact (md5 n=128/512, harness bit-identical, +wrap-safe); **speed-neutral** — the old path already fused HC into the +shared-down kernel, so only a 16 KiB round trip was removed against the +26 MB expert stream. Kept: one less dispatch and a simpler tail. + +**Probe 2 — KV-staging direct read** (landed, default OFF; opt-in +`DS4_METAL_ENABLE_DECODE_RAW_DIRECT_KV`, hard-off +`DS4_METAL_DISABLE_DECODE_RAW_DIRECT_KV`, loud +`DS4_METAL_REQUIRE_DECODE_RAW_DIRECT_KV`): packed32 attention variant reading +the raw F32 ring and comp F16 caches directly (per-row source selection in +QK and PV loops, in-place tail, wrap-safe). Bit-exact everywhere including a +probe-verified ring-wrap run — **but ~7% slower** (42.1 vs 45.3 t/s): the F32 +raw rows (2048 B) are re-read and re-converted by all 64 head threadgroups +every token, ~2× raw-region traffic versus the one-time staging conversion +that amortizes across heads (+30 µs/layer on `attn_inv_rope`). Confirms the +round-1 lesson: more parallel traffic throttles this GPU. May still win on +devices with different L2/ALU balance. + +Round-3 verdict: the handoff's remaining-headroom list is exhausted — HC +fusion neutral, direct-KV negative, and the concurrent shared-expert stream +(IQ2-gated) is contraindicated by the same throttling evidence. The +remaining wall−GPU gap is ~0.1 ms/token; further gains need a cheaper MoE or +attention core, not packaging. + ### Metal decode schedule A/B Build the balanced, same-engine Metal decode comparison with: @@ -54,6 +257,25 @@ legacy decode path, including token selection, use: --tokens 1024 ``` +The pre-M5 one-token Flash attention-output LOW projection also has an exact +fixed-shape Q8_0 kernel. Compare it with the generic rollback using: + +``` +./speed-bench/metal_decode_schedule_bench \ + --candidate-env DS4_METAL_DISABLE_PRE_M5_ATTN_OUT_LOW_Q8_STATIC \ + --include-selection \ + --tokens 512 +``` + +Balanced M3 Ultra A/B runs favored the fixed-shape kernel by 0.53%, 0.53%, +and 0.58% at a 2K-token prefix (43.03/42.80, 42.96/42.73, and 43.24/43.00 +tok/s) and by 0.48% at an 8K-token prefix (38.44/38.26 tok/s). An independent +IQ2/Q2-model run at 2K gave 44.47/44.23 tok/s (+0.56%). All 1,909 compared +rows, 246,795,520 full-vocabulary logits, and 1,904 selected token IDs were +bit-identical. +Performance was measured on M3 Ultra; the exact host gate covers the shared +M1-M4 Flash shape and otherwise retains the generic kernel. + ### Metal prefill variant A/B Build the balanced prefill comparison. To compare the default resident pre-M5 @@ -74,6 +296,231 @@ default, use its down-specific rollback as the candidate: --candidate-env DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_MM_ID_DOWN_TAIL_SIMDGROUP_CULL ``` +For resident pre-M5 MXFP4 prefills of 32 through 2047 tokens, the exact +scatter map, compact pair tile, pair/down tail culls, and down half-LUT now +use the same defaults as the established 2K+ path. Compare the complete +short-prefill extension with its aggregate rollback using: + +``` +./speed-bench/metal_prefill_variant_bench \ + --prefix-tokens 256 \ + --candidate-env DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_SMALL_PREFILL +``` + +Balanced M3 Ultra A/B medians for the tuned path versus the rollback were +112.43/89.25 tok/s at 32 tokens (+26.0%), 111.95/89.30 at 33 (+25.4%), +180.48/142.59 at 64 (+26.6%), 275.72/219.98 at 128 (+25.3%), +389.27/316.93 at 256 (+22.8%), 512.94/439.39 at 512 (+16.7%), +173.94/157.08 at 1024 (+10.7%), and 632.30/593.70 at 2047 (+6.5%). Every +one of the 64 measured runs produced bit-identical full-vocabulary logits. +Performance was measured on M3 Ultra; the guarded default also covers the +shared resident M1-M4 path. + +### Metal batch indexer-query pruning A/B + +On zero-prefix ratio-4 layers, the indexer query and its per-head weights are +not consumed until the compressed cache grows beyond the 512-row top-k. The +resident pre-M5 path now skips the otherwise dead Q projection, RoPE, QAT, and +weight projection for batches of at least 32 tokens while the final compressed +count remains at or below top-k. The indexer compressor and its persistent +cache/state updates are unchanged. Compare the pruned path with its rollback: + +``` +./speed-bench/metal_prefill_variant_bench \ + --prefix-tokens 2048 \ + --warmup-tokens 2048 \ + --repeats 4 \ + --candidate-env DS4_METAL_DISABLE_PRE_M5_BATCH_INDEXER_QUERY_PRUNE +``` + +Balanced M3 Ultra A/B throughput for the pruned path versus rollback was +115.00/112.98 tok/s at 32 tokens (+1.79%), 277.42/273.81 at 128 (+1.32%), +509.19/502.28 at 512 (+1.38%), 606.13/597.65 at 1024 (+1.42%), and +660.17/651.29 at 2048 (+1.36%). The last eligible prefix, 2051 tokens, gained +1.41%; 2052 tokens was flat, confirming that the query path remains enabled +when the 513th compressed row first makes top-k selection necessary. All 56 +prefill runs and 7,239,680 compared full-vocabulary logits were bit-identical. +A 2051-token prefix followed across the row-513 transition also matched three +full-vocabulary rows and two selected token IDs exactly. Performance was +measured on M3 Ultra; the guarded path covers resident single-device M1-M4. + +### M5 Max port of batch indexer-query pruning + +The pruning eligibility now uses the shared M5 port helper, so resident +single-device M5 decode graphs take the same skip while the compressed cache +stays at or below top-k. The M5 rollback is +`DS4_METAL_DISABLE_M5_BATCH_INDEXER_QUERY_PRUNE` (the pre-M5 name still +controls M1-M4 together with the aggregate `DS4_METAL_DISABLE_PRE_M5_DECODE_PORTS` +switch). Balanced M5 Max A/B (IQ2_XXS/Q2_K `ds4flash.gguf`) for the pruned +path versus rollback was 753.62/732.83 tok/s at 2048 tokens (+2.84%); all 8 +runs and 1,034,240 compared full-vocabulary logits were bit-identical. The +standard Prometti-sposi sweep's virgin first frontier measured 790.3 vs +727.1 tok/s (+8.7%) with decode unchanged; frontiers past the eligibility +window stayed within the ±1-2% noise band. + +Two other pre-M5 wins were measured on M5 Max and rejected: the indexed +prefill RB4/heads16-dual-RB4 kernels (bit-exact, -0.27% to -1.56% versus the +MPP dual-heads default across two balanced sessions — M5 neural accelerators +keep the MPP path ahead), and the Q2 2/32 decode command-buffer split +(bit-exact, -0.62%/-0.79% versus the 4/none default at 512 and 1024 tokens). +Do not re-port these without new evidence. + +### Metal batch Q/KV finalizer A/B + +The M3 resident Flash prefill path now follows vLLM's horizontal Q/KV +finalization schedule while retaining DS4's existing Q-head PSO. It dispatches +that exact Q RMSNorm+RoPE kernel concurrently with a KV-only kernel that folds +KV RoPE, the FP8 round trip, F16 rounding, and raw-ring insertion together. +The default covers 128 through 4096 tokens per dispatch; longer contexts use +the normal 4096-token chunks. A 32-token dispatch regressed and 64 tokens did +not clear the 0.3% acceptance threshold. Compare the retained path with its +serial rollback using: + +``` +./speed-bench/metal_prefill_variant_bench \ + --prefix-tokens 512 \ + --warmup-tokens 512 \ + --repeats 4 \ + --candidate-env DS4_METAL_DISABLE_PRE_M5_BATCH_QKV_FINALIZE +``` + +Balanced M3 Ultra A/B throughput for the concurrent path versus rollback was +273.60/272.63 tok/s at 128 tokens (+0.36%), 502.44/500.37 at 512 (+0.41%), +598.07/595.87 at 1024 (+0.37%), 651.27/649.11 at 2048 (+0.33%), and +615.06/612.37 at 8192 (+0.44%). All 56 prefill runs and 7,239,680 compared +full-vocabulary logits were bit-identical. +A 512-token prefill followed by decode also matched 73 full-vocabulary rows and +72 selected token IDs exactly, covering the persisted raw-cache state. The +schedule is based on vLLM's +[fused DeepSeek V4 finalizer](https://github.com/vllm-project/vllm/blob/c8de519917ce549f72132952116185e38b37c95d/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu#L382-L603), +but keeps Q and KV in separate Metal pipeline states to preserve DS4's exact Q +fast-math code generation. + +### Metal prefill indexed-attention four-row staging A/B + +The resident single-device pre-M5 indexed-attention prefill path now stages +four raw or F16-compressed K/V rows per threadgroup barrier while consuming +them in the original scalar-kernel row order. The specialization covers +non-quality batches of at least 32 tokens with 64 heads, 512-wide heads, +top-k 512, a 128-row raw window, ratio 4, and more than 512 compressed rows; +SSD streaming, TP2, decode, M5, and other shapes keep the original path. +Compare it with the +one-row staging rollback using: + +``` +./speed-bench/metal_prefill_variant_bench \ + --prefix-tokens 8192 \ + --warmup-tokens 4096 \ + --repeats 1 \ + --candidate-env DS4_METAL_DISABLE_PRE_M5_INDEXED_ATTN_PREFILL_RB4 +``` + +Balanced M3 Ultra A/B throughput for four-row staging versus rollback was +631.53/626.50 tok/s at 4100 tokens (+0.80%), +620.39/614.76 at 8192 (+0.92%), 594.11/589.36 at 16384 (+0.81%), and +557.91/553.71 at 32768 (+0.76%). The 2048-token boundary, where compressed +rows do not yet exceed top-k, was flat. All 16 active-path prefill runs and +2,068,480 compared full-vocabulary logits were bit-identical. A direct forced +kernel oracle also matched all 1,048,576 attention-output floats exactly while +covering raw-ring wrap, visibility stops, and one- through three-row tails. +Separate 8K and 32K prefix-to-decode checks matched 106 full-vocabulary +frontiers, 13,703,680 floats, and 104 selected token IDs. Eight-row staging was +also exact but was 0.17-0.21% slower at the first two screened sizes, so only +the four-row specialization is retained. Performance was measured on M3 +Ultra; the guarded default covers the shared resident M1-M4 path. + +The same guarded path now groups sixteen heads per 256-thread workgroup by +having each SIMDgroup update two heads from every staged four-row block. This +halves K/V staging and workgroup count without changing either head's row +order or online-softmax arithmetic. Compare it with the accepted eight-head +RB4 rollback using: + +``` +./speed-bench/metal_prefill_variant_bench \ + --prefix-tokens 8192 \ + --warmup-tokens 4096 \ + --repeats 1 \ + --candidate-env DS4_METAL_DISABLE_PRE_M5_INDEXED_ATTN_PREFILL_HEADS16_DUAL_RB4 +``` + +Balanced M3 Ultra A/B throughput for heads16 dual RB4 versus heads8 RB4 was +638.45/629.00 tok/s at 2052 tokens (+1.50%), 651.08/637.46 at 4100 +(+2.14%), 640.33/626.12 at 8192 (+2.27%), and 613.72/600.81 at 16384 +(+2.15%). The first matched 32768-token pair was 576.10/563.42 (+2.25%); +later slots in that process encountered severe system throttling and are not +used for the comparison. All 28 timed full-vocabulary rows were bit-identical. +The forced direct oracle matched all 1,048,576 attention-output floats and +proved both the dedicated rollback and new selector. Separate 8K and 32K +prefix-to-decode checks matched 94 full-vocabulary frontiers, 12,152,320 +floats, and 92 selected token IDs. + +### Metal batch MoE sum6-to-HC4 epilogue A/B + +The resident single-device pre-M5 MXFP4 prefill path now consumes the six +routed expert-down rows directly in the HC4 expand/add/split epilogue. The +kernel preserves the original `s0 + s1 + ... + s5`, shared-expert add, and HC4 +post/comb accumulation order while removing the standalone sum6 dispatch and +the routed-output F32 materialization. The default covers batches of 32 through +4096 tokens; longer prefixes use the normal 4096-token chunks. Debug, profiling, +steering, SSD, TP2, quality, decode, and other tensor shapes keep the original +path. Compare the fused path with its rollback using: + +``` +./speed-bench/metal_prefill_variant_bench \ + --prefix-tokens 8192 \ + --warmup-tokens 4096 \ + --repeats 4 \ + --candidate-env DS4_METAL_DISABLE_PRE_M5_BATCH_MOE_SUM6_HC_FUSION +``` + +Balanced M3 Ultra A/B throughput for the fused path versus rollback was +663.24/660.81 tok/s at 2048 tokens (+0.37%), 621.83/619.49 at 8192 +(+0.38%), 596.50/594.76 at 16384 (+0.29%), and 560.30/558.54 at 32768 +(+0.32%). The 32-token point was flat (-0.03%); 128, 512, and 4096 tokens +gained 0.19%, 0.18%, and 0.28%, respectively. All 64 prefill runs and +8,273,920 compared full-vocabulary logits were bit-identical. A direct oracle +also matched all 527,372 HC output floats exactly for a tail shape and the +production 32-by-4096 shape. Separate 8K and 32K prefix-to-decode checks +matched 146 full-vocabulary frontiers, 18,874,880 floats, and 144 selected +token IDs exactly. Performance was measured on M3 Ultra; the guarded default +covers the shared resident M1-M4 shape. + +### Metal batch Q8 attention-output-to-HC4 epilogue A/B + +The resident single-device pre-M5 Q8 attention-output path now feeds the +aligned 8192-to-4096 output-B matmul tile directly into the HC4 expand/split +epilogue. The specialization preserves the legacy Q8_0 dequantization and +simdgroup-MMA order, stages every 64-by-32 result through an 8 KiB F32 +threadgroup tile, and then repeats the original scalar HC4 post/comb +accumulation order. This removes the global F32 `attn_out` write/read and the +standalone HC dispatch without removing the materialized F32 rounding +boundary. The default is deliberately limited to aligned batches of 512 +through 4096 tokens; shorter batches, tails, debug/profiling/steering, +quality, SSD, TP2, M5, and other shapes retain the original path. Compare the +fused path with rollback using: + +``` +./speed-bench/metal_prefill_variant_bench \ + --prefix-tokens 8192 \ + --warmup-tokens 4096 \ + --repeats 2 \ + --candidate-env DS4_METAL_DISABLE_PRE_M5_BATCH_ATTN_OUT_HC_FUSION +``` + +Balanced M3 Ultra A/B throughput for the fused path versus rollback was +526.16/523.77 tok/s at 512 tokens (+0.46%), 668.43/663.66 at 2048 +(+0.72%), 626.47/621.49 at 8192 (+0.80%), and 599.57/596.92 at 16384 +(+0.44%). All 44 retained timing runs and 5,688,320 compared +full-vocabulary logits were bit-identical. A direct production-shape oracle +also matched all 262,144 low-rank values and 524,288 HC outputs exactly and +confirmed that the fused path did not touch the dead `attn_out` buffer. +Separate 8K and 32K prefix-to-decode checks matched 14 full-vocabulary +frontiers, 1,809,920 floats, and 12 selected token IDs exactly. A standalone +32K timing process was discarded because sustained thermal throttling changed +slot time from 58 to 142 seconds; the 32K run is correctness evidence only. +Performance was measured on M3 Ultra; the guarded default covers the shared +resident M1-M4 path. + The harness uses one Metal engine and fresh sessions for every run. It warms both variants with at least 32 tokens, alternates control/candidate order in ABBA and BAAB blocks, poisons host logit buffers before copying, and aborts diff --git a/speed-bench/m3_ultra_mxfp4.csv b/speed-bench/m3_ultra_mxfp4.csv new file mode 100644 index 000000000..76978425d --- /dev/null +++ b/speed-bench/m3_ultra_mxfp4.csv @@ -0,0 +1,33 @@ +ctx_tokens,prefill_tokens,prefill_tps,gen_tokens,gen_tps,gen_first_ms,gen_steady_tokens,gen_steady_tps,kvcache_bytes +2048,2048,190.08,128,10.26,414.017,127,10.59,52184460 +4096,2048,545.01,128,34.66,29.753,127,35.37,80373132 +6144,2048,540.43,128,34.46,30.823,127,35.18,108561804 +8192,2048,533.47,128,34.20,31.794,127,34.92,136750476 +10240,2048,525.39,128,33.98,32.349,127,34.68,164939148 +12288,2048,520.21,128,33.98,32.144,127,34.67,193127820 +14336,2048,514.57,128,33.82,32.465,127,34.53,221316492 +16384,2048,509.90,128,33.57,32.177,127,34.26,249505164 +18432,2048,502.02,128,33.40,32.594,127,34.04,277693836 +20480,2048,497.13,128,33.35,32.434,127,33.97,305882508 +22528,2048,493.48,128,33.21,32.826,127,33.88,334071180 +24576,2048,488.36,128,33.16,32.771,127,33.83,362259852 +26624,2048,479.97,128,33.03,33.129,127,33.68,390448524 +28672,2048,476.35,128,33.04,33.123,127,33.64,418637196 +30720,2048,472.39,128,32.74,33.512,127,33.34,446825868 +32768,2048,468.42,128,32.49,33.459,127,33.08,475014540 +34816,2048,461.24,128,32.30,33.388,127,32.88,503203212 +36864,2048,458.64,128,32.31,34.161,127,32.89,531391884 +38912,2048,454.83,128,32.18,33.691,127,32.75,559580556 +40960,2048,450.11,128,32.08,33.546,127,32.69,587769228 +43008,2048,443.67,128,31.72,35.117,127,32.33,615957900 +45056,2048,439.56,128,31.79,34.573,127,32.40,644146572 +47104,2048,436.40,128,31.63,34.108,127,32.24,672335244 +49152,2048,432.63,128,31.64,34.833,127,32.25,700523916 +51200,2048,427.73,128,31.44,35.283,127,32.05,728712588 +53248,2048,424.36,128,31.53,35.141,127,32.13,756901260 +55296,2048,420.70,128,31.18,35.228,127,31.77,785089932 +57344,2048,417.16,128,31.39,34.652,127,32.00,813278604 +59392,2048,411.15,128,31.08,35.906,127,31.67,841467276 +61440,2048,409.27,128,30.96,35.877,127,31.54,869655948 +63488,2048,406.26,128,30.83,35.807,127,31.37,897844620 +65536,2048,403.30,128,30.51,37.810,127,31.04,0 diff --git a/speed-bench/m3_ultra_mxfp4_r2.csv b/speed-bench/m3_ultra_mxfp4_r2.csv new file mode 100644 index 000000000..351dc79c0 --- /dev/null +++ b/speed-bench/m3_ultra_mxfp4_r2.csv @@ -0,0 +1,33 @@ +ctx_tokens,prefill_tokens,prefill_tps,gen_tokens,gen_tps,gen_first_ms,gen_steady_tokens,gen_steady_tps,kvcache_bytes +2048,2048,596.78,128,37.32,31.197,127,38.17,52184460 +4096,2048,544.09,128,34.51,30.047,127,35.22,80373132 +6144,2048,538.51,128,34.28,31.884,127,35.00,108561804 +8192,2048,532.10,128,34.07,32.041,127,34.78,136750476 +10240,2048,523.59,128,33.87,32.217,127,34.57,164939148 +12288,2048,518.54,128,33.88,32.447,127,34.58,193127820 +14336,2048,513.78,128,33.72,32.436,127,34.41,221316492 +16384,2048,508.52,128,33.43,32.503,127,34.11,249505164 +18432,2048,500.69,128,33.29,32.775,127,33.96,277693836 +20480,2048,496.36,128,33.22,32.883,127,33.89,305882508 +22528,2048,492.10,128,33.13,32.898,127,33.80,334071180 +24576,2048,487.48,128,33.12,32.896,127,33.80,362259852 +26624,2048,479.50,128,32.99,33.844,127,33.66,390448524 +28672,2048,474.97,128,32.93,33.054,127,33.59,418637196 +30720,2048,470.72,128,32.81,33.275,127,33.46,446825868 +32768,2048,466.85,128,32.55,34.221,127,33.22,475014540 +34816,2048,459.73,128,32.33,33.734,127,32.98,503203212 +36864,2048,456.94,128,32.31,34.346,127,32.95,531391884 +38912,2048,453.63,128,32.17,33.835,127,32.80,559580556 +40960,2048,449.36,128,32.13,34.748,127,32.77,587769228 +43008,2048,443.00,128,31.96,34.147,127,32.60,615957900 +45056,2048,438.73,128,31.95,34.086,127,32.58,644146572 +47104,2048,435.34,128,31.80,34.587,127,32.42,672335244 +49152,2048,432.29,128,31.76,35.155,127,32.38,700523916 +51200,2048,427.21,128,31.64,34.803,127,32.25,728712588 +53248,2048,423.87,128,31.58,34.824,127,32.21,756901260 +55296,2048,420.08,128,31.43,35.012,127,32.04,785089932 +57344,2048,416.89,128,31.41,35.428,127,32.02,813278604 +59392,2048,410.89,128,31.25,35.365,127,31.84,841467276 +61440,2048,408.86,128,31.20,36.338,127,31.81,869655948 +63488,2048,406.04,128,31.05,35.654,127,31.66,897844620 +65536,2048,402.63,128,30.72,39.509,127,31.33,0 diff --git a/speed-bench/m3_ultra_mxfp4_r2_ts.svg b/speed-bench/m3_ultra_mxfp4_r2_ts.svg new file mode 100644 index 000000000..34ea378f1 --- /dev/null +++ b/speed-bench/m3_ultra_mxfp4_r2_ts.svg @@ -0,0 +1,50 @@ + + + + +M3 Ultra (512GB) MXFP4 t/s — round 2 + +0 + +200 + +400 + +600 + +800 + +1k +0 +10 +20 +30 +40 + +0 + +20k + +40k + +60k + + + +ctx size +prefill t/s +generation t/s + + + + +prefill + +generation + diff --git a/speed-bench/m3_ultra_mxfp4_rebased.csv b/speed-bench/m3_ultra_mxfp4_rebased.csv new file mode 100644 index 000000000..266c9313d --- /dev/null +++ b/speed-bench/m3_ultra_mxfp4_rebased.csv @@ -0,0 +1,33 @@ +ctx_tokens,prefill_tokens,prefill_tps,gen_tokens,gen_tps,gen_first_ms,gen_steady_tokens,gen_steady_tps,kvcache_bytes +2048,2048,632.90,128,39.49,30.221,127,39.59,52184460 +4096,2048,575.43,128,36.44,29.733,127,36.50,80373132 +6144,2048,569.95,128,36.23,31.643,127,36.31,108561804 +8192,2048,563.11,128,35.95,30.938,127,36.02,136750476 +10240,2048,553.98,128,35.70,32.155,127,35.77,164939148 +12288,2048,548.25,128,35.66,31.949,127,35.74,193127820 +14336,2048,542.53,128,35.51,32.251,127,35.59,221316492 +16384,2048,536.77,128,35.20,32.410,127,35.27,249505164 +18432,2048,528.44,128,35.02,32.626,127,35.09,277693836 +20480,2048,523.02,128,34.97,32.869,127,35.05,305882508 +22528,2048,518.07,128,34.80,33.007,127,34.87,334071180 +24576,2048,512.70,128,34.76,33.027,127,34.83,362259852 +26624,2048,503.79,128,34.63,32.795,127,34.70,390448524 +28672,2048,499.18,128,34.57,33.196,127,34.64,418637196 +30720,2048,494.72,128,34.47,33.246,127,34.54,446825868 +32768,2048,489.92,128,34.13,33.251,127,34.20,475014540 +34816,2048,482.51,128,33.81,33.568,127,33.88,503203212 +36864,2048,479.43,128,24.03,33.415,127,24.00,531391884 +38912,2048,438.82,128,21.62,45.645,127,21.63,559580556 +40960,2048,178.79,128,2.30,414.794,127,2.30,587769228 +43008,2048,175.60,128,3.49,491.360,127,3.51,615957900 +45056,2048,425.36,128,10.53,45.177,127,10.49,644146572 +47104,2048,246.46,128,26.80,260.694,127,28.14,672335244 +49152,2048,452.05,128,33.32,34.099,127,33.39,700523916 +51200,2048,445.11,128,33.14,34.970,127,33.21,728712588 +53248,2048,442.10,128,33.10,34.409,127,33.17,756901260 +55296,2048,438.93,128,32.89,35.231,127,32.96,785089932 +57344,2048,435.16,128,32.91,34.891,127,32.97,813278604 +59392,2048,428.50,128,32.74,35.300,127,32.81,841467276 +61440,2048,426.39,128,32.60,35.516,127,32.66,869655948 +63488,2048,417.20,128,32.33,35.660,127,32.39,897844620 +65536,2048,414.10,128,31.95,38.584,127,32.04,0 diff --git a/speed-bench/m3_ultra_mxfp4_rebased_r2.csv b/speed-bench/m3_ultra_mxfp4_rebased_r2.csv new file mode 100644 index 000000000..d045ee148 --- /dev/null +++ b/speed-bench/m3_ultra_mxfp4_rebased_r2.csv @@ -0,0 +1,10 @@ +ctx_tokens,prefill_tokens,prefill_tps,gen_tokens,gen_tps,gen_first_ms,gen_steady_tokens,gen_steady_tps,kvcache_bytes +32768,32768,552.18,128,34.10,33.792,127,34.17,475014540 +34816,2048,482.52,128,33.87,33.740,127,33.93,503203212 +36864,2048,479.43,128,33.85,33.983,127,33.92,531391884 +38912,2048,474.52,128,33.69,33.852,127,33.76,559580556 +40960,2048,470.10,128,33.62,34.175,127,33.69,587769228 +43008,2048,463.59,128,33.56,34.674,127,33.63,615957900 +45056,2048,459.43,128,33.41,34.405,127,33.49,644146572 +47104,2048,455.52,128,33.34,34.806,127,33.41,672335244 +49152,2048,451.45,128,33.31,38.422,127,33.41,0 diff --git a/speed-bench/m3_ultra_mxfp4_ts.svg b/speed-bench/m3_ultra_mxfp4_ts.svg new file mode 100644 index 000000000..eb97aa7d2 --- /dev/null +++ b/speed-bench/m3_ultra_mxfp4_ts.svg @@ -0,0 +1,50 @@ + + + + +M3 Ultra (512GB) MXFP4 t/s + +0 + +200 + +400 + +600 + +800 + +1k +0 +10 +20 +30 +40 + +0 + +20k + +40k + +60k + + + +ctx size +prefill t/s +generation t/s + + + + +prefill + +generation + diff --git a/speed-bench/metal_decode_schedule_bench.c b/speed-bench/metal_decode_schedule_bench.c index 5866b49db..575fc8944 100644 --- a/speed-bench/metal_decode_schedule_bench.c +++ b/speed-bench/metal_decode_schedule_bench.c @@ -418,8 +418,12 @@ int main(int argc, char **argv) { .len = cfg.prefix_tokens, .cap = cfg.prefix_tokens, }; + /* Build each session's persistent KV/compressor state under its own + * variant. This also makes the first frontier comparison cover prefill + * features instead of comparing two default-prefix sessions. */ for (int i = 0; i < VARIANT_COUNT; i++) { - if (ds4_session_create(&sessions[i], engine, cfg.ctx) != 0 || + if (select_variant(&cfg, i) != 0 || + ds4_session_create(&sessions[i], engine, cfg.ctx) != 0 || ds4_session_sync(sessions[i], &prefix, err, sizeof(err)) != 0) { fprintf(stderr, "metal-decode-schedule-bench: session %d prefill failed: %s\n", diff --git a/tests/ds4_test.c b/tests/ds4_test.c index cf8bca2c5..fd16f6735 100644 --- a/tests/ds4_test.c +++ b/tests/ds4_test.c @@ -386,6 +386,7 @@ static void test_fill_copy_f32_patterns(void *dst, uint32_t n, uint32_t salt) { memcpy(bytes + (uint64_t)i * sizeof(bits), &bits, sizeof(bits)); } } + #endif static uint16_t test_float_to_f16(float f) { @@ -863,6 +864,214 @@ static void test_metal_store_raw_kv_batch_wrap(void) { ds4_gpu_tensor_free(raw); } +#if defined(__APPLE__) +static void test_metal_batch_qkv_finalizer_exact_case( + uint32_t n_tokens, + uint32_t raw_cap, + uint32_t pos0, + uint32_t seed, + bool yarn, + bool batch_commands, + bool force_pos_buffer) { + const uint32_t n_head = 64u; + const uint32_t head_dim = 512u; + const uint32_t n_rot = 64u; + const uint64_t q_count = + (uint64_t)n_tokens * n_head * head_dim; + const uint64_t kv_count = (uint64_t)n_tokens * head_dim; + const uint64_t raw_count = (uint64_t)raw_cap * head_dim; + const uint64_t q_bytes = q_count * sizeof(float); + const uint64_t kv_bytes = kv_count * sizeof(float); + const uint64_t raw_bytes = raw_count * sizeof(float); + + ds4_gpu_tensor *ref_q = ds4_gpu_tensor_alloc(q_bytes); + ds4_gpu_tensor *fused_q = ds4_gpu_tensor_alloc(q_bytes); + ds4_gpu_tensor *ref_kv = ds4_gpu_tensor_alloc(kv_bytes); + ds4_gpu_tensor *fused_kv = ds4_gpu_tensor_alloc(kv_bytes); + ds4_gpu_tensor *ref_raw = ds4_gpu_tensor_alloc(raw_bytes); + ds4_gpu_tensor *fused_raw = ds4_gpu_tensor_alloc(raw_bytes); + float *q_input = malloc((size_t)q_bytes); + float *kv_input = malloc((size_t)kv_bytes); + float *raw_input = malloc((size_t)raw_bytes); + float *ref_host = malloc((size_t)q_bytes); + float *fused_host = malloc((size_t)q_bytes); + + TEST_ASSERT(ref_q != NULL); + TEST_ASSERT(fused_q != NULL); + TEST_ASSERT(ref_kv != NULL); + TEST_ASSERT(fused_kv != NULL); + TEST_ASSERT(ref_raw != NULL); + TEST_ASSERT(fused_raw != NULL); + TEST_ASSERT(q_input != NULL); + TEST_ASSERT(kv_input != NULL); + TEST_ASSERT(raw_input != NULL); + TEST_ASSERT(ref_host != NULL); + TEST_ASSERT(fused_host != NULL); + + const bool allocated = ref_q && fused_q && ref_kv && fused_kv && + ref_raw && fused_raw && q_input && kv_input && raw_input && + ref_host && fused_host; + if (allocated) { + for (uint64_t i = 0; i < q_count; i++) { + const int value = + (int)((i * 29u + (i >> 5u) * 17u + seed * 31u) % 509u) - + 254; + q_input[i] = (float)value / 47.0f; + } + for (uint64_t i = 0; i < kv_count; i++) { + const int value = + (int)((i * 37u + (i >> 4u) * 13u + seed * 19u) % 521u) - + 260; + kv_input[i] = (float)value / 53.0f; + } + const uint32_t negative_zero = 0x80000000u; + memcpy(q_input + (seed % q_count), &negative_zero, + sizeof(negative_zero)); + memcpy(kv_input + (seed % kv_count), &negative_zero, + sizeof(negative_zero)); + for (uint64_t i = 0; i < raw_count; i++) { + const uint32_t poison = + 0x7fc00001u + (uint32_t)(i & 0x3ffu); + memcpy(raw_input + i, &poison, sizeof(poison)); + } + + TEST_ASSERT(ds4_gpu_tensor_write( + ref_q, 0, q_input, q_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_q, 0, q_input, q_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + ref_kv, 0, kv_input, kv_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_kv, 0, kv_input, kv_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + ref_raw, 0, raw_input, raw_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_raw, 0, raw_input, raw_bytes) != 0); + + const float freq_base = yarn ? 160000.0f : 10000.0f; + const float freq_scale = yarn ? 1.0f / 16.0f : 1.0f; + const float ext_factor = yarn ? 1.0f : 0.0f; + const uint32_t n_ctx_orig = yarn ? 65536u : 0u; + const float attn_factor = yarn + ? 1.0f / (1.0f + 0.1f * logf(1.0f / freq_scale)) + : 1.0f; + const float beta_fast = 32.0f; + const float beta_slow = 1.0f; + const float eps = 1.0e-6f; + + TEST_ASSERT(ds4_gpu_head_rms_norm_rope_tail_tensor( + ref_q, n_tokens, n_head, head_dim, n_rot, pos0, + n_ctx_orig, + false, freq_base, freq_scale, ext_factor, + attn_factor, beta_fast, beta_slow, eps) != 0); + TEST_ASSERT(ds4_gpu_rope_tail_tensor( + ref_kv, n_tokens, 1u, head_dim, n_rot, pos0, + n_ctx_orig, + false, freq_base, freq_scale, ext_factor, + attn_factor, beta_fast, beta_slow) != 0); + TEST_ASSERT(ds4_gpu_dsv4_fp8_kv_quantize_tensor( + ref_kv, n_tokens, head_dim, n_rot) != 0); + TEST_ASSERT(ds4_gpu_store_raw_kv_batch_tensor( + ref_raw, ref_kv, raw_cap, pos0, + n_tokens, head_dim) != 0); + + const char *pos_buffer_env = + "DS4_GPU_TEST_BATCH_QKV_FINALIZE_POS_BUFFER"; + char *saved_pos_buffer_env = force_pos_buffer + ? test_save_env(pos_buffer_env) : NULL; + if (force_pos_buffer) { + TEST_ASSERT(setenv(pos_buffer_env, "1", 1) == 0); + } + const int begun = batch_commands ? ds4_gpu_begin_commands() : 1; + TEST_ASSERT(begun != 0); + const int finalized = begun + ? ds4_gpu_dsv4_batch_qnorm_rope_kv_finalize_tensor( + fused_q, fused_kv, fused_raw, raw_cap, n_tokens, + n_head, head_dim, n_rot, pos0, n_ctx_orig, false, + freq_base, freq_scale, ext_factor, attn_factor, + beta_fast, beta_slow, eps) + : 0; + const int ended = batch_commands && begun + ? ds4_gpu_end_commands() : begun; + if (force_pos_buffer) { + test_restore_env(pos_buffer_env, saved_pos_buffer_env); + } + TEST_ASSERT(finalized != 0); + TEST_ASSERT(ended != 0); + + TEST_ASSERT(ds4_gpu_tensor_read( + ref_q, 0, ref_host, q_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + fused_q, 0, fused_host, q_bytes) != 0); + const test_float_compare_stats q_stats = + test_compare_float_bits(ref_host, fused_host, (size_t)q_count); + + TEST_ASSERT(ds4_gpu_tensor_read( + ref_kv, 0, ref_host, kv_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + fused_kv, 0, fused_host, kv_bytes) != 0); + const test_float_compare_stats kv_stats = + test_compare_float_bits(ref_host, fused_host, (size_t)kv_count); + + TEST_ASSERT(ds4_gpu_tensor_read( + ref_raw, 0, ref_host, raw_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + fused_raw, 0, fused_host, raw_bytes) != 0); + const test_float_compare_stats raw_stats = + test_compare_float_bits(ref_host, fused_host, (size_t)raw_count); + + fprintf(stderr, + "ds4-test: batch Q/KV finalizer rows=%u yarn=%d batch=%d " + "posbuf=%d " + "q=%zu kv=%zu raw=%zu " + "max_ulp=%u/%u/%u\n", + n_tokens, + yarn ? 1 : 0, + batch_commands ? 1 : 0, + force_pos_buffer ? 1 : 0, + q_stats.mismatch_count, + kv_stats.mismatch_count, + raw_stats.mismatch_count, + q_stats.max_ulp, + kv_stats.max_ulp, + raw_stats.max_ulp); + TEST_ASSERT(q_stats.mismatch_count == 0); + TEST_ASSERT(kv_stats.mismatch_count == 0); + TEST_ASSERT(raw_stats.mismatch_count == 0); + } + + free(fused_host); + free(ref_host); + free(raw_input); + free(kv_input); + free(q_input); + ds4_gpu_tensor_free(fused_raw); + ds4_gpu_tensor_free(ref_raw); + ds4_gpu_tensor_free(fused_kv); + ds4_gpu_tensor_free(ref_kv); + ds4_gpu_tensor_free(fused_q); + ds4_gpu_tensor_free(ref_q); +} + +static void test_metal_batch_qkv_finalizer_exact(void) { + if (!ds4_gpu_dsv4_batch_qnorm_rope_kv_finalize_available()) { + fprintf(stderr, + "ds4-test: batch Q/KV finalizer unavailable; skipping exact oracle\n"); + return; + } + test_metal_batch_qkv_finalizer_exact_case( + 7u, 13u, 10u, 17u, false, false, false); + test_metal_batch_qkv_finalizer_exact_case( + 33u, 43u, 39u, 29u, false, false, false); + test_metal_batch_qkv_finalizer_exact_case( + 17u, 29u, 32761u, 43u, true, false, false); + /* Production-sized admission, the batch-owned concurrent encoder, and + * the same MTLBuffer binding path used by >1024-token position arrays. */ + test_metal_batch_qkv_finalizer_exact_case( + 128u, 139u, 131u, 59u, true, true, true); +} +#endif + static void test_dspark_cache_window_crop(void) { TEST_ASSERT(ds4_test_dspark_cache_window_crop()); } @@ -3121,6 +3330,185 @@ static void test_metal_contiguous_compressed_f16_attention_exact(void) { ds4_gpu_tensor_free(raw); } +static void test_metal_indexed_attention_prefill_heads16_dual_rb4_exact(void) { + const uint32_t n_tokens = 32; + const uint32_t n_head = 64; + const uint32_t head_dim = 512; + const uint32_t top_k = 512; + const uint32_t n_comp = 520; + const uint32_t ratio = 4; + const uint32_t window = 128; + const uint32_t pos0 = 2035; + const uint32_t n_raw = 128; + const uint32_t raw_cap = 137; + const uint32_t raw_start = 133; + const uint64_t q_count = + (uint64_t)n_tokens * n_head * head_dim; + const uint64_t raw_count = (uint64_t)raw_cap * head_dim; + const uint64_t comp_count = (uint64_t)n_comp * head_dim; + const uint64_t topk_count = (uint64_t)n_tokens * top_k; + const uint64_t q_bytes = q_count * sizeof(float); + const uint64_t raw_bytes = raw_count * sizeof(float); + const uint64_t comp_bytes = comp_count * sizeof(uint16_t); + const uint64_t topk_bytes = topk_count * sizeof(int32_t); + const uint64_t page = (uint64_t)getpagesize(); + const char *disable_env = + "DS4_METAL_DISABLE_PRE_M5_INDEXED_ATTN_PREFILL_RB4"; + const char *require_env = + "DS4_METAL_REQUIRE_PRE_M5_INDEXED_ATTN_PREFILL_RB4"; + const char *dual_disable_env = + "DS4_METAL_DISABLE_PRE_M5_INDEXED_ATTN_PREFILL_HEADS16_DUAL_RB4"; + const char *dual_require_env = + "DS4_METAL_REQUIRE_PRE_M5_INDEXED_ATTN_PREFILL_HEADS16_DUAL_RB4"; + char *saved_disable = test_save_env(disable_env); + char *saved_require = test_save_env(require_env); + char *saved_dual_disable = test_save_env(dual_disable_env); + char *saved_dual_require = test_save_env(dual_require_env); + const int saved_quality = ds4_gpu_test_get_quality(); + + void *model_raw = NULL; + TEST_ASSERT(posix_memalign(&model_raw, (size_t)page, (size_t)page) == 0); + ds4_gpu_tensor *q = ds4_gpu_tensor_alloc(q_bytes); + ds4_gpu_tensor *raw = ds4_gpu_tensor_alloc(raw_bytes); + ds4_gpu_tensor *comp = ds4_gpu_tensor_alloc(comp_bytes); + ds4_gpu_tensor *topk = ds4_gpu_tensor_alloc(topk_bytes); + ds4_gpu_tensor *reference = ds4_gpu_tensor_alloc(q_bytes); + ds4_gpu_tensor *candidate = ds4_gpu_tensor_alloc(q_bytes); + float *q_host = malloc((size_t)q_bytes); + float *raw_host = malloc((size_t)raw_bytes); + uint16_t *comp_host = malloc((size_t)comp_bytes); + int32_t *topk_host = malloc((size_t)topk_bytes); + float *reference_host = malloc((size_t)q_bytes); + float *candidate_host = malloc((size_t)q_bytes); + TEST_ASSERT(model_raw != NULL); + TEST_ASSERT(q != NULL); + TEST_ASSERT(raw != NULL); + TEST_ASSERT(comp != NULL); + TEST_ASSERT(topk != NULL); + TEST_ASSERT(reference != NULL); + TEST_ASSERT(candidate != NULL); + TEST_ASSERT(q_host != NULL); + TEST_ASSERT(raw_host != NULL); + TEST_ASSERT(comp_host != NULL); + TEST_ASSERT(topk_host != NULL); + TEST_ASSERT(reference_host != NULL); + TEST_ASSERT(candidate_host != NULL); + + const bool allocated = model_raw && q && raw && comp && topk && + reference && candidate && q_host && raw_host && comp_host && + topk_host && reference_host && candidate_host; + test_float_compare_stats stats = {0}; + if (allocated) { + memset(model_raw, 0, (size_t)page); + float *sinks = model_raw; + for (uint32_t head = 0; head < n_head; head++) { + const int value = (int)((head * 29u + 7u) % 61u) - 30; + sinks[head] = (float)value / 32.0f; + } + for (uint64_t i = 0; i < q_count; i++) { + const int value = + (int)((i * 37u + (i ^ (i >> 5u)) * 11u) % 251u) - 125; + q_host[i] = (float)value / 128.0f; + } + for (uint64_t i = 0; i < raw_count; i++) { + const int value = + (int)((i * 19u + (i ^ (i >> 4u)) * 7u) % 233u) - 116; + raw_host[i] = (float)value / 128.0f; + } + for (uint64_t i = 0; i < comp_count; i++) { + const int value = + (int)((i * 23u + (i ^ (i >> 3u)) * 13u) % 227u) - 113; + comp_host[i] = test_float_to_f16((float)value / 128.0f); + } + for (uint32_t token = 0; token < n_tokens; token++) { + uint32_t visible = (pos0 + token + 1u) / ratio; + if (visible > n_comp) visible = n_comp; + const uint32_t valid = visible < top_k ? visible : top_k; + for (uint32_t i = 0; i < top_k; i++) { + const uint32_t desc = top_k - 1u - i; + const uint32_t index = desc < valid + ? desc + : visible + desc - valid; + topk_host[(uint64_t)token * top_k + i] = (int32_t)index; + } + } + memset(reference_host, 0xa5, (size_t)q_bytes); + memset(candidate_host, 0x5a, (size_t)q_bytes); + TEST_ASSERT(ds4_gpu_tensor_write(q, 0, q_host, q_bytes) != 0); + 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(topk, 0, topk_host, topk_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + reference, 0, reference_host, q_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + candidate, 0, candidate_host, q_bytes) != 0); + TEST_ASSERT(ds4_gpu_set_model_map(model_raw, page) != 0); + + /* REQUIRE plus the dedicated rollback proves that the reference is + * the accepted heads8 RB4 path, not the new heads16 specialization. */ + ds4_gpu_set_quality(false); + ds4_gpu_test_set_flags(DS4_GPU_TEST_INDEXED_ATTN_PREFILL_RB4); + TEST_ASSERT(unsetenv(disable_env) == 0); + TEST_ASSERT(setenv(require_env, "1", 1) == 0); + TEST_ASSERT(setenv(dual_disable_env, "1", 1) == 0); + TEST_ASSERT(setenv(dual_require_env, "1", 1) == 0); + TEST_ASSERT(ds4_gpu_attention_indexed_mixed_batch_heads_tensor( + reference, model_raw, page, 0, q, raw, comp, 1, topk, + n_tokens, pos0, n_raw, raw_cap, raw_start, n_comp, top_k, + window, ratio, n_head, head_dim) == 0); + TEST_ASSERT(unsetenv(dual_require_env) == 0); + TEST_ASSERT(ds4_gpu_attention_indexed_mixed_batch_heads_tensor( + reference, model_raw, page, 0, q, raw, comp, 1, topk, + n_tokens, pos0, n_raw, raw_cap, raw_start, n_comp, top_k, + window, ratio, n_head, head_dim) != 0); + + /* Removing only the new rollback must select heads16 dual RB4. */ + TEST_ASSERT(unsetenv(dual_disable_env) == 0); + TEST_ASSERT(setenv(dual_require_env, "1", 1) == 0); + TEST_ASSERT(ds4_gpu_attention_indexed_mixed_batch_heads_tensor( + candidate, model_raw, page, 0, q, raw, comp, 1, topk, + n_tokens, pos0, n_raw, raw_cap, raw_start, n_comp, top_k, + window, ratio, n_head, head_dim) != 0); + + TEST_ASSERT(ds4_gpu_tensor_read( + reference, 0, reference_host, q_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + candidate, 0, candidate_host, q_bytes) != 0); + stats = test_compare_float_bits( + reference_host, candidate_host, (size_t)q_count); + } + + ds4_gpu_test_set_flags(0); + ds4_gpu_set_quality(saved_quality != 0); + test_restore_env(dual_require_env, saved_dual_require); + test_restore_env(dual_disable_env, saved_dual_disable); + test_restore_env(require_env, saved_require); + test_restore_env(disable_env, saved_disable); + fprintf(stderr, + "ds4-test: indexed-attention prefill heads16 dual RB4 exactness " + "mismatches=%zu/%llu max_ulp=%u max_abs=%g\n", + stats.mismatch_count, + (unsigned long long)q_count, + stats.max_ulp, + stats.max_abs); + TEST_ASSERT(stats.mismatch_count == 0); + TEST_ASSERT(stats.max_ulp == 0); + + free(candidate_host); + free(reference_host); + free(topk_host); + free(comp_host); + free(raw_host); + free(q_host); + ds4_gpu_tensor_free(candidate); + ds4_gpu_tensor_free(reference); + ds4_gpu_tensor_free(topk); + ds4_gpu_tensor_free(comp); + ds4_gpu_tensor_free(raw); + ds4_gpu_tensor_free(q); + free(model_raw); +} + static void test_metal_persistent_zero_attention_mask_exact_case( uint32_t raw_cap, uint32_t n_raw, @@ -3641,6 +4029,419 @@ static void test_metal_zero_prefix_prefill_mask_cache_exact(void) { #endif #if defined(__APPLE__) +static void test_metal_moe_sum6_hc_expand_exact_case( + uint32_t n_tokens, + uint32_t n_embd, + uint32_t seed) { + const uint32_t n_hc = 4; + const uint32_t guard_count = 32; + const uint32_t ref_poison = 0x7fc0a5a5u; + const uint32_t fused_poison = 0x7fc05a5au; + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + const uint64_t expert_count = (uint64_t)n_tokens * 6u * n_embd; + const uint64_t routed_count = (uint64_t)n_tokens * n_embd; + const uint64_t shared_count = routed_count; + const uint64_t residual_count = (uint64_t)n_tokens * n_hc * n_embd; + const uint64_t split_count = (uint64_t)n_tokens * mix_hc; + const uint64_t out_count = residual_count; + const uint64_t guarded_out_count = out_count + guard_count; + const uint64_t expert_bytes = expert_count * sizeof(float); + const uint64_t routed_bytes = routed_count * sizeof(float); + const uint64_t shared_bytes = shared_count * sizeof(float); + const uint64_t residual_bytes = residual_count * sizeof(float); + const uint64_t split_bytes = split_count * sizeof(float); + const uint64_t out_bytes = out_count * sizeof(float); + const uint64_t guarded_out_bytes = guarded_out_count * sizeof(float); + + ds4_gpu_tensor *expert_down = ds4_gpu_tensor_alloc(expert_bytes); + ds4_gpu_tensor *routed = ds4_gpu_tensor_alloc(routed_bytes); + ds4_gpu_tensor *shared = ds4_gpu_tensor_alloc(shared_bytes); + ds4_gpu_tensor *residual = ds4_gpu_tensor_alloc(residual_bytes); + ds4_gpu_tensor *split = ds4_gpu_tensor_alloc(split_bytes); + ds4_gpu_tensor *ref_base = ds4_gpu_tensor_alloc(guarded_out_bytes); + ds4_gpu_tensor *fused_base = ds4_gpu_tensor_alloc(guarded_out_bytes); + ds4_gpu_tensor *ref = ref_base ? + ds4_gpu_tensor_view(ref_base, 0, out_bytes) : NULL; + ds4_gpu_tensor *fused = fused_base ? + ds4_gpu_tensor_view(fused_base, 0, out_bytes) : NULL; + float *expert_host = malloc((size_t)expert_bytes); + float *shared_host = malloc((size_t)shared_bytes); + float *residual_host = malloc((size_t)residual_bytes); + float *split_host = malloc((size_t)split_bytes); + float *ref_host = malloc((size_t)guarded_out_bytes); + float *fused_host = malloc((size_t)guarded_out_bytes); + TEST_ASSERT(expert_down != NULL); + TEST_ASSERT(routed != NULL); + TEST_ASSERT(shared != NULL); + TEST_ASSERT(residual != NULL); + TEST_ASSERT(split != NULL); + TEST_ASSERT(ref_base != NULL && ref != NULL); + TEST_ASSERT(fused_base != NULL && fused != NULL); + TEST_ASSERT(expert_host != NULL); + TEST_ASSERT(shared_host != NULL); + TEST_ASSERT(residual_host != NULL); + TEST_ASSERT(split_host != NULL); + TEST_ASSERT(ref_host != NULL); + TEST_ASSERT(fused_host != NULL); + + const bool allocated = expert_down && routed && shared && residual && split && + ref_base && ref && fused_base && fused && expert_host && shared_host && + residual_host && split_host && ref_host && fused_host; + test_float_compare_stats stats = {0}; + size_t guard_mismatches = 0; + size_t ref_poisoned = 0; + size_t fused_poisoned = 0; + if (allocated) { + for (uint32_t t = 0; t < n_tokens; t++) { + for (uint32_t d = 0; d < n_embd; d++) { + const uint64_t cell = (uint64_t)t * n_embd + d; + const int base_i = + (int)((d * 37u + t * 101u + seed * 17u) % 127u) - 63; + const float large = (float)base_i * 256.0f; + const int small_i = + (int)((d * 19u + t * 43u + seed * 11u) % 31u) - 15; + expert_host[((uint64_t)t * 6u + 0u) * n_embd + d] = large; + expert_host[((uint64_t)t * 6u + 1u) * n_embd + d] = + (float)small_i / 64.0f; + expert_host[((uint64_t)t * 6u + 2u) * n_embd + d] = -large; + expert_host[((uint64_t)t * 6u + 3u) * n_embd + d] = + (float)((int)((d * 7u + seed) % 23u) - 11) / 8.0f; + expert_host[((uint64_t)t * 6u + 4u) * n_embd + d] = + (float)((int)((d * 13u + t) % 29u) - 14) / 16.0f; + expert_host[((uint64_t)t * 6u + 5u) * n_embd + d] = + (float)((int)((d * 5u + t * 3u + seed) % 19u) - 9) / 32.0f; + shared_host[cell] = + (float)((int)((d * 11u + t * 17u + seed) % 41u) - 20) / 16.0f; + } + } + for (uint64_t i = 0; i < residual_count; i++) { + const int value = + (int)((i * 29u + (i ^ (i >> 4u)) * 7u + seed) % 97u) - 48; + residual_host[i] = (float)value / 32.0f; + } + for (uint64_t i = 0; i < split_count; i++) { + const int value = + (int)((i * 17u + (i ^ (i >> 3u)) * 5u + seed * 3u) % 61u) - 30; + split_host[i] = (float)value / 32.0f; + } + for (uint64_t i = 0; i < guarded_out_count; i++) { + memcpy(ref_host + i, &ref_poison, sizeof(ref_poison)); + memcpy(fused_host + i, &fused_poison, sizeof(fused_poison)); + } + + TEST_ASSERT(ds4_gpu_tensor_write( + expert_down, 0, expert_host, expert_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + shared, 0, shared_host, shared_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + residual, 0, residual_host, residual_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + split, 0, split_host, split_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + ref_base, 0, ref_host, guarded_out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_base, 0, fused_host, guarded_out_bytes) != 0); + TEST_ASSERT(ds4_gpu_moe_sum6_hc_expand_available() != 0); + + TEST_ASSERT(ds4_gpu_test_moe_sum6_tensor( + routed, expert_down, n_embd, n_tokens) != 0); + TEST_ASSERT(ds4_gpu_hc_expand_add_split_tensor( + ref, routed, shared, residual, split, + n_embd, n_hc) != 0); + TEST_ASSERT(ds4_gpu_moe_sum6_hc_expand_split_tensor( + fused, expert_down, shared, residual, split, + n_embd, n_hc) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + ref_base, 0, ref_host, guarded_out_bytes) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + fused_base, 0, fused_host, guarded_out_bytes) != 0); + + stats = test_compare_float_bits(ref_host, fused_host, (size_t)out_count); + for (uint64_t i = 0; i < out_count; i++) { + uint32_t ref_bits = 0; + uint32_t fused_bits = 0; + memcpy(&ref_bits, ref_host + i, sizeof(ref_bits)); + memcpy(&fused_bits, fused_host + i, sizeof(fused_bits)); + if (ref_bits == ref_poison) ref_poisoned++; + if (fused_bits == fused_poison) fused_poisoned++; + } + for (uint32_t i = 0; i < guard_count; i++) { + uint32_t ref_bits = 0; + uint32_t fused_bits = 0; + memcpy(&ref_bits, ref_host + out_count + i, sizeof(ref_bits)); + memcpy(&fused_bits, fused_host + out_count + i, sizeof(fused_bits)); + if (ref_bits != ref_poison || fused_bits != fused_poison) { + guard_mismatches++; + } + } + } + + fprintf(stderr, + "ds4-test: MoE sum6 HC exact tokens=%u width=%u " + "mismatches=%zu/%llu max_ulp=%u max_abs=%g poison=%zu/%zu guard=%zu\n", + n_tokens, + n_embd, + stats.mismatch_count, + (unsigned long long)out_count, + stats.max_ulp, + stats.max_abs, + ref_poisoned, + fused_poisoned, + guard_mismatches); + TEST_ASSERT(stats.mismatch_count == 0); + TEST_ASSERT(stats.max_ulp == 0); + TEST_ASSERT(ref_poisoned == 0); + TEST_ASSERT(fused_poisoned == 0); + TEST_ASSERT(guard_mismatches == 0); + + free(fused_host); + free(ref_host); + free(split_host); + free(residual_host); + free(shared_host); + free(expert_host); + ds4_gpu_tensor_free(fused); + ds4_gpu_tensor_free(ref); + ds4_gpu_tensor_free(fused_base); + ds4_gpu_tensor_free(ref_base); + ds4_gpu_tensor_free(split); + ds4_gpu_tensor_free(residual); + ds4_gpu_tensor_free(shared); + ds4_gpu_tensor_free(routed); + ds4_gpu_tensor_free(expert_down); +} + +static void test_metal_moe_sum6_hc_expand_exact(void) { + test_metal_moe_sum6_hc_expand_exact_case(3, 257, 17); + test_metal_moe_sum6_hc_expand_exact_case(32, 4096, 31); +} + +static void test_metal_batch_attn_out_hc_fusion_exact(void) { + const uint32_t n_tokens = 32u; + const uint32_t group_dim = 4096u; + const uint32_t rank = 1024u; + const uint32_t n_groups = 8u; + const uint32_t low_dim = n_groups * rank; + const uint32_t out_dim = 4096u; + const uint32_t n_hc = 4u; + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + const uint64_t row_a_bytes = (uint64_t)(group_dim / 32u) * 34u; + const uint64_t row_b_bytes = (uint64_t)(low_dim / 32u) * 34u; + const uint64_t out_a_bytes = (uint64_t)n_groups * rank * row_a_bytes; + const uint64_t out_b_offset = test_round_up_u64(out_a_bytes, 64u); + const uint64_t out_b_bytes = (uint64_t)out_dim * row_b_bytes; + const uint64_t page = (uint64_t)getpagesize(); + const uint64_t model_bytes = + test_round_up_u64(out_b_offset + out_b_bytes, page); + const uint64_t heads_count = + (uint64_t)n_tokens * n_groups * group_dim; + const uint64_t low_count = (uint64_t)n_tokens * low_dim; + const uint64_t out_count = (uint64_t)n_tokens * out_dim; + const uint64_t hc_count = (uint64_t)n_tokens * n_hc * out_dim; + const uint64_t split_count = (uint64_t)n_tokens * mix_hc; + + void *model_raw = NULL; + TEST_ASSERT(posix_memalign(&model_raw, (size_t)page, + (size_t)model_bytes) == 0); + float *heads_host = malloc((size_t)heads_count * sizeof(float)); + float *residual_host = malloc((size_t)hc_count * sizeof(float)); + float *split_host = malloc((size_t)split_count * sizeof(float)); + float *poison_out_host = malloc((size_t)out_count * sizeof(float)); + float *ref_low_host = malloc((size_t)low_count * sizeof(float)); + float *fused_low_host = malloc((size_t)low_count * sizeof(float)); + float *ref_hc_host = malloc((size_t)hc_count * sizeof(float)); + float *fused_hc_host = malloc((size_t)hc_count * sizeof(float)); + + ds4_gpu_tensor *heads = ds4_gpu_tensor_alloc(heads_count * sizeof(float)); + ds4_gpu_tensor *ref_low = ds4_gpu_tensor_alloc(low_count * sizeof(float)); + ds4_gpu_tensor *fused_low = ds4_gpu_tensor_alloc(low_count * sizeof(float)); + ds4_gpu_tensor *ref_out = ds4_gpu_tensor_alloc(out_count * sizeof(float)); + ds4_gpu_tensor *fused_out = ds4_gpu_tensor_alloc(out_count * sizeof(float)); + ds4_gpu_tensor *residual = ds4_gpu_tensor_alloc(hc_count * sizeof(float)); + ds4_gpu_tensor *split = ds4_gpu_tensor_alloc(split_count * sizeof(float)); + ds4_gpu_tensor *ref_hc = ds4_gpu_tensor_alloc(hc_count * sizeof(float)); + ds4_gpu_tensor *fused_hc = ds4_gpu_tensor_alloc(hc_count * sizeof(float)); + ds4_gpu_tensor *group_tmp = ds4_gpu_tensor_alloc(sizeof(float)); + ds4_gpu_tensor *low_tmp = ds4_gpu_tensor_alloc(sizeof(float)); + + const bool allocated = model_raw && heads_host && residual_host && + split_host && poison_out_host && ref_low_host && fused_low_host && + ref_hc_host && fused_hc_host && heads && ref_low && fused_low && + ref_out && fused_out && residual && split && ref_hc && fused_hc && + group_tmp && low_tmp; + TEST_ASSERT(allocated); + + const char *disable_env = + "DS4_METAL_DISABLE_PRE_M5_BATCH_ATTN_OUT_HC_FUSION"; + const char *require_env = + "DS4_METAL_REQUIRE_PRE_M5_BATCH_ATTN_OUT_HC_FUSION"; + char *saved_disable = test_save_env(disable_env); + char *saved_require = test_save_env(require_env); + const int saved_quality = ds4_gpu_test_get_quality(); + test_float_compare_stats low_stats = {0}; + test_float_compare_stats hc_stats = {0}; + size_t out_poison_mismatches = 0; + + if (allocated) { + memset(model_raw, 0, (size_t)model_bytes); + uint8_t *model_u8 = model_raw; + for (uint32_t row = 0; row < n_groups * rank; ++row) { + for (uint32_t block = 0; block < group_dim / 32u; ++block) { + uint8_t *q = model_u8 + (uint64_t)row * row_a_bytes + + (uint64_t)block * 34u; + const float scale = + (float)(1u + ((row * 3u + block * 5u) % 7u)) / 128.0f; + const uint16_t d = test_float_to_f16(scale); + memcpy(q, &d, sizeof(d)); + for (uint32_t j = 0; j < 32u; ++j) { + q[2u + j] = (uint8_t)(int8_t) + ((int)((row * 11u + block * 7u + j * 5u) % 15u) - 7); + } + } + } + for (uint32_t row = 0; row < out_dim; ++row) { + for (uint32_t block = 0; block < low_dim / 32u; ++block) { + uint8_t *q = model_u8 + out_b_offset + + (uint64_t)row * row_b_bytes + (uint64_t)block * 34u; + const float scale = + (float)(1u + ((row * 7u + block * 3u) % 5u)) / 192.0f; + const uint16_t d = test_float_to_f16(scale); + memcpy(q, &d, sizeof(d)); + for (uint32_t j = 0; j < 32u; ++j) { + q[2u + j] = (uint8_t)(int8_t) + ((int)((row * 5u + block * 13u + j * 3u) % 13u) - 6); + } + } + } + for (uint64_t i = 0; i < heads_count; ++i) { + const int v = (int)((i * 17u + (i ^ (i >> 7u)) * 3u) % 127u) - 63; + heads_host[i] = (float)v / 96.0f; + } + for (uint64_t i = 0; i < hc_count; ++i) { + const int v = (int)((i * 19u + (i ^ (i >> 5u)) * 7u) % 149u) - 74; + residual_host[i] = (float)v / 80.0f; + } + for (uint32_t t = 0; t < n_tokens; ++t) { + float *row = split_host + (uint64_t)t * mix_hc; + for (uint32_t h = 0; h < n_hc; ++h) { + row[h] = 0.0f; + row[n_hc + h] = 0.55f + (float)((t + h * 3u) % 11u) / 32.0f; + } + for (uint32_t dst_hc = 0; dst_hc < n_hc; ++dst_hc) { + for (uint32_t src_hc = 0; src_hc < n_hc; ++src_hc) { + const int v = (int)((t * 5u + dst_hc * 7u + + src_hc * 11u) % 17u) - 8; + row[2u * n_hc + dst_hc * n_hc + src_hc] = + (float)v / 24.0f; + } + } + } + uint32_t *poison_bits = (uint32_t *)poison_out_host; + for (uint64_t i = 0; i < out_count; ++i) { + poison_bits[i] = 0x7fc12345u; + } + + TEST_ASSERT(ds4_gpu_tensor_write( + heads, 0, heads_host, + heads_count * sizeof(float)) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + residual, 0, residual_host, + hc_count * sizeof(float)) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + split, 0, split_host, + split_count * sizeof(float)) != 0); + TEST_ASSERT(ds4_gpu_tensor_write( + fused_out, 0, poison_out_host, + out_count * sizeof(float)) != 0); + TEST_ASSERT(ds4_gpu_set_model_map(model_raw, model_bytes) != 0); + + ds4_gpu_set_quality(true); + ds4_gpu_test_set_flags(DS4_GPU_TEST_BATCH_ATTN_OUT_HC_FUSION); + TEST_ASSERT(setenv(disable_env, "1", 1) == 0); + TEST_ASSERT(setenv(require_env, "1", 1) == 0); + const int rejected = ds4_gpu_attention_output_q8_batch_hc_tensor( + fused_out, fused_hc, residual, split, fused_low, + group_tmp, low_tmp, model_raw, model_bytes, 0u, out_b_offset, + group_dim, rank, n_groups, out_dim, heads, n_tokens, n_hc); + TEST_ASSERT(rejected == 0); + + TEST_ASSERT(ds4_gpu_attention_output_q8_batch_tensor( + ref_out, ref_low, group_tmp, low_tmp, model_raw, model_bytes, + 0u, out_b_offset, group_dim, rank, n_groups, out_dim, + heads, n_tokens) != 0); + TEST_ASSERT(ds4_gpu_hc_expand_split_tensor( + ref_hc, ref_out, residual, split, out_dim, n_hc) != 0); + + TEST_ASSERT(unsetenv(disable_env) == 0); + const int fused_ok = ds4_gpu_attention_output_q8_batch_hc_tensor( + fused_out, fused_hc, residual, split, fused_low, + group_tmp, low_tmp, model_raw, model_bytes, 0u, out_b_offset, + group_dim, rank, n_groups, out_dim, heads, n_tokens, n_hc); + TEST_ASSERT(fused_ok == 1); + + TEST_ASSERT(ds4_gpu_tensor_read( + ref_low, 0, ref_low_host, + low_count * sizeof(float)) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + fused_low, 0, fused_low_host, + low_count * sizeof(float)) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + ref_hc, 0, ref_hc_host, + hc_count * sizeof(float)) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + fused_hc, 0, fused_hc_host, + hc_count * sizeof(float)) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + fused_out, 0, poison_out_host, + out_count * sizeof(float)) != 0); + + low_stats = test_compare_float_bits( + ref_low_host, fused_low_host, (size_t)low_count); + hc_stats = test_compare_float_bits( + ref_hc_host, fused_hc_host, (size_t)hc_count); + for (uint64_t i = 0; i < out_count; ++i) { + if (((uint32_t *)poison_out_host)[i] != 0x7fc12345u) { + out_poison_mismatches++; + } + } + } + + ds4_gpu_test_set_flags(0u); + ds4_gpu_set_quality(saved_quality != 0); + test_restore_env(disable_env, saved_disable); + test_restore_env(require_env, saved_require); + fprintf(stderr, + "ds4-test: batch attention-output HC fusion exact " + "low=%zu/%llu hc=%zu/%llu max_ulp=%u/%u out_writes=%zu\n", + low_stats.mismatch_count, (unsigned long long)low_count, + hc_stats.mismatch_count, (unsigned long long)hc_count, + low_stats.max_ulp, hc_stats.max_ulp, out_poison_mismatches); + TEST_ASSERT(low_stats.mismatch_count == 0); + TEST_ASSERT(hc_stats.mismatch_count == 0); + TEST_ASSERT(out_poison_mismatches == 0); + + ds4_gpu_tensor_free(low_tmp); + ds4_gpu_tensor_free(group_tmp); + ds4_gpu_tensor_free(fused_hc); + ds4_gpu_tensor_free(ref_hc); + ds4_gpu_tensor_free(split); + ds4_gpu_tensor_free(residual); + ds4_gpu_tensor_free(fused_out); + ds4_gpu_tensor_free(ref_out); + ds4_gpu_tensor_free(fused_low); + ds4_gpu_tensor_free(ref_low); + ds4_gpu_tensor_free(heads); + free(fused_hc_host); + free(ref_hc_host); + free(fused_low_host); + free(ref_low_host); + free(poison_out_host); + free(split_host); + free(residual_host); + free(heads_host); + free(model_raw); +} + 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. */ @@ -4730,6 +5531,7 @@ static void test_metal_kernel_group(void) { test_dspark_cache_window_crop(); test_metal_q8_0_decode_pair_exact(); #if defined(__APPLE__) + test_metal_batch_qkv_finalizer_exact(); test_metal_f16_compressor_pair_state_store_exact(); test_metal_compressor_ape_add_exact(); test_metal_compressor_ratio4_pack_exact(); @@ -4739,8 +5541,11 @@ static void test_metal_kernel_group(void) { test_metal_contiguous_f32_f16_roundtrip_exact(); test_metal_gathered_kv_stage_exact(); test_metal_contiguous_compressed_f16_attention_exact(); + test_metal_indexed_attention_prefill_heads16_dual_rb4_exact(); test_metal_persistent_zero_attention_mask_exact(); test_metal_zero_prefix_prefill_mask_cache_exact(); + test_metal_moe_sum6_hc_expand_exact(); + test_metal_batch_attn_out_hc_fusion_exact(); test_metal_hc_split_weighted_sum_norm_batch_exact(); test_metal_output_hc_weights4_exact(); test_metal_hc_rms_scale_project_f16_exact(); diff --git a/tests/test_mxfp4_metal.c b/tests/test_mxfp4_metal.c index bb1f37fed..f5d5a9cc3 100644 --- a/tests/test_mxfp4_metal.c +++ b/tests/test_mxfp4_metal.c @@ -16,12 +16,22 @@ #define N_EXPERT 6u #define DIM 256u #define BATCH_TOKENS 48u +#define ATTN_GROUPS 8u +#define ATTN_GROUP_DIM 4096u +#define ATTN_RANK 1024u typedef struct { uint8_t e; uint8_t qs[QK_MXFP4 / 2u]; } block_mxfp4; +typedef struct { + _Float16 d; + int8_t qs[32]; +} block_q8_0; +_Static_assert(sizeof(block_q8_0) == 34u, + "Q8_0 test fixture must match the Metal block stride"); + static const float mxfp4_values[16] = { 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, -0.0f,-0.5f,-1.0f,-1.5f,-2.0f,-3.0f,-4.0f,-6.0f, @@ -36,6 +46,33 @@ static uint64_t align_up(uint64_t value, uint64_t alignment) { return (value + alignment - 1u) / alignment * alignment; } +typedef struct { + const char *name; + char *value; + bool had_value; +} saved_env; + +static int save_env(saved_env *saved, const char *name) { + const char *value = getenv(name); + saved->name = name; + saved->had_value = value != NULL; + saved->value = value ? strdup(value) : NULL; + return !value || saved->value != NULL; +} + +static int restore_env(saved_env *saved) { + int rc = 0; + if (saved->had_value) { + if (!saved->value) return 0; + rc = setenv(saved->name, saved->value, 1); + } else { + rc = unsetenv(saved->name); + } + free(saved->value); + saved->value = NULL; + return rc == 0; +} + static float e8m0_to_f32(uint8_t e) { uint32_t bits = e == 0 ? 0x00400000u : (uint32_t)e << 23u; float value; @@ -116,7 +153,14 @@ int main(void) { const uint64_t gate_offset = 0; const uint64_t up_offset = align_up(tensor_bytes, page); const uint64_t down_offset = align_up(up_offset + tensor_bytes, page); - const uint64_t model_size = align_up(down_offset + tensor_bytes, page); + const uint64_t attn_row_bytes = + (ATTN_GROUP_DIM / 32u) * sizeof(block_q8_0); + const uint64_t attn_weight_bytes = + (uint64_t)ATTN_GROUPS * ATTN_RANK * attn_row_bytes; + const uint64_t attn_offset = + align_up(down_offset + tensor_bytes, page); + const uint64_t model_size = + align_up(attn_offset + attn_weight_bytes, page); void *model = NULL; if (posix_memalign(&model, (size_t)page, (size_t)model_size) != 0) { fprintf(stderr, "MXFP4 Metal test model allocation failed\n"); @@ -126,6 +170,17 @@ int main(void) { fill_matrix((block_mxfp4 *)((uint8_t *)model + gate_offset), 1u); fill_matrix((block_mxfp4 *)((uint8_t *)model + up_offset), 5u); fill_matrix((block_mxfp4 *)((uint8_t *)model + down_offset), 9u); + block_q8_0 *attn_matrix = + (block_q8_0 *)((uint8_t *)model + attn_offset); + const uint64_t attn_blocks = + (uint64_t)ATTN_GROUPS * ATTN_RANK * (ATTN_GROUP_DIM / 32u); + for (uint64_t block = 0; block < attn_blocks; block++) { + attn_matrix[block].d = (_Float16)(1.0f / 128.0f); + for (uint32_t i = 0; i < 32u; i++) { + attn_matrix[block].qs[i] = + (int8_t)((int32_t)((block * 11u + i * 7u) % 31u) - 15); + } + } float x[DIM]; int32_t selected[N_EXPERT] = { 0, 2, 3, 5, 6, 7 }; @@ -263,7 +318,7 @@ int main(void) { MXFP4_TYPE, MXFP4_TYPE, expert_bytes, row_bytes, expert_bytes, row_bytes, DIM, DIM, DIM, selected_tensor, weights_tensor, N_TOTAL_EXPERT, N_EXPERT, - 7.0f, x_tensor, NULL, 0u, true); + 7.0f, x_tensor, NULL, 0u, true, NULL, NULL, NULL, NULL, NULL); ok = ok && ds4_gpu_tensor_read( gate_tensor, 0, gate_gpu, pair_count * sizeof(float)); ok = ok && ds4_gpu_tensor_read( @@ -299,7 +354,7 @@ int main(void) { MXFP4_TYPE, MXFP4_TYPE, expert_bytes, row_bytes, expert_bytes, row_bytes, DIM, DIM, DIM, selected_tensor, weights_tensor, N_TOTAL_EXPERT, N_EXPERT, - 7.0f, x_tensor, NULL, 0u, true) && + 7.0f, x_tensor, NULL, 0u, true, NULL, NULL, NULL, NULL, NULL) && ds4_gpu_tensor_read( gate_tensor, 0, gate_fast, pair_count * sizeof(float)) && ds4_gpu_tensor_read( @@ -347,6 +402,95 @@ int main(void) { compare_values("out", out_gpu, out_ref, DIM, 2.0e-4f); } + /* Compare the exact fixed-shape attention-output LOW kernel directly + * against its generic rollback. Force the static PSO through the test + * flag so this remains coverage on newer Apple GPUs, and poison the full + * destination before each run to catch partial writes. */ + const uint64_t attn_heads_count = + (uint64_t)ATTN_GROUPS * ATTN_GROUP_DIM; + const uint64_t attn_low_count = + (uint64_t)ATTN_GROUPS * ATTN_RANK; + const size_t attn_heads_bytes = + (size_t)attn_heads_count * sizeof(float); + const size_t attn_low_bytes = + (size_t)attn_low_count * sizeof(float); + float *attn_heads = malloc(attn_heads_bytes); + uint8_t *attn_poison = malloc(attn_low_bytes); + uint8_t *attn_generic = malloc(attn_low_bytes); + uint8_t *attn_static = malloc(attn_low_bytes); + ds4_gpu_tensor *attn_heads_tensor = + ds4_gpu_tensor_alloc(attn_heads_bytes); + ds4_gpu_tensor *attn_low_tensor = + ds4_gpu_tensor_alloc(attn_low_bytes); + saved_env decode_ports_env = { 0 }; + saved_env attn_static_env = { 0 }; + bool attn_env_ok = save_env( + &decode_ports_env, "DS4_METAL_DISABLE_PRE_M5_DECODE_PORTS") != 0; + attn_env_ok = save_env( + &attn_static_env, + "DS4_METAL_DISABLE_PRE_M5_ATTN_OUT_LOW_Q8_STATIC") != 0 && + attn_env_ok; + bool attn_exact = ok && attn_env_ok && attn_heads && attn_poison && + attn_generic && attn_static && attn_heads_tensor && attn_low_tensor; + if (attn_heads) { + for (uint64_t i = 0; i < attn_heads_count; i++) { + attn_heads[i] = + (float)((int32_t)((i * 17u) % 257u) - 128) / 256.0f; + } + } + if (attn_poison) memset(attn_poison, 0xa5, attn_low_bytes); + attn_exact = attn_exact && ds4_gpu_tensor_write( + attn_heads_tensor, 0, attn_heads, attn_heads_bytes); + + if (attn_env_ok && + setenv("DS4_METAL_DISABLE_PRE_M5_ATTN_OUT_LOW_Q8_STATIC", "1", 1) != 0) { + attn_exact = false; + } + ds4_gpu_test_set_flags(0); + attn_exact = attn_exact && ds4_gpu_tensor_write( + attn_low_tensor, 0, attn_poison, attn_low_bytes); + attn_exact = attn_exact && ds4_gpu_attention_output_low_q8_tensor( + attn_low_tensor, model, model_size, attn_offset, + ATTN_GROUP_DIM, ATTN_RANK, ATTN_GROUPS, + attn_heads_tensor); + attn_exact = attn_exact && ds4_gpu_tensor_read( + attn_low_tensor, 0, attn_generic, attn_low_bytes); + + if (attn_env_ok && + (unsetenv("DS4_METAL_DISABLE_PRE_M5_DECODE_PORTS") != 0 || + unsetenv("DS4_METAL_DISABLE_PRE_M5_ATTN_OUT_LOW_Q8_STATIC") != 0)) { + attn_exact = false; + } + ds4_gpu_test_set_flags(DS4_GPU_TEST_ATTN_OUT_LOW_Q8_STATIC); + attn_exact = attn_exact && ds4_gpu_tensor_write( + attn_low_tensor, 0, attn_poison, attn_low_bytes); + attn_exact = attn_exact && ds4_gpu_attention_output_low_q8_tensor( + attn_low_tensor, model, model_size, attn_offset, + ATTN_GROUP_DIM, ATTN_RANK, ATTN_GROUPS, + attn_heads_tensor); + attn_exact = attn_exact && ds4_gpu_tensor_read( + attn_low_tensor, 0, attn_static, attn_low_bytes); + ds4_gpu_test_set_flags(0); + + if (attn_exact && memcmp(attn_generic, attn_static, attn_low_bytes) != 0) { + fprintf(stderr, + "MXFP4 Metal attention-output LOW static/generic A/B mismatch\n"); + attn_exact = false; + } else if (attn_exact) { + fprintf(stderr, + "MXFP4 Metal attention-output LOW static/generic A/B exact\n"); + } + const bool attn_static_env_restored = restore_env(&attn_static_env) != 0; + const bool decode_ports_env_restored = restore_env(&decode_ports_env) != 0; + ok = ok && attn_exact && + attn_static_env_restored && decode_ports_env_restored; + ds4_gpu_tensor_free(attn_low_tensor); + ds4_gpu_tensor_free(attn_heads_tensor); + free(attn_static); + free(attn_generic); + free(attn_poison); + free(attn_heads); + /* The production large-prefill path stores its fused SwiGLU result as * FP16 before the down projection. Compare that independent grouped-MMA * path against the same scalar reference with the documented rounding. */ @@ -363,11 +507,15 @@ int main(void) { (size_t)batch_pairs, sizeof(_Float16)); _Float16 *mid_batch_storage = calloc( (size_t)batch_pairs, sizeof(_Float16)); + _Float16 *mid_batch_half_lut_baseline = calloc( + (size_t)batch_pairs, sizeof(_Float16)); float *out_batch_expected = calloc((size_t)batch_out_count, sizeof(float)); float *out_batch_baseline = calloc( (size_t)batch_out_count, sizeof(float)); + float *out_batch_half_lut_baseline = calloc( + (size_t)batch_out_count, sizeof(float)); float *out_batch_actual = calloc((size_t)batch_out_count, sizeof(float)); - float *experts_batch_baseline = calloc( + float *experts_batch_half_lut_baseline = calloc( (size_t)batch_out_count * N_EXPERT, sizeof(float)); float *experts_batch_actual = calloc( (size_t)batch_out_count * N_EXPERT, sizeof(float)); @@ -419,9 +567,11 @@ int main(void) { bool mid_is_f16 = false; ok = ok && x_batch && selected_batch && weights_batch && mid_batch_expected && mid_batch_actual && mid_batch_baseline && - mid_batch_storage && out_batch_expected && out_batch_baseline && - out_batch_actual && experts_batch_baseline && - experts_batch_actual && batch_poison && + mid_batch_storage && mid_batch_half_lut_baseline && + out_batch_expected && out_batch_baseline && + out_batch_half_lut_baseline && out_batch_actual && + experts_batch_half_lut_baseline && experts_batch_actual && + batch_poison && x_batch_tensor && selected_batch_tensor && weights_batch_tensor && gate_batch_tensor && up_batch_tensor && mid_batch_tensor && experts_batch_tensor && out_batch_tensor; @@ -437,11 +587,25 @@ int main(void) { memset(batch_poison, 0xa5, (size_t)batch_out_count * N_EXPERT * sizeof(float)); } + /* Keep the established per-feature A/B checks isolated now that the + * production defaults also cover this 48-token shape. The widened + * dispatcher gets its own aggregate rollback comparison below. */ + const char *small_prefill_rollback_name = + "DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_SMALL_PREFILL"; + const char *small_prefill_rollback_previous = + getenv(small_prefill_rollback_name); + char *small_prefill_rollback_saved = small_prefill_rollback_previous ? + strdup(small_prefill_rollback_previous) : NULL; + if ((small_prefill_rollback_previous && !small_prefill_rollback_saved) || + setenv(small_prefill_rollback_name, "1", 1) != 0) { + fprintf(stderr, + "MXFP4 Metal could not isolate short-prefill feature tests\n"); + ok = 0; + } /* With 48 identical routes, every occupied expert has a 32-row tile plus - * a 16-row tail tile. Run the established, pair-culling, and down-culling - * pipelines in the same process and require exact FP16 intermediates and - * F32 outputs. */ - ds4_gpu_test_set_flags(DS4_GPU_TEST_MXFP4_DOWN_TAIL_CULL); + * a 16-row tail tile. First capture the generic rollback as the baseline + * for the independent pair, compact-pair, and down-tail checks. */ + ds4_gpu_test_set_flags(0); ok = ok && ds4_gpu_tensor_write( mid_batch_tensor, 0, batch_poison, batch_pairs * sizeof(_Float16)); @@ -459,7 +623,7 @@ int main(void) { expert_bytes, row_bytes, DIM, DIM, DIM, selected_batch_tensor, weights_batch_tensor, N_TOTAL_EXPERT, N_EXPERT, 7.0f, x_batch_tensor, - 0u, BATCH_TOKENS, &mid_is_f16, true); + 0u, BATCH_TOKENS, &mid_is_f16, false, true); ok = ok && mid_is_f16; ok = ok && ds4_gpu_tensor_read( mid_batch_tensor, 0, mid_batch_baseline, @@ -467,8 +631,38 @@ int main(void) { ok = ok && ds4_gpu_tensor_read( out_batch_tensor, 0, out_batch_baseline, batch_out_count * sizeof(float)); + + /* The half-LUT candidate retains the down-tail kernel, so give it a + * separate down-tail baseline instead of conflating two features. */ + mid_is_f16 = false; + ds4_gpu_test_set_flags(DS4_GPU_TEST_MXFP4_DOWN_TAIL_CULL); + ok = ok && ds4_gpu_tensor_write( + mid_batch_tensor, 0, batch_poison, + batch_pairs * sizeof(_Float16)); + ok = ok && ds4_gpu_tensor_write( + experts_batch_tensor, 0, batch_poison, + batch_out_count * N_EXPERT * sizeof(float)); + ok = ok && ds4_gpu_tensor_write( + out_batch_tensor, 0, batch_poison, + batch_out_count * sizeof(float)); + ok = ok && ds4_gpu_routed_moe_batch_tensor( + out_batch_tensor, gate_batch_tensor, up_batch_tensor, + mid_batch_tensor, experts_batch_tensor, + model, model_size, gate_offset, up_offset, down_offset, + MXFP4_TYPE, MXFP4_TYPE, expert_bytes, row_bytes, + expert_bytes, row_bytes, DIM, DIM, DIM, + selected_batch_tensor, weights_batch_tensor, + N_TOTAL_EXPERT, N_EXPERT, 7.0f, x_batch_tensor, + 0u, BATCH_TOKENS, &mid_is_f16, false, true); + ok = ok && mid_is_f16; + ok = ok && ds4_gpu_tensor_read( + mid_batch_tensor, 0, mid_batch_half_lut_baseline, + batch_pairs * sizeof(_Float16)); + ok = ok && ds4_gpu_tensor_read( + out_batch_tensor, 0, out_batch_half_lut_baseline, + batch_out_count * sizeof(float)); ok = ok && ds4_gpu_tensor_read( - experts_batch_tensor, 0, experts_batch_baseline, + experts_batch_tensor, 0, experts_batch_half_lut_baseline, batch_out_count * N_EXPERT * sizeof(float)); /* Force the exact half-result dequantization table only for the resident @@ -498,7 +692,7 @@ int main(void) { expert_bytes, row_bytes, DIM, DIM, DIM, selected_batch_tensor, weights_batch_tensor, N_TOTAL_EXPERT, N_EXPERT, 7.0f, x_batch_tensor, - 0u, BATCH_TOKENS, &mid_is_f16, true); + 0u, BATCH_TOKENS, &mid_is_f16, false, true); ok = ok && mid_is_f16; ok = ok && ds4_gpu_tensor_read( mid_batch_tensor, 0, mid_batch_storage, @@ -510,11 +704,11 @@ int main(void) { out_batch_tensor, 0, out_batch_actual, batch_out_count * sizeof(float)); if (ok && - (memcmp(mid_batch_storage, mid_batch_baseline, + (memcmp(mid_batch_storage, mid_batch_half_lut_baseline, batch_pairs * sizeof(_Float16)) != 0 || - memcmp(experts_batch_actual, experts_batch_baseline, + memcmp(experts_batch_actual, experts_batch_half_lut_baseline, batch_out_count * N_EXPERT * sizeof(float)) != 0 || - memcmp(out_batch_actual, out_batch_baseline, + memcmp(out_batch_actual, out_batch_half_lut_baseline, batch_out_count * sizeof(float)) != 0)) { fprintf(stderr, "MXFP4 Metal down half-LUT poisoned A/B mismatch on repetition %u\n", @@ -547,7 +741,7 @@ int main(void) { expert_bytes, row_bytes, DIM, DIM, DIM, selected_batch_tensor, weights_batch_tensor, N_TOTAL_EXPERT, N_EXPERT, 7.0f, x_batch_tensor, - 0u, BATCH_TOKENS, &mid_is_f16, true); + 0u, BATCH_TOKENS, &mid_is_f16, false, true); ok = ok && mid_is_f16; ok = ok && ds4_gpu_tensor_read( mid_batch_tensor, 0, mid_batch_storage, @@ -592,7 +786,7 @@ int main(void) { expert_bytes, row_bytes, DIM, DIM, DIM, selected_batch_tensor, weights_batch_tensor, N_TOTAL_EXPERT, N_EXPERT, 7.0f, x_batch_tensor, - 0u, BATCH_TOKENS, &mid_is_f16, true); + 0u, BATCH_TOKENS, &mid_is_f16, false, true); ok = ok && mid_is_f16; ok = ok && ds4_gpu_tensor_read( mid_batch_tensor, 0, mid_batch_storage, @@ -632,7 +826,7 @@ int main(void) { expert_bytes, row_bytes, DIM, DIM, DIM, selected_batch_tensor, weights_batch_tensor, N_TOTAL_EXPERT, N_EXPERT, 7.0f, x_batch_tensor, - 0u, BATCH_TOKENS, &mid_is_f16, true); + 0u, BATCH_TOKENS, &mid_is_f16, false, true); ok = ok && mid_is_f16; ok = ok && ds4_gpu_tensor_read( mid_batch_tensor, 0, mid_batch_storage, @@ -689,7 +883,7 @@ int main(void) { expert_bytes, row_bytes, DIM, DIM, DIM, selected_batch_tensor, weights_batch_tensor, N_TOTAL_EXPERT, N_EXPERT, 7.0f, x_batch_tensor, - 0u, BATCH_TOKENS, &mid_is_f16, true); + 0u, BATCH_TOKENS, &mid_is_f16, false, true); ok = ok && mid_is_f16; ok = ok && ds4_gpu_tensor_read( mid_batch_tensor, 0, mid_batch_baseline, @@ -720,7 +914,7 @@ int main(void) { expert_bytes, row_bytes, DIM, DIM, DIM, selected_batch_tensor, weights_batch_tensor, N_TOTAL_EXPERT, N_EXPERT, 7.0f, x_batch_tensor, - 0u, BATCH_TOKENS, &mid_is_f16, true); + 0u, BATCH_TOKENS, &mid_is_f16, false, true); ok = ok && mid_is_f16; ok = ok && ds4_gpu_tensor_read( mid_batch_tensor, 0, mid_batch_storage, @@ -744,6 +938,101 @@ int main(void) { } ds4_gpu_test_set_flags(0); + /* Finally compare the complete automatic 32..2047-token default bundle + * with its aggregate rollback. Reuse the uneven routes above so the map + * contains full tiles plus 16-, 15-, and 1-row tails. */ + mid_is_f16 = false; + ok = ok && ds4_gpu_tensor_write( + mid_batch_tensor, 0, batch_poison, + batch_pairs * sizeof(_Float16)); + ok = ok && ds4_gpu_tensor_write( + experts_batch_tensor, 0, batch_poison, + batch_out_count * N_EXPERT * sizeof(float)); + ok = ok && ds4_gpu_tensor_write( + out_batch_tensor, 0, batch_poison, + batch_out_count * sizeof(float)); + ok = ok && ds4_gpu_routed_moe_batch_tensor( + out_batch_tensor, gate_batch_tensor, up_batch_tensor, + mid_batch_tensor, experts_batch_tensor, + model, model_size, gate_offset, up_offset, down_offset, + MXFP4_TYPE, MXFP4_TYPE, expert_bytes, row_bytes, + expert_bytes, row_bytes, DIM, DIM, DIM, + selected_batch_tensor, weights_batch_tensor, + N_TOTAL_EXPERT, N_EXPERT, 7.0f, x_batch_tensor, + 0u, BATCH_TOKENS, &mid_is_f16, false, true); + ok = ok && mid_is_f16; + ok = ok && ds4_gpu_tensor_read( + mid_batch_tensor, 0, mid_batch_baseline, + batch_pairs * sizeof(_Float16)); + ok = ok && ds4_gpu_tensor_read( + experts_batch_tensor, 0, experts_batch_half_lut_baseline, + batch_out_count * N_EXPERT * sizeof(float)); + ok = ok && ds4_gpu_tensor_read( + out_batch_tensor, 0, out_batch_baseline, + batch_out_count * sizeof(float)); + + if (unsetenv(small_prefill_rollback_name) != 0) { + fprintf(stderr, + "MXFP4 Metal could not enable short-prefill defaults\n"); + ok = 0; + } + mid_is_f16 = false; + ok = ok && ds4_gpu_tensor_write( + mid_batch_tensor, 0, batch_poison, + batch_pairs * sizeof(_Float16)); + ok = ok && ds4_gpu_tensor_write( + experts_batch_tensor, 0, batch_poison, + batch_out_count * N_EXPERT * sizeof(float)); + ok = ok && ds4_gpu_tensor_write( + out_batch_tensor, 0, batch_poison, + batch_out_count * sizeof(float)); + ok = ok && ds4_gpu_routed_moe_batch_tensor( + out_batch_tensor, gate_batch_tensor, up_batch_tensor, + mid_batch_tensor, experts_batch_tensor, + model, model_size, gate_offset, up_offset, down_offset, + MXFP4_TYPE, MXFP4_TYPE, expert_bytes, row_bytes, + expert_bytes, row_bytes, DIM, DIM, DIM, + selected_batch_tensor, weights_batch_tensor, + N_TOTAL_EXPERT, N_EXPERT, 7.0f, x_batch_tensor, + 0u, BATCH_TOKENS, &mid_is_f16, false, true); + ok = ok && mid_is_f16; + ok = ok && ds4_gpu_tensor_read( + mid_batch_tensor, 0, mid_batch_storage, + batch_pairs * sizeof(_Float16)); + ok = ok && ds4_gpu_tensor_read( + experts_batch_tensor, 0, experts_batch_actual, + batch_out_count * N_EXPERT * sizeof(float)); + ok = ok && ds4_gpu_tensor_read( + out_batch_tensor, 0, out_batch_actual, + batch_out_count * sizeof(float)); + if (ok && (memcmp(mid_batch_storage, mid_batch_baseline, + batch_pairs * sizeof(_Float16)) != 0 || + memcmp(experts_batch_actual, + experts_batch_half_lut_baseline, + batch_out_count * N_EXPERT * sizeof(float)) != 0 || + memcmp(out_batch_actual, out_batch_baseline, + batch_out_count * sizeof(float)) != 0)) { + fprintf(stderr, + "MXFP4 Metal short-prefill default bundle A/B mismatch\n"); + ok = 0; + } else if (ok) { + fprintf(stderr, + "MXFP4 Metal short-prefill default bundle A/B exact\n"); + } + if (small_prefill_rollback_saved) { + if (setenv(small_prefill_rollback_name, + small_prefill_rollback_saved, 1) != 0) { + fprintf(stderr, + "MXFP4 Metal could not restore short-prefill rollback\n"); + ok = 0; + } + } else if (unsetenv(small_prefill_rollback_name) != 0) { + fprintf(stderr, + "MXFP4 Metal could not restore short-prefill environment\n"); + ok = 0; + } + free(small_prefill_rollback_saved); + ds4_gpu_tensor_free(out_batch_tensor); ds4_gpu_tensor_free(experts_batch_tensor); ds4_gpu_tensor_free(mid_batch_tensor); @@ -753,12 +1042,14 @@ int main(void) { ds4_gpu_tensor_free(selected_batch_tensor); ds4_gpu_tensor_free(x_batch_tensor); free(out_batch_actual); + free(out_batch_half_lut_baseline); free(out_batch_baseline); free(out_batch_expected); free(experts_batch_actual); - free(experts_batch_baseline); + free(experts_batch_half_lut_baseline); free(batch_poison); free(mid_batch_storage); + free(mid_batch_half_lut_baseline); free(mid_batch_baseline); free(mid_batch_actual); free(mid_batch_expected);