From 15941eefa8606be18d89faf21a53321d3e95eb56 Mon Sep 17 00:00:00 2001 From: ivanfioravanti Date: Thu, 20 Aug 2026 11:57:02 +0200 Subject: [PATCH 01/16] metal: speed up pre-M5 decode and short prefill --- ds4_gpu.h | 1 + ds4_metal.m | 61 ++++++-- metal/moe.metal | 92 +++++++++++ speed-bench/README.md | 39 +++++ tests/test_mxfp4_metal.c | 319 +++++++++++++++++++++++++++++++++++++-- 5 files changed, 486 insertions(+), 26 deletions(-) diff --git a/ds4_gpu.h b/ds4_gpu.h index 21d0160191..00eed4e606 100644 --- a/ds4_gpu.h +++ b/ds4_gpu.h @@ -186,6 +186,7 @@ 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, }; void ds4_gpu_test_set_flags(uint32_t flags); void ds4_gpu_release_zero_prefix_prefill_mask_cache(void); diff --git a/ds4_metal.m b/ds4_metal.m index 3363d7df5e..552456fbe9 100644 --- a/ds4_metal.m +++ b/ds4_metal.m @@ -191,6 +191,7 @@ static id g_dsv4_indexed_attention_heads16_dual_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; @@ -10354,6 +10355,7 @@ void ds4_gpu_cleanup(void) { g_dsv4_indexed_attention_heads16_dual_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; @@ -25268,8 +25270,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, @@ -41461,6 +41492,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 +41510,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 +41526,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 +41545,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 +41556,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 +41573,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 && diff --git a/metal/moe.metal b/metal/moe.metal index 7aeb9d9222..27f11ad6fc 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, diff --git a/speed-bench/README.md b/speed-bench/README.md index 6d4188e72c..28a0638f05 100644 --- a/speed-bench/README.md +++ b/speed-bench/README.md @@ -54,6 +54,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 +93,26 @@ 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. + 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/tests/test_mxfp4_metal.c b/tests/test_mxfp4_metal.c index bb1f37fed2..8d036c1197 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 }; @@ -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)); @@ -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, 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( - experts_batch_tensor, 0, experts_batch_baseline, + 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_half_lut_baseline, batch_out_count * N_EXPERT * sizeof(float)); /* Force the exact half-result dequantization table only for the resident @@ -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", @@ -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, 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, 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); From 8d0c86b0ba99539acce6ca9512861be94258f65a Mon Sep 17 00:00:00 2001 From: ivanfioravanti Date: Thu, 20 Aug 2026 12:49:19 +0200 Subject: [PATCH 02/16] metal: overlap batch Q and KV finalizers on M3 --- ds4.c | 90 ++++++++++++--- ds4_gpu.h | 22 ++++ ds4_metal.m | 256 +++++++++++++++++++++++++++++++++++++++--- metal/dsv4_rope.metal | 100 +++++++++++++++++ speed-bench/README.md | 31 +++++ tests/ds4_test.c | 209 ++++++++++++++++++++++++++++++++++ 6 files changed, 674 insertions(+), 34 deletions(-) diff --git a/ds4.c b/ds4.c index b54075539e..38a1cb2b5a 100644 --- a/ds4.c +++ b/ds4.c @@ -29331,6 +29331,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 +29382,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 +29517,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 +29535,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 +29552,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 +29615,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], @@ -30297,12 +30351,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) { diff --git a/ds4_gpu.h b/ds4_gpu.h index 00eed4e606..72e37b605f 100644 --- a/ds4_gpu.h +++ b/ds4_gpu.h @@ -1120,6 +1120,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, diff --git a/ds4_metal.m b/ds4_metal.m index 552456fbe9..35e8360d36 100644 --- a/ds4_metal.m +++ b/ds4_metal.m @@ -168,6 +168,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; @@ -5142,6 +5144,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) { @@ -5752,7 +5782,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 @@ -10332,6 +10364,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; @@ -21307,22 +21341,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), @@ -21364,6 +21383,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, diff --git a/metal/dsv4_rope.metal b/metal/dsv4_rope.metal index 279365c8f1..b772435843 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 diff --git a/speed-bench/README.md b/speed-bench/README.md index 28a0638f05..8ba64af46b 100644 --- a/speed-bench/README.md +++ b/speed-bench/README.md @@ -113,6 +113,37 @@ 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 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. + 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/tests/ds4_test.c b/tests/ds4_test.c index cf8bca2c52..3545aa7693 100644 --- a/tests/ds4_test.c +++ b/tests/ds4_test.c @@ -863,6 +863,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()); } @@ -4730,6 +4938,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(); From 9408fd16423010e43ae6a10568a50eac9c2ba012 Mon Sep 17 00:00:00 2001 From: ivanfioravanti Date: Thu, 20 Aug 2026 13:49:24 +0200 Subject: [PATCH 03/16] metal: prune unused batch indexer queries on pre-M5 --- ds4.c | 28 +++++++++++++++++++---- speed-bench/README.md | 28 +++++++++++++++++++++++ speed-bench/metal_decode_schedule_bench.c | 6 ++++- 3 files changed, 56 insertions(+), 6 deletions(-) diff --git a/ds4.c b/ds4.c index 38a1cb2b5a..75c18141cb 100644 --- a/ds4.c +++ b/ds4.c @@ -29976,6 +29976,25 @@ 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 && + ds4_gpu_device_is_pre_m5_apple_silicon() && + getenv("DS4_METAL_DISABLE_PRE_M5_BATCH_INDEXER_QUERY_PRUNE") == NULL; +#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 || @@ -30012,14 +30031,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, @@ -30033,10 +30052,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, @@ -30478,7 +30497,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)); diff --git a/speed-bench/README.md b/speed-bench/README.md index 8ba64af46b..609c6d8119 100644 --- a/speed-bench/README.md +++ b/speed-bench/README.md @@ -113,6 +113,34 @@ 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. + ### Metal batch Q/KV finalizer A/B The M3 resident Flash prefill path now follows vLLM's horizontal Q/KV diff --git a/speed-bench/metal_decode_schedule_bench.c b/speed-bench/metal_decode_schedule_bench.c index 5866b49db4..575fc8944c 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", From 364109be8692125c7cca27f8cfec7000d2f95678 Mon Sep 17 00:00:00 2001 From: ivanfioravanti Date: Thu, 20 Aug 2026 16:06:00 +0200 Subject: [PATCH 04/16] metal: stage indexed prefill attention rows on pre-M5 --- ds4_gpu.h | 2 + ds4_metal.m | 43 ++++++++++- metal/dsv4_misc.metal | 118 +++++++++++++++++++++++++++++ speed-bench/README.md | 33 +++++++++ tests/ds4_test.c | 168 ++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 360 insertions(+), 4 deletions(-) diff --git a/ds4_gpu.h b/ds4_gpu.h index 72e37b605f..a07610c24d 100644 --- a/ds4_gpu.h +++ b/ds4_gpu.h @@ -187,7 +187,9 @@ enum { 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, }; +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 diff --git a/ds4_metal.m b/ds4_metal.m index 35e8360d36..ce5a24f3fd 100644 --- a/ds4_metal.m +++ b/ds4_metal.m @@ -191,6 +191,7 @@ 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_heads8_split_pipeline; static id g_dsv4_indexed_attention_heads8_split_reduce_pipeline; static bool g_attn_out_low_q8_static_unavailable; @@ -8496,6 +8497,8 @@ 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_heads8_split_pipeline = ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads8_split"); g_dsv4_indexed_attention_heads8_split_reduce_pipeline = @@ -8768,6 +8771,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; @@ -10387,6 +10394,7 @@ 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_heads8_split_pipeline = nil; g_dsv4_indexed_attention_heads8_split_reduce_pipeline = nil; g_attn_out_low_q8_static_unavailable = false; @@ -29208,6 +29216,29 @@ 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; + 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; @@ -29219,7 +29250,10 @@ 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_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, @@ -29356,13 +29390,14 @@ 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 + (use_prefill_dual_heads ? 15u : 7u)) / + (use_prefill_dual_heads ? 16u : 8u), 1) threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; ds4_gpu_end_compute_encoder(cb, enc); diff --git a/metal/dsv4_misc.metal b/metal/dsv4_misc.metal index 5620b24432..eb31c8d110 100644 --- a/metal/dsv4_misc.metal +++ b/metal/dsv4_misc.metal @@ -5936,6 +5936,124 @@ 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; +} + // 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/speed-bench/README.md b/speed-bench/README.md index 609c6d8119..7ee141c181 100644 --- a/speed-bench/README.md +++ b/speed-bench/README.md @@ -172,6 +172,39 @@ schedule is based on vLLM's 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 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/tests/ds4_test.c b/tests/ds4_test.c index 3545aa7693..98a7fe5b90 100644 --- a/tests/ds4_test.c +++ b/tests/ds4_test.c @@ -3329,6 +3329,173 @@ static void test_metal_contiguous_compressed_f16_attention_exact(void) { ds4_gpu_tensor_free(raw); } +static void test_metal_indexed_attention_prefill_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"; + 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(); + + 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); + + /* Quality mode selects the original heads8 kernel even on M5, where + * the normal fast fallback is heads16_dual. */ + ds4_gpu_test_set_flags(0); + ds4_gpu_set_quality(true); + TEST_ASSERT(setenv(disable_env, "1", 1) == 0); + TEST_ASSERT(unsetenv(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); + + /* The force bit changes only the Apple-generation gate. REQUIRE makes + * this call fail instead of silently comparing heads8 with itself. */ + 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(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(require_env, saved_require); + test_restore_env(disable_env, saved_disable); + fprintf(stderr, + "ds4-test: indexed-attention prefill 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, @@ -4948,6 +5115,7 @@ 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_rb4_exact(); test_metal_persistent_zero_attention_mask_exact(); test_metal_zero_prefix_prefill_mask_cache_exact(); test_metal_hc_split_weighted_sum_norm_batch_exact(); From 8b7c01241335d44e1ebcc0199fd92af62a5f1f2e Mon Sep 17 00:00:00 2001 From: ivanfioravanti Date: Thu, 20 Aug 2026 23:23:22 +0200 Subject: [PATCH 05/16] metal: accelerate pre-M5 long prefill --- ds4.c | 86 +++++++- ds4_cuda.cu | 3 +- ds4_gpu.h | 44 ++++ ds4_metal.m | 438 ++++++++++++++++++++++++++++++++++++- metal/dsv4_hc.metal | 227 +++++++++++++++++++ metal/dsv4_misc.metal | 146 +++++++++++++ speed-bench/README.md | 92 ++++++++ tests/ds4_test.c | 458 +++++++++++++++++++++++++++++++++++++-- tests/test_mxfp4_metal.c | 20 +- 9 files changed, 1478 insertions(+), 36 deletions(-) diff --git a/ds4.c b/ds4.c index 75c18141cb..6934aee9f2 100644 --- a/ds4.c +++ b/ds4.c @@ -30851,6 +30851,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). @@ -30865,7 +30901,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 @@ -30973,7 +31009,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), @@ -31283,10 +31321,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, \ @@ -31537,6 +31605,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) { @@ -31629,6 +31698,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), @@ -44437,6 +44516,7 @@ static int glm_graph_routed_moe_batch_dispatch( il, n_tokens, &g->batch_routed_mid_is_f16, + false, force_resident); } diff --git a/ds4_cuda.cu b/ds4_cuda.cu index fd36d7014e..37ea93e1d1 100644 --- a/ds4_cuda.cu +++ b/ds4_cuda.cu @@ -26014,8 +26014,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_gpu.h b/ds4_gpu.h index a07610c24d..8c33e7b658 100644 --- a/ds4_gpu.h +++ b/ds4_gpu.h @@ -188,6 +188,7 @@ enum { 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); @@ -2275,6 +2276,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, @@ -2743,6 +2769,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); /* ========================================================================= @@ -2924,6 +2951,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 ce5a24f3fd..0c2752bcaf 100644 --- a/ds4_metal.m +++ b/ds4_metal.m @@ -192,6 +192,7 @@ 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; @@ -269,6 +270,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 { @@ -8499,6 +8502,8 @@ int ds4_gpu_init(void) { 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 = @@ -8644,6 +8649,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 || @@ -10395,6 +10404,7 @@ void ds4_gpu_cleanup(void) { 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; @@ -10473,6 +10483,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; @@ -24458,7 +24470,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, @@ -24472,7 +24562,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 || @@ -24786,9 +24877,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"); @@ -24800,6 +24896,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, @@ -29230,6 +29455,26 @@ int ds4_gpu_attention_indexed_mixed_batch_heads_tensor( 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) { @@ -29250,6 +29495,9 @@ 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_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") : @@ -29396,8 +29644,9 @@ int ds4_gpu_attention_indexed_mixed_batch_heads_tensor( atIndex:0]; [enc dispatchThreadgroups: MTLSizeMake((NSUInteger)n_tokens, - ((NSUInteger)n_head + (use_prefill_dual_heads ? 15u : 7u)) / - (use_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); @@ -32135,6 +32384,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, @@ -41386,6 +41675,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; @@ -41415,6 +41705,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 @@ -41935,6 +42226,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) { @@ -42614,6 +42915,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) { @@ -44041,6 +44343,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 c467946754..fab24b396e 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 eb31c8d110..c949838dfa 100644 --- a/metal/dsv4_misc.metal +++ b/metal/dsv4_misc.metal @@ -6054,6 +6054,152 @@ kernel void kernel_dsv4_indexed_mixed_attention_heads8_rb4( 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/speed-bench/README.md b/speed-bench/README.md index 7ee141c181..446b242130 100644 --- a/speed-bench/README.md +++ b/speed-bench/README.md @@ -205,6 +205,98 @@ 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/tests/ds4_test.c b/tests/ds4_test.c index 98a7fe5b90..fd16f67355 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) { @@ -3329,7 +3330,7 @@ static void test_metal_contiguous_compressed_f16_attention_exact(void) { ds4_gpu_tensor_free(raw); } -static void test_metal_indexed_attention_prefill_rb4_exact(void) { +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; @@ -3355,8 +3356,14 @@ static void test_metal_indexed_attention_prefill_rb4_exact(void) { "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; @@ -3437,23 +3444,27 @@ static void test_metal_indexed_attention_prefill_rb4_exact(void) { candidate, 0, candidate_host, q_bytes) != 0); TEST_ASSERT(ds4_gpu_set_model_map(model_raw, page) != 0); - /* Quality mode selects the original heads8 kernel even on M5, where - * the normal fast fallback is heads16_dual. */ - ds4_gpu_test_set_flags(0); - ds4_gpu_set_quality(true); - TEST_ASSERT(setenv(disable_env, "1", 1) == 0); - TEST_ASSERT(unsetenv(require_env) == 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); - /* The force bit changes only the Apple-generation gate. REQUIRE makes - * this call fail instead of silently comparing heads8 with itself. */ - 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); + /* 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, @@ -3469,10 +3480,12 @@ static void test_metal_indexed_attention_prefill_rb4_exact(void) { 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 RB4 exactness " + "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, @@ -4016,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. */ @@ -5115,9 +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_rb4_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 8d036c1197..15c73203ba 100644 --- a/tests/test_mxfp4_metal.c +++ b/tests/test_mxfp4_metal.c @@ -623,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, @@ -653,7 +653,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_half_lut_baseline, @@ -692,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, @@ -741,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, @@ -786,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, @@ -826,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, @@ -883,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, @@ -914,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, @@ -959,7 +959,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, @@ -994,7 +994,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, From eb18f5725324764d3513b76e3884b4e79a997ad3 Mon Sep 17 00:00:00 2001 From: ivanfioravanti Date: Fri, 21 Aug 2026 12:04:12 +0200 Subject: [PATCH 06/16] metal: add commit-only GPU stage-counter profiler The end-and-wait decode stage profiler serializes the token, so its stage times mix in per-boundary CPU waits and disable the concurrent shared-expert overlap. Add a diagnostic that keeps the production schedule: each stage boundary commits the open batch command buffer without waiting and records it, and after the token every stage's GPU busy span (GPUEndTime minus GPUStartTime) is printed. Only the serializing profiler now disqualifies concurrent dispatch, so DS4_METAL_STAGE_COUNTERS=1 measures the real schedule; per-token total-cb-busy matches production GPU-busy time and generated tokens are unchanged. Document the attribution it produced on M3 Ultra: dense Q8_0 matvecs already stream at the memory wall, while the per-layer HC pre/post kernels, the Q/KV norm-RoPE-store dispatch, and the routed-MoE window hold the remaining slack. --- ds4.c | 23 ++++++++++++++--- ds4_gpu.h | 7 ++++++ ds4_metal.m | 58 ++++++++++++++++++++++++++++++++++++++++++- speed-bench/README.md | 28 +++++++++++++++++++++ 4 files changed, 112 insertions(+), 4 deletions(-) diff --git a/ds4.c b/ds4.c index 6934aee9f2..9b6fb640f1 100644 --- a/ds4.c +++ b/ds4.c @@ -22740,6 +22740,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 +22767,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 && @@ -25159,7 +25164,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 +25183,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) && @@ -28841,6 +28846,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 +28975,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 +29000,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, @@ -31984,12 +31998,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); 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; diff --git a/ds4_gpu.h b/ds4_gpu.h index 8c33e7b658..8374c5952f 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); diff --git a/ds4_metal.m b/ds4_metal.m index 0c2752bcaf..cd68c7fb51 100644 --- a/ds4_metal.m +++ b/ds4_metal.m @@ -999,6 +999,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 @@ -1010,7 +1066,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; diff --git a/speed-bench/README.md b/speed-bench/README.md index 446b242130..072330ec6d 100644 --- a/speed-bench/README.md +++ b/speed-bench/README.md @@ -27,6 +27,34 @@ 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`. +### 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 schedule A/B Build the balanced, same-engine Metal decode comparison with: From 709d4b4102bd2f52aef5f0841d65a5279cb791f0 Mon Sep 17 00:00:00 2001 From: ivanfioravanti Date: Fri, 21 Aug 2026 15:07:53 +0200 Subject: [PATCH 07/16] docs: record DSpark-on-Metal measurements and work order Non-strict DSpark peaks at 37.8 t/s on M3 Ultra (vs 43.7 plain): the batch verify costs 46-60 ms against a ~29-33 ms two-row floor (uniform 1.5-2.5x per-stage excess, decomposed with the commit-only stage counters), plain decode inside a DSpark session costs ~51 ms/token from hidden-state capture, and the draft proposes on 54-85% of cycles. Scheduler env tuning was swept and documented; making speculation profitable needs the N<=6 microbatch verifier, capture-light plain decode, and a cheaper propose chain. --- speed-bench/README.md | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/speed-bench/README.md b/speed-bench/README.md index 072330ec6d..a9acf4c40f 100644 --- a/speed-bench/README.md +++ b/speed-bench/README.md @@ -27,6 +27,33 @@ 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. +2. Plain decode inside a DSpark session costs ~51 ms/token (23 ms decode + + per-layer hidden-state capture and session bookkeeping), taxing every + non-speculated token. +3. The draft propose chain costs ~3-8 ms/cycle; with the current + confidence gating it proposes on only ~54-85% of cycles. + +Scheduler 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), but peak measured generation stayed +at 37.8 t/s because of the three costs above. 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 From bb1d354d3466ebfbb18970d609c84bd8b76af039 Mon Sep 17 00:00:00 2001 From: ivanfioravanti Date: Fri, 21 Aug 2026 15:17:11 +0200 Subject: [PATCH 08/16] docs: correct DSpark plain-decode cost and record knob sweep Token-level profiling inside a DSpark session shows plain decode at a normal ~23.3 ms; the earlier ~51 ms figure misattributed post-accept first-token evals. The real per-cycle costs are the confidence-gated propose (3-8 ms, declining 45-75% of cycles) and ~4-5 ms of bookkeeping. Best swept configuration: 39.5 t/s at --dspark-confidence 0.75 with NO_DRAFT_SKIP=0, still below the 43.7 t/s plain equilibrium. --- speed-bench/README.md | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/speed-bench/README.md b/speed-bench/README.md index a9acf4c40f..06d183ffa8 100644 --- a/speed-bench/README.md +++ b/speed-bench/README.md @@ -41,16 +41,20 @@ decode (~43.7 t/s) today; scheduler tuning alone cannot fix them: 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. -2. Plain decode inside a DSpark session costs ~51 ms/token (23 ms decode + - per-layer hidden-state capture and session bookkeeping), taxing every - non-speculated token. -3. The draft propose chain costs ~3-8 ms/cycle; with the current - confidence gating it proposes on only ~54-85% of cycles. - -Scheduler 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), but peak measured generation stayed -at 37.8 t/s because of the three costs above. Strict mode (`--dspark-strict`) +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. 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. From 14409484fa550bfdbdb7dcae566680a18c73d857 Mon Sep 17 00:00:00 2001 From: ivanfioravanti Date: Fri, 21 Aug 2026 15:36:13 +0200 Subject: [PATCH 09/16] docs: record dispatch-routing microbatch increments as measured negative Per-row routed MoE via the single-token static kernels (bit-identical, verify_layer unchanged) and per-row HC pre via the decode fused producer (valid non-strict output, 39.3 vs 40.2 t/s) both confirm the batch-verify excess is inside the batch kernels' execution rather than dispatch routing. The microbatch verifier needs genuine small-N batched kernels. --- speed-bench/README.md | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/speed-bench/README.md b/speed-bench/README.md index 06d183ffa8..7dc584c388 100644 --- a/speed-bench/README.md +++ b/speed-bench/README.md @@ -54,7 +54,15 @@ 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. Strict mode (`--dspark-strict`) +43.7 t/s plain-decode equilibrium. + +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. From 06654cda4705f9168df00784d1ab1cec3c6d2d71 Mon Sep 17 00:00:00 2001 From: ivanfioravanti Date: Fri, 21 Aug 2026 15:54:31 +0200 Subject: [PATCH 10/16] docs: record dual-row HC-pre verify kernel result and strict oracle The N=2 dual-row producer was bit-exact under the --dspark-strict oracle (after fixing a kernel-argument misbinding the oracle caught) but measured no per-verify gain and collapsed non-strict drafting on the test prompt. The strict-mode oracle is documented as the validation tool for future verify-kernel work; the HC-pre excess was overestimated by serialized stage decomposition. --- speed-bench/README.md | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/speed-bench/README.md b/speed-bench/README.md index 7dc584c388..9ef85c265b 100644 --- a/speed-bench/README.md +++ b/speed-bench/README.md @@ -56,6 +56,17 @@ 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 From 0dd11a5486332f8322fabefe12bc4a7283c787e1 Mon Sep 17 00:00:00 2001 From: ivanfioravanti Date: Fri, 21 Aug 2026 15:58:12 +0200 Subject: [PATCH 11/16] docs: recalibrate DSpark verify floor after third negative increment Per-row MoE, per-row HC pre, and a strict-oracle-verified dual-row HC-pre kernel each recovered nothing measurable, and the batch MoE already runs at its distinct-expert floor. The perfect-sharing verify floor (~29-33 ms) looks undeliverable on this GPU: the N=2 verify near 50 ms is close to its real floor, so speculation is unlikely to beat plain decode on M3 Ultra. --- speed-bench/README.md | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/speed-bench/README.md b/speed-bench/README.md index 9ef85c265b..2f45865e20 100644 --- a/speed-bench/README.md +++ b/speed-bench/README.md @@ -41,6 +41,13 @@ decode (~43.7 t/s) today; scheduler tuning alone cannot fix them: 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 From 0cdb7f0bed7cbed77560e25072694f85acb76ee5 Mon Sep 17 00:00:00 2001 From: ivanfioravanti Date: Fri, 21 Aug 2026 16:06:24 +0200 Subject: [PATCH 12/16] docs: add decode-campaign handoff for next session Cold-restart kit: verified baselines with reproduction commands, the 22.6 ms token ledger, the closed-avenue list (8 kernel variants, knob sweeps, 3 microbatch increments), tool usage (commit-only stage profiler, strict-mode oracle, balanced A/B harness), hard-won gotchas, and the go/grind/accept/new-hardware decision that must be made first. --- speed-bench/DECODE-CAMPAIGN-HANDOFF.md | 129 +++++++++++++++++++++++++ speed-bench/README.md | 5 + 2 files changed, 134 insertions(+) create mode 100644 speed-bench/DECODE-CAMPAIGN-HANDOFF.md diff --git a/speed-bench/DECODE-CAMPAIGN-HANDOFF.md b/speed-bench/DECODE-CAMPAIGN-HANDOFF.md new file mode 100644 index 0000000000..1961b788d8 --- /dev/null +++ b/speed-bench/DECODE-CAMPAIGN-HANDOFF.md @@ -0,0 +1,129 @@ +# Decode >45 t/s campaign — handoff (session of Aug 21, M3 Ultra) + +Goal: push greedy decode past **45 tokens/s** (towards 50) on the MXFP4 +`ds4flash.gguf` model, bit-exact. Final verified state: **43.2–44.0 t/s**, +all paths in agreement. Goal not reached; every in-session path was built, +measured, and closed. This file is the cold-restart kit: state, evidence, +tools, closed avenues, and the decision the next session must make first. + +## Where everything is + +- Working tree: clean at `814a933`; tests `make test` PASS 44/44; bit-exact + output md5 for the standard prompt unchanged (`db0c504c…`). +- Campaign commits (this session): `11689e1` (commit-only GPU stage profiler + + docs), `b40f33c`/`5f9da0a`/`06ca424`/`c2084a1`/`814a933` (DSpark + measurements, corrections, negative-result records, floor recalibration). +- `speed-bench/README.md` holds the full measurement record: per-stage decode + ledger, every A/B protocol, and the DSpark-on-Metal work order. Read the + sections "Metal decode stage GPU counters" and "DSpark speculation on M3 + Ultra" before doing anything. + +## Verified numbers to trust (and how to reproduce) + +| metric | value | command | +|---|---|---| +| CLI decode | 43.2–43.9 t/s | `./ds4 -m ds4flash.gguf -p "Write a short story about a lighthouse keeper." -c 8192 -n 128 --temp 0` | +| repo bench | 43.77 steady @ctx2048 | `./ds4-bench -m ds4flash.gguf --prompt-file speed-bench/promessi_sposi.txt --ctx-start 2048 --ctx-max 2048 --gen-tokens 96` | +| balanced harness | 43.38 t/s | `make metal-decode-schedule-bench && ./speed-bench/metal_decode_schedule_bench -m ds4flash.gguf --include-selection --tokens 512` | +| best DSpark | 39.5–40.2 t/s | `./ds4 -m ds4flash.gguf --mtp gguf/DeepSeek-V4-Flash-DSpark-support.gguf --dspark --dspark-confidence 0.75` (+`DS4_DSPARK_SCHEDULER_NO_DRAFT_SKIP=0`) | + +Thermal envelope is ±2%: sustained runs sit ~43.2–43.5, first run on a cool +machine reaches 43.9–44.01. Always compare via the interleaved harness, and +let the machine idle ~60s after heavy runs (a transient 2–10× slowdown right +after sustained benching was observed repeatedly; it recovers by itself). + +## The token ledger (22.6 ms GPU busy; encode 0.7 ms hidden by split-flush) + +Per layer (µs, commit-only counters, short ctx): routed MoE 139 (~floor for +6×12.6 MB experts at ~550 GB/s effective), attention core 44.3 (flat vs +context; latency floor), attn output A+B 111 (645 GB/s ≈ wall), Q-lora path +120 (quad kernel 41 + q_b 59 + norm/rope 21), HC pre 2×19.4 (structural +floor), shared/router overlapped 44, KV staging 7.8. Weights memory floor +≈11.6 ms/token; the ~10 ms above it is distributed latency that resists +every single-kernel fix tried (see "Closed avenues"). Boundary tail: +CPU argmax 35.7 µs, loop ~0.05 ms recoverable, remainder wake/launch +latency. Bit-exact recoverable stack sums to ~0.55 ms < the ~0.7 ms needed +for 45 t/s. + +## Closed avenues — do not redo (details + numbers in README) + +1. Eight bit-exact kernel variants, all validated bit-exact, all ≤0.03%: + HC tgstash, HC rows12 (10-TG sibling), quad NR1, packed attention sg16 + and sg32 (both +8 ms/token systemically — more parallelism throttles the + whole token on this GPU), FP8 block one-pass amax, down r4 (slots6 and + static paths), plus an earlier q8 nr0 tune (warm-cache artifact only). +2. DSpark strict: 43.07 t/s (no gain by design). Non-strict: peaks 40.2; + knob space fully swept (confidence 0.75 optimal, scheduler pauses off). +3. Three microbatch increments: per-row routed MoE (bit-identical, + verify_layer unchanged), per-row HC pre (slower), and a genuine dual-row + HC-pre kernel — proven bit-exact via the strict oracle, no per-verify + gain. Conclusion: the N=2 verify near 50 ms is close to its real floor; + the "perfect sharing" 29–33 ms floor is likely undeliverable here, so + speculation probably never beats plain decode on M3 Ultra. +4. Split schedule (`DS4_METAL_GRAPH_TOKEN_SPLIT_LAYERS`): default 4 optimal. +5. Readback-bubble premise: disproved — MXFP4 static-trip reads expert ids + on-device; there is no CPU readback in decode. +6. Stale CSVs in speed-bench/ predate current code; regenerate before + comparing historical numbers. + +## Tools built this session (reuse them) + +- **Commit-only stage profiler** (committed): + `DS4_METAL_DECODE_STAGE_PROFILE=1 DS4_METAL_STAGE_COUNTERS=1 ./ds4 …` + prints per-stage GPU busy spans whose sum matches production GPU-busy + (~22.5 ms/token) — trustworthy, unlike the end-and-wait profiler which + inflates stages ~6× and serializes the schedule. For the batch/verify + path also export `DS4_METAL_LAYER_STAGE_PROFILE=1` (and see the reset/ + report wrap pattern used around the verify call in the session log). +- **Strict-mode oracle** for any verify-kernel work: `--dspark-strict` + output must match plain decode md5 exactly. It caught a real 5-argument + kernel misbinding that ordinary testing missed. +- **Balanced A/B harness** (`speed-bench/metal_decode_schedule_bench`) with + `--candidate-env NAME --include-selection`: interleaved, bit-exactness + enforced over full-vocab logits. Acceptance threshold the project used: + ≥0.3%. +- Microbench pattern (cold-cycling 6 weight buffers to defeat L2) for any + standalone kernel work — note the lesson: warm-cache microbench wins + (e.g. q8 nr0=8) evaporate in production cold streaming. + +## Gotchas learned the hard way + +- Adding Metal kernel parameters: insert new buffer args **after** existing + ones or renumber the host bindings to match — a mid-signature insertion + silently misbinds everything downstream (caught only by the strict oracle). +- The end-and-wait stage profiler changes the schedule (it disables the + concurrent shared-expert overlap); the commit-only mode does not. +- `misc/` is gitignored; anything that must survive belongs in a tracked + path like speed-bench/. +- One `ds4` instance at a time (instance lock is intentional; 145 GB + resident model). +- DSpark stats: per-verify averages = verify_layer/(full+partial), not + per-cycle; the two disagree wildly when no_draft is high. + +## Decision needed before any work resumes + +The session's two objectives — "bit exact" (first message) and ">45 t/s" +(goal) — are jointly infeasible on this hardware per the committed +arithmetic. Next session must pick first: + +1. **go** — multi-day small-N batched-kernel DSpark verifier build, dropping + bit-exactness. Fair warning: triple-confirmed evidence says it likely + tops out below 45 on M3 Ultra anyway. If attempted, start from the + per-stage work order in README's DSpark section; batched output + projection and N≤2 KV-sharing attention are the only stages left + untried, and expectations should be low. +2. **grind** — bit-exact persistent-kernel work (HC epilogue tail-fusions + ~0.3 ms + KV-staging elimination ~0.25 ms + boundary ~0.05 ms ≈ 0.6 ms + best case → ~44.5 t/s). The KV-staging direct-read kernel was started + once (raw-rows-first layout, zero-mask semantics worked out) but + diverged bit-wise and was cut; the addressing notes are in the session + history — the layout facts in README are verified. +3. **accept** — 43.3–44.0 t/s is the M3 Ultra equilibrium; the campaign + artifacts stand as the deliverable. +4. **different hardware** — M5-class parts change the latency-floor math + (more L2, different power behavior); the profiler + ledger apply + as-is there. + +Quick first command next session: +`make && ./ds4 -m ds4flash.gguf -p "Write a short story about a lighthouse keeper." -c 8192 -n 128 --temp 0` +→ expect ~43.3–43.9 t/s on a cool machine; then decide. diff --git a/speed-bench/README.md b/speed-bench/README.md index 2f45865e20..854e37a106 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: From 28a09c0baebfe32e248884f20e4637443374448d Mon Sep 17 00:00:00 2001 From: ivanfioravanti Date: Sat, 22 Aug 2026 08:20:03 +0200 Subject: [PATCH 13/16] metal: greedy chain decode for sessions/ds4-eval, decode tail probes Round 2 (carried in working tree): raw-layer gathered attention parity (+2.7%) and CLI greedy chain decode (+1.75%), 43.46 -> 45.51 t/s, bit-exact. Round 3: ds4-eval decodes in session chain bursts (ds4_session_eval_chain_greedy) outside the think-close window, 44.5 -> 45.2 t/s, traces bit-identical; DS4_DISABLE_GREEDY_CHAIN=1 covers both paths. MoE down-sum6+HC4 tail fusion landed bit-exact, speed-neutral (rollback DS4_METAL_DISABLE_DECODE_MOE_HC_FUSION). KV-staging direct read landed bit-exact but -7% (per-head F32 re-read/re-convert amplification), gated off (opt-in DS4_METAL_ENABLE_DECODE_RAW_DIRECT_KV). make test 44/44; CLI transcript md5 db0c504c... unchanged. --- ds4.c | 516 ++++++++++++++++++++-- ds4.h | 14 + ds4_cuda.cu | 15 +- ds4_eval.c | 129 ++++++ ds4_gpu.h | 28 +- ds4_metal.m | 367 ++++++++++++++- metal/dsv4_rope.metal | 61 +++ metal/flash_attn.metal | 317 +++++++++++++ metal/moe.metal | 128 ++++++ rocm/ds4_rocm_glm.cuh | 3 +- rocm/ds4_rocm_moe_launch.cuh | 8 +- speed-bench/DECODE-CAMPAIGN-HANDOFF.md | 244 +++++----- speed-bench/README.md | 113 +++++ speed-bench/m3_ultra_mxfp4.csv | 33 ++ speed-bench/m3_ultra_mxfp4_r2.csv | 33 ++ speed-bench/m3_ultra_mxfp4_r2_ts.svg | 50 +++ speed-bench/m3_ultra_mxfp4_rebased.csv | 33 ++ speed-bench/m3_ultra_mxfp4_rebased_r2.csv | 10 + speed-bench/m3_ultra_mxfp4_ts.svg | 50 +++ tests/test_mxfp4_metal.c | 4 +- 20 files changed, 1985 insertions(+), 171 deletions(-) create mode 100644 speed-bench/m3_ultra_mxfp4.csv create mode 100644 speed-bench/m3_ultra_mxfp4_r2.csv create mode 100644 speed-bench/m3_ultra_mxfp4_r2_ts.svg create mode 100644 speed-bench/m3_ultra_mxfp4_rebased.csv create mode 100644 speed-bench/m3_ultra_mxfp4_rebased_r2.csv create mode 100644 speed-bench/m3_ultra_mxfp4_ts.svg diff --git a/ds4.c b/ds4.c index 9b6fb640f1..a7bac3e2bb 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 @@ -24678,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, @@ -24693,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, @@ -24999,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__) @@ -25080,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__) @@ -25266,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), @@ -25472,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), @@ -25568,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), @@ -25589,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(); @@ -25797,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), @@ -25840,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, @@ -25876,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 */ @@ -28363,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; @@ -28394,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 @@ -31577,7 +31694,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); @@ -32002,7 +32120,7 @@ static bool metal_graph_eval_token_raw_swa( 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; @@ -32036,6 +32154,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) { @@ -32491,7 +32774,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, @@ -32552,7 +32835,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), @@ -44395,7 +44678,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, @@ -53128,6 +53412,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, @@ -53262,7 +53559,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); @@ -53270,6 +53611,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); @@ -69001,6 +69345,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, @@ -69451,7 +69911,8 @@ static int ds4_sessions_eval_batch_metal( items[i].token, pos, true, - false); + false, + NULL); } } } @@ -73553,7 +74014,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 e6dae1b9f0..e8a09e651b 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 37ea93e1d1..3d3746ba9c 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; diff --git a/ds4_eval.c b/ds4_eval.c index bcb351453e..bd7fda42bd 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 8374c5952f..6f2a311ea6 100644 --- a/ds4_gpu.h +++ b/ds4_gpu.h @@ -2443,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, @@ -2745,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, diff --git a/ds4_metal.m b/ds4_metal.m index cd68c7fb51..3f9513009d 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; @@ -5603,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; @@ -7889,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 || @@ -7897,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; @@ -10412,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; @@ -26418,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; } @@ -26507,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, @@ -28258,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 && @@ -28328,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, @@ -28370,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), @@ -28395,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, @@ -28471,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]; @@ -29938,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; @@ -31211,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, @@ -38916,6 +39072,101 @@ 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), + NULL, + hash_rows, + 1, + n_expert, + n_expert_used, + expert_weight_scale, + has_bias && !hash_mode, + hash_mode); + 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, @@ -39151,10 +39402,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 @@ -39209,6 +39466,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) { @@ -39384,6 +39664,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; @@ -39429,6 +39721,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 ? @@ -41640,6 +41934,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, @@ -41655,6 +42000,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, @@ -41891,7 +42237,8 @@ int ds4_gpu_routed_moe_batch_tensor( x, NULL, layer_index, - false); + false, + NULL, NULL, NULL, NULL, NULL); } @autoreleasepool { diff --git a/metal/dsv4_rope.metal b/metal/dsv4_rope.metal index b772435843..81f57db0d9 100644 --- a/metal/dsv4_rope.metal +++ b/metal/dsv4_rope.metal @@ -981,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 e3232a72d8..cc0cfcc466 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 27f11ad6fc..9ec961fa95 100644 --- a/metal/moe.metal +++ b/metal/moe.metal @@ -6700,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 e0bface992..f9692a0d13 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 16ad7bcb4b..40b72ce94b 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 index 1961b788d8..88c040385d 100644 --- a/speed-bench/DECODE-CAMPAIGN-HANDOFF.md +++ b/speed-bench/DECODE-CAMPAIGN-HANDOFF.md @@ -1,129 +1,117 @@ -# Decode >45 t/s campaign — handoff (session of Aug 21, M3 Ultra) - -Goal: push greedy decode past **45 tokens/s** (towards 50) on the MXFP4 -`ds4flash.gguf` model, bit-exact. Final verified state: **43.2–44.0 t/s**, -all paths in agreement. Goal not reached; every in-session path was built, -measured, and closed. This file is the cold-restart kit: state, evidence, -tools, closed avenues, and the decision the next session must make first. - -## Where everything is - -- Working tree: clean at `814a933`; tests `make test` PASS 44/44; bit-exact - output md5 for the standard prompt unchanged (`db0c504c…`). -- Campaign commits (this session): `11689e1` (commit-only GPU stage profiler - + docs), `b40f33c`/`5f9da0a`/`06ca424`/`c2084a1`/`814a933` (DSpark - measurements, corrections, negative-result records, floor recalibration). -- `speed-bench/README.md` holds the full measurement record: per-stage decode - ledger, every A/B protocol, and the DSpark-on-Metal work order. Read the - sections "Metal decode stage GPU counters" and "DSpark speculation on M3 - Ultra" before doing anything. - -## Verified numbers to trust (and how to reproduce) - -| metric | value | command | +# 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 | |---|---|---| -| CLI decode | 43.2–43.9 t/s | `./ds4 -m ds4flash.gguf -p "Write a short story about a lighthouse keeper." -c 8192 -n 128 --temp 0` | -| repo bench | 43.77 steady @ctx2048 | `./ds4-bench -m ds4flash.gguf --prompt-file speed-bench/promessi_sposi.txt --ctx-start 2048 --ctx-max 2048 --gen-tokens 96` | -| balanced harness | 43.38 t/s | `make metal-decode-schedule-bench && ./speed-bench/metal_decode_schedule_bench -m ds4flash.gguf --include-selection --tokens 512` | -| best DSpark | 39.5–40.2 t/s | `./ds4 -m ds4flash.gguf --mtp gguf/DeepSeek-V4-Flash-DSpark-support.gguf --dspark --dspark-confidence 0.75` (+`DS4_DSPARK_SCHEDULER_NO_DRAFT_SKIP=0`) | - -Thermal envelope is ±2%: sustained runs sit ~43.2–43.5, first run on a cool -machine reaches 43.9–44.01. Always compare via the interleaved harness, and -let the machine idle ~60s after heavy runs (a transient 2–10× slowdown right -after sustained benching was observed repeatedly; it recovers by itself). - -## The token ledger (22.6 ms GPU busy; encode 0.7 ms hidden by split-flush) - -Per layer (µs, commit-only counters, short ctx): routed MoE 139 (~floor for -6×12.6 MB experts at ~550 GB/s effective), attention core 44.3 (flat vs -context; latency floor), attn output A+B 111 (645 GB/s ≈ wall), Q-lora path -120 (quad kernel 41 + q_b 59 + norm/rope 21), HC pre 2×19.4 (structural -floor), shared/router overlapped 44, KV staging 7.8. Weights memory floor -≈11.6 ms/token; the ~10 ms above it is distributed latency that resists -every single-kernel fix tried (see "Closed avenues"). Boundary tail: -CPU argmax 35.7 µs, loop ~0.05 ms recoverable, remainder wake/launch -latency. Bit-exact recoverable stack sums to ~0.55 ms < the ~0.7 ms needed -for 45 t/s. - -## Closed avenues — do not redo (details + numbers in README) - -1. Eight bit-exact kernel variants, all validated bit-exact, all ≤0.03%: - HC tgstash, HC rows12 (10-TG sibling), quad NR1, packed attention sg16 - and sg32 (both +8 ms/token systemically — more parallelism throttles the - whole token on this GPU), FP8 block one-pass amax, down r4 (slots6 and - static paths), plus an earlier q8 nr0 tune (warm-cache artifact only). -2. DSpark strict: 43.07 t/s (no gain by design). Non-strict: peaks 40.2; - knob space fully swept (confidence 0.75 optimal, scheduler pauses off). -3. Three microbatch increments: per-row routed MoE (bit-identical, - verify_layer unchanged), per-row HC pre (slower), and a genuine dual-row - HC-pre kernel — proven bit-exact via the strict oracle, no per-verify - gain. Conclusion: the N=2 verify near 50 ms is close to its real floor; - the "perfect sharing" 29–33 ms floor is likely undeliverable here, so - speculation probably never beats plain decode on M3 Ultra. -4. Split schedule (`DS4_METAL_GRAPH_TOKEN_SPLIT_LAYERS`): default 4 optimal. -5. Readback-bubble premise: disproved — MXFP4 static-trip reads expert ids - on-device; there is no CPU readback in decode. -6. Stale CSVs in speed-bench/ predate current code; regenerate before - comparing historical numbers. - -## Tools built this session (reuse them) - -- **Commit-only stage profiler** (committed): - `DS4_METAL_DECODE_STAGE_PROFILE=1 DS4_METAL_STAGE_COUNTERS=1 ./ds4 …` - prints per-stage GPU busy spans whose sum matches production GPU-busy - (~22.5 ms/token) — trustworthy, unlike the end-and-wait profiler which - inflates stages ~6× and serializes the schedule. For the batch/verify - path also export `DS4_METAL_LAYER_STAGE_PROFILE=1` (and see the reset/ - report wrap pattern used around the verify call in the session log). -- **Strict-mode oracle** for any verify-kernel work: `--dspark-strict` - output must match plain decode md5 exactly. It caught a real 5-argument - kernel misbinding that ordinary testing missed. -- **Balanced A/B harness** (`speed-bench/metal_decode_schedule_bench`) with - `--candidate-env NAME --include-selection`: interleaved, bit-exactness - enforced over full-vocab logits. Acceptance threshold the project used: - ≥0.3%. -- Microbench pattern (cold-cycling 6 weight buffers to defeat L2) for any - standalone kernel work — note the lesson: warm-cache microbench wins - (e.g. q8 nr0=8) evaporate in production cold streaming. - -## Gotchas learned the hard way - -- Adding Metal kernel parameters: insert new buffer args **after** existing - ones or renumber the host bindings to match — a mid-signature insertion - silently misbinds everything downstream (caught only by the strict oracle). -- The end-and-wait stage profiler changes the schedule (it disables the - concurrent shared-expert overlap); the commit-only mode does not. -- `misc/` is gitignored; anything that must survive belongs in a tracked - path like speed-bench/. -- One `ds4` instance at a time (instance lock is intentional; 145 GB - resident model). -- DSpark stats: per-verify averages = verify_layer/(full+partial), not - per-cycle; the two disagree wildly when no_draft is high. - -## Decision needed before any work resumes - -The session's two objectives — "bit exact" (first message) and ">45 t/s" -(goal) — are jointly infeasible on this hardware per the committed -arithmetic. Next session must pick first: - -1. **go** — multi-day small-N batched-kernel DSpark verifier build, dropping - bit-exactness. Fair warning: triple-confirmed evidence says it likely - tops out below 45 on M3 Ultra anyway. If attempted, start from the - per-stage work order in README's DSpark section; batched output - projection and N≤2 KV-sharing attention are the only stages left - untried, and expectations should be low. -2. **grind** — bit-exact persistent-kernel work (HC epilogue tail-fusions - ~0.3 ms + KV-staging elimination ~0.25 ms + boundary ~0.05 ms ≈ 0.6 ms - best case → ~44.5 t/s). The KV-staging direct-read kernel was started - once (raw-rows-first layout, zero-mask semantics worked out) but - diverged bit-wise and was cut; the addressing notes are in the session - history — the layout facts in README are verified. -3. **accept** — 43.3–44.0 t/s is the M3 Ultra equilibrium; the campaign - artifacts stand as the deliverable. -4. **different hardware** — M5-class parts change the latency-floor math - (more L2, different power behavior); the profiler + ledger apply - as-is there. - -Quick first command next session: -`make && ./ds4 -m ds4flash.gguf -p "Write a short story about a lighthouse keeper." -c 8192 -n 128 --temp 0` -→ expect ~43.3–43.9 t/s on a cool machine; then decide. +| 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 854e37a106..576518922c 100644 --- a/speed-bench/README.md +++ b/speed-bench/README.md @@ -117,6 +117,119 @@ 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: diff --git a/speed-bench/m3_ultra_mxfp4.csv b/speed-bench/m3_ultra_mxfp4.csv new file mode 100644 index 0000000000..76978425dd --- /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 0000000000..351dc79c0f --- /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 0000000000..34ea378f1c --- /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 0000000000..266c9313d8 --- /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 0000000000..d045ee148e --- /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 0000000000..eb97aa7d2a --- /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/tests/test_mxfp4_metal.c b/tests/test_mxfp4_metal.c index 15c73203ba..f5d5a9cc36 100644 --- a/tests/test_mxfp4_metal.c +++ b/tests/test_mxfp4_metal.c @@ -318,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( @@ -354,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( From 909d708604744750663def10d297d7b2cdb18ab6 Mon Sep 17 00:00:00 2001 From: Ivan Fioravanti Date: Tue, 1 Sep 2026 18:13:03 +0200 Subject: [PATCH 14/16] metal: fix stale ds4_gpu_encode_router_select call in dev-token path ds4_gpu_router_select_tensor_devtoken still passed 20 arguments after ds4_gpu_encode_router_select gained visual_bias, single_token, vocab_size, and mixed_visual parameters, breaking the build. Pass nil/0/NULL/0/false for the non-visual dev-token path, matching the other three call sites. --- ds4_metal.m | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/ds4_metal.m b/ds4_metal.m index 3f9513009d..9f32ebff5a 100644 --- a/ds4_metal.m +++ b/ds4_metal.m @@ -39150,14 +39150,18 @@ int ds4_gpu_router_select_tensor_devtoken( 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); + hash_mode, + false); if (!had_batch) { ok = ds4_gpu_end_commands() != 0 && ok; } From 20b2877fc8dba0c9fdd0c6008ce26e381d13136a Mon Sep 17 00:00:00 2001 From: Ivan Fioravanti Date: Tue, 1 Sep 2026 18:47:58 +0200 Subject: [PATCH 15/16] metal: port batch indexer-query pruning to M5 Swap the pre-M5 device gate in metal_graph_encode_layer_attention_batch for the shared M5 port helper so resident single-device M5 graphs also skip the four dead indexer-query dispatches on zero-prefix ratio-4 batches while the compressed cache stays at or below top-k. Rollback on M5 is DS4_METAL_DISABLE_M5_BATCH_INDEXER_QUERY_PRUNE; pre-M5 keeps its existing per-feature and aggregate rollbacks. Balanced M5 Max A/B (IQ2_XXS/Q2_K ds4flash.gguf): 753.62/732.83 tok/s at 2048 tokens (+2.84%), 8/8 runs bit-identical (1,034,240 logits). The standard sweep's first frontier measured +8.7% prefill with decode and all ineligible frontiers unchanged. Also measured and rejected on M5 Max (see speed-bench/README.md): indexed prefill RB4 kernels (-0.27..-1.56% vs the MPP dual-heads default) and the Q2 2/32 decode split (-0.6..-0.8% vs the 4/none default). --- ds4.c | 5 +++-- speed-bench/README.md | 21 +++++++++++++++++++++ 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/ds4.c b/ds4.c index a7bac3e2bb..e07ea47177 100644 --- a/ds4.c +++ b/ds4.c @@ -30120,8 +30120,9 @@ static bool metal_graph_encode_layer_attention_batch( !topk_prefill_needed && !g->quality && !g->ssd_streaming && !g->ssd_streaming_cold && g->placement == NULL && g->tp_world < 2u && - ds4_gpu_device_is_pre_m5_apple_silicon() && - getenv("DS4_METAL_DISABLE_PRE_M5_BATCH_INDEXER_QUERY_PRUNE") == NULL; + 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 diff --git a/speed-bench/README.md b/speed-bench/README.md index 576518922c..9433c27730 100644 --- a/speed-bench/README.md +++ b/speed-bench/README.md @@ -344,6 +344,27 @@ 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 From e04243b93f717038109756c6633cc1e0ebe20d33 Mon Sep 17 00:00:00 2001 From: Paperino Date: Mon, 31 Aug 2026 13:11:21 +0200 Subject: [PATCH 16/16] dspark: make the capture notice opt-in via DS4_DSPARK_VERBOSE The line is printed on every session creation, and it reports a build and weights decision the caller did not make and cannot act on. An embedder that creates a session per turn (or per aside, or per sub-agent) has it land in the middle of the user's screen each time. Gated behind DS4_DSPARK_VERBOSE, matching the other DS4_* diagnostic switches. The failure path beside it stays unconditional: that one is news. Co-Authored-By: Claude Opus 5 --- ds4.c | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/ds4.c b/ds4.c index e07ea47177..d30db5a70b 100644 --- a/ds4.c +++ b/ds4.c @@ -65078,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,