diff --git a/ds4_metal.m b/ds4_metal.m index 84af7563f0..415c39fea8 100644 --- a/ds4_metal.m +++ b/ds4_metal.m @@ -159,7 +159,7 @@ static id g_soft_max_f32_4_pipeline; static id g_argsort_f32_i32_desc_pipeline; static id g_argsort_merge_f32_i32_desc_pipeline; -static id g_topk_stream512_pipeline; +static id g_topk_fused512_pipeline; static id g_sum_rows_f32_f32_pipeline; static id g_dsv4_topk_mask_pipeline; static id g_dsv4_topk_mask_scatter_pipeline; @@ -170,6 +170,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_heads16_dual_rb16_pipeline; static id g_dsv4_indexed_attention_heads8_split_pipeline; static id g_dsv4_indexed_attention_heads8_split_reduce_pipeline; static id g_dsv4_softplus_sqrt_pipeline; @@ -7295,16 +7296,16 @@ int ds4_gpu_init(void) { return 0; } - fn = [library newFunctionWithName:@"kernel_topk_stream512"]; + fn = [library newFunctionWithName:@"kernel_topk_fused512"]; if (!fn) { - fprintf(stderr, "ds4: Metal kernel_topk_stream512 function not found\n"); + fprintf(stderr, "ds4: Metal kernel_topk_fused512 function not found\n"); g_queue = nil; g_device = nil; return 0; } - g_topk_stream512_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; - if (!g_topk_stream512_pipeline) { - fprintf(stderr, "ds4: Metal kernel_topk_stream512 pipeline failed: %s\n", + g_topk_fused512_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_topk_fused512_pipeline) { + fprintf(stderr, "ds4: Metal kernel_topk_fused512 pipeline failed: %s\n", [[error localizedDescription] UTF8String]); g_queue = nil; g_device = nil; @@ -7707,6 +7708,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_heads16_dual_rb16_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads16_dual_rb16"); 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 = @@ -7849,6 +7852,7 @@ int ds4_gpu_init(void) { !g_dsv4_indexed_attention_heads8_pipeline || !g_dsv4_indexed_attention_heads8_rb16_pipeline || !g_dsv4_indexed_attention_heads16_dual_pipeline || + !g_dsv4_indexed_attention_heads16_dual_rb16_pipeline || !g_dsv4_indexed_attention_heads8_split_pipeline || !g_dsv4_indexed_attention_heads8_split_reduce_pipeline || !g_dsv4_softplus_sqrt_pipeline || @@ -9191,7 +9195,7 @@ void ds4_gpu_cleanup(void) { g_soft_max_f32_4_pipeline = nil; g_argsort_f32_i32_desc_pipeline = nil; g_argsort_merge_f32_i32_desc_pipeline = nil; - g_topk_stream512_pipeline = nil; + g_topk_fused512_pipeline = nil; g_sum_rows_f32_f32_pipeline = nil; g_dsv4_topk_mask_pipeline = nil; g_dsv4_topk_mask_scatter_pipeline = nil; @@ -9202,6 +9206,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_heads16_dual_rb16_pipeline = nil; g_dsv4_indexed_attention_heads8_split_pipeline = nil; g_dsv4_indexed_attention_heads8_split_reduce_pipeline = nil; g_dsv4_softplus_sqrt_pipeline = nil; @@ -16650,7 +16655,15 @@ int ds4_gpu_indexer_topk_tensor( fprintf(stderr, "ds4: Metal graph indexer top-k received undersized buffers\n"); return 0; } - if (top_k == 512u && n_comp > 1024u && n_tokens >= 32u) { + /* Fuse the two-run M5 case. Wider rows benefit from the canonical + * path's parallel block sorts, so leave them on argsort/merge. */ + if (top_k == 512u && n_comp > 1024u && n_comp <= 2048u && + n_tokens >= 32u && + (ds4_gpu_device_name_contains("M5") || + getenv("DS4_METAL_ENABLE_TOPK_FUSED512") != NULL) && + g_argsort_f32_i32_desc_pipeline.maxTotalThreadsPerThreadgroup >= 1024u && + g_topk_fused512_pipeline.maxTotalThreadsPerThreadgroup >= 1024u && + getenv("DS4_METAL_DISABLE_TOPK_FUSED512") == NULL) { ds4_gpu_kargs_argsort args = { .ne00 = (int32_t)n_comp, .ne01 = (int32_t)n_tokens, @@ -16670,14 +16683,14 @@ int ds4_gpu_indexer_topk_tensor( id cb = ds4_gpu_command_buffer(&owned); if (!cb) return 0; id enc = ds4_gpu_compute_encoder(cb); - [enc setComputePipelineState:g_topk_stream512_pipeline]; + [enc setComputePipelineState:g_topk_fused512_pipeline]; [enc setBytes:&args length:sizeof(args) atIndex:0]; [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:1]; [enc setBuffer:selbuf offset:ds4_gpu_tensor_offset(selected) atIndex:2]; [enc dispatchThreadgroups:MTLSizeMake(n_tokens, 1, 1) - threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + threadsPerThreadgroup:MTLSizeMake(1024, 1, 1)]; ds4_gpu_end_compute_encoder(cb, enc); - return ds4_gpu_finish_command_buffer(cb, owned, "indexer streaming top-k"); + return ds4_gpu_finish_command_buffer(cb, owned, "indexer fused top-k"); } NSUInteger max_threads = g_argsort_f32_i32_desc_pipeline.maxTotalThreadsPerThreadgroup; if (max_threads == 0) max_threads = 256; @@ -26479,6 +26492,11 @@ 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 prefill_dual_heads_rb16 = + prefill_dual_heads && n_tokens >= 32u && + (ds4_gpu_device_name_contains("M5 Max") || + getenv("DS4_METAL_ENABLE_INDEXED_ATTN_DUAL_RB16") != NULL) && + getenv("DS4_METAL_DISABLE_INDEXED_ATTN_DUAL_RB16") == NULL; const uint32_t decode_splits = decode_one_token && !g_quality_mode ? 12u : 1u; const bool split_decode = decode_splits > 1u; @@ -26490,6 +26508,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_dual_heads_rb16 ? + ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads16_dual_rb16_pipeline, + "kernel_dsv4_indexed_mixed_attention_heads16_dual_rb16") : prefill_dual_heads ? ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads16_dual_pipeline, "kernel_dsv4_indexed_mixed_attention_heads16_dual") : @@ -26627,7 +26648,8 @@ 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 || prefill_dual_heads_rb16 ? 16u : 1u) * 128u * 4u * sizeof(uint16_t) atIndex:0]; [enc dispatchThreadgroups: diff --git a/metal/argsort.metal b/metal/argsort.metal index 1a7ecac7e8..618135169b 100644 --- a/metal/argsort.metal +++ b/metal/argsort.metal @@ -274,105 +274,104 @@ kernel void kernel_argsort_merge_f32_i32( // Host-visible merge variant used by DS4 top-k selection. template [[host_name("kernel_argsort_merge_f32_i32_desc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32; -// Exact top-512 selection for wide prefill rows. Scores and row indices are -// packed into one key, giving a deterministic descending total order. The -// threshold is the 512th-best key seen so far, so discarded entries cannot -// enter the final result. -kernel void kernel_topk_stream512( +// Exact fused top-512 selection for wide M5 prefill rows. This reproduces the +// canonical 1024-wide block sort and left-biased merge order, including its +// tie behavior, while keeping intermediate indices in threadgroup memory. +kernel void kernel_topk_fused512( constant ds4_metal_args_argsort & args [[buffer(0)]], device const float *scores [[buffer(1)]], device uint32_t *selected [[buffer(2)]], uint token [[threadgroup_position_in_grid]], - uint tid [[thread_index_in_threadgroup]], - uint ntg [[threads_per_threadgroup]]) { - constexpr uint stream_cap = 1024u; + uint tid [[thread_index_in_threadgroup]]) { + constexpr uint block_width = 1024u; constexpr uint keep_count = 512u; - threadgroup uint64_t keys[stream_cap]; - threadgroup atomic_uint count; - threadgroup uint64_t threshold; + threadgroup int32_t block_idx[block_width]; + threadgroup float block_scores[block_width]; + threadgroup int32_t retained[keep_count]; + threadgroup int32_t merged[keep_count]; if (token >= (uint)args.ne01) return; - if (tid == 0u) { - atomic_store_explicit(&count, 0u, memory_order_relaxed); - threshold = 0u; - } - threadgroup_barrier(mem_flags::mem_threadgroup); - device const float *row = scores + (uint64_t)token * (uint)args.ne00; const uint n_comp = (uint)args.ne00; - const uint start = (uint)(((uint64_t)(token + 1u) * 0x9E3779B9u) % n_comp); - - for (uint base = 0u; base < n_comp; base += ntg) { - const uint i = base + tid; - if (i < n_comp) { - uint col = start + i; - if (col >= n_comp) col -= n_comp; - const uint bits = as_type(row[col]); - const uint ordered = (bits & 0x80000000u) ? ~bits : - (bits ^ 0x80000000u); - const uint64_t key = ((uint64_t)ordered << 32u) | - (uint64_t)(0xffffffffu - col); - if (key > threshold) { - const uint pos = atomic_fetch_add_explicit( - &count, 1u, memory_order_relaxed); - keys[pos] = key; - } - } - threadgroup_barrier(mem_flags::mem_threadgroup); - const uint used = atomic_load_explicit(&count, memory_order_relaxed); - if (used > stream_cap - ntg) { - for (uint j = used + tid; j < stream_cap; j += ntg) keys[j] = 0u; - threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint base = 0u; base < n_comp; base += block_width) { + const uint col = base + tid; + block_idx[tid] = (int32_t)col; + if (col < n_comp) block_scores[tid] = row[col]; + threadgroup_barrier(mem_flags::mem_threadgroup); - for (uint k = 2u; k <= stream_cap; k <<= 1u) { - for (uint j = k >> 1u; j > 0u; j >>= 1u) { - for (uint p = tid; p < stream_cap; p += ntg) { - const uint other = p ^ j; - if (other > p) { - const uint64_t a = keys[p]; - const uint64_t b = keys[other]; - const bool descending = (p & k) == 0u; - if ((descending && a < b) || (!descending && a > b)) { - keys[p] = b; - keys[other] = a; - } + // Keep this comparison network identical to + // kernel_argsort_f32_i32_desc above. In particular, strict float + // comparisons preserve its deterministic, non-stable tie order. + for (uint k = 2u; k <= block_width; k <<= 1u) { + for (uint j = k >> 1u; j > 0u; j >>= 1u) { + const uint other = tid ^ j; + if (other > tid) { + const int32_t lhs = block_idx[tid]; + const int32_t rhs = block_idx[other]; + if ((tid & k) == 0u) { + if (lhs >= (int32_t)n_comp || + (rhs < (int32_t)n_comp && + block_scores[lhs - (int32_t)base] < + block_scores[rhs - (int32_t)base])) { + block_idx[tid] = rhs; + block_idx[other] = lhs; + } + } else { + if (rhs >= (int32_t)n_comp || + (lhs < (int32_t)n_comp && + block_scores[lhs - (int32_t)base] > + block_scores[rhs - (int32_t)base])) { + block_idx[tid] = rhs; + block_idx[other] = lhs; } } - threadgroup_barrier(mem_flags::mem_threadgroup); } + threadgroup_barrier(mem_flags::mem_threadgroup); } - if (tid == 0u) { - threshold = keys[keep_count - 1u]; - atomic_store_explicit(&count, keep_count, memory_order_relaxed); - } + } + + const uint block_len = min(block_width, n_comp - base); + const uint run_len = min(keep_count, block_len); + if (base == 0u) { + if (tid < keep_count) retained[tid] = block_idx[tid]; threadgroup_barrier(mem_flags::mem_threadgroup); + continue; } - } - const uint used = atomic_load_explicit(&count, memory_order_relaxed); - for (uint j = used + tid; j < stream_cap; j += ntg) keys[j] = 0u; - threadgroup_barrier(mem_flags::mem_threadgroup); - for (uint k = 2u; k <= stream_cap; k <<= 1u) { - for (uint j = k >> 1u; j > 0u; j >>= 1u) { - for (uint p = tid; p < stream_cap; p += ntg) { - const uint other = p ^ j; - if (other > p) { - const uint64_t a = keys[p]; - const uint64_t b = keys[other]; - const bool descending = (p & k) == 0u; - if ((descending && a < b) || (!descending && a > b)) { - keys[p] = b; - keys[other] = a; - } + if (tid < keep_count) { + const int k0 = (int)tid; + int low = k0 > (int)run_len ? k0 - (int)run_len : 0; + int high = min(k0, (int)keep_count); + while (low < high) { + const int mid = (low + high) >> 1; + const int32_t idx0 = retained[mid]; + const int32_t idx1 = block_idx[k0 - mid - 1]; + if (row[idx0] >= row[idx1]) { + low = mid + 1; + } else { + high = mid; } } - threadgroup_barrier(mem_flags::mem_threadgroup); + + const int i = low; + const int j = k0 - i; + if (i >= (int)keep_count) { + merged[tid] = block_idx[j]; + } else if (j >= (int)run_len) { + merged[tid] = retained[i]; + } else if (row[retained[i]] >= row[block_idx[j]]) { + merged[tid] = retained[i]; + } else { + merged[tid] = block_idx[j]; + } } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < keep_count) retained[tid] = merged[tid]; + threadgroup_barrier(mem_flags::mem_threadgroup); } - device uint32_t *out = selected + (uint64_t)token * keep_count; - for (uint i = tid; i < keep_count; i += ntg) { - out[i] = 0xffffffffu - (uint)keys[i]; + if (tid < keep_count) { + selected[(uint64_t)token * keep_count + tid] = (uint32_t)retained[tid]; } } diff --git a/metal/dsv4_misc.metal b/metal/dsv4_misc.metal index 8a0025ca39..c62e665e79 100644 --- a/metal/dsv4_misc.metal +++ b/metal/dsv4_misc.metal @@ -5321,7 +5321,8 @@ kernel void kernel_dsv4_indexed_mixed_attention_heads8( // Each simdgroup owns two heads and updates both from one staged K/V row. // This doubles row reuse without increasing the 256-thread workgroup. -kernel void kernel_dsv4_indexed_mixed_attention_heads16_dual( +template +kernel void kernel_dsv4_indexed_mixed_attention_heads16_dual_impl( constant ds4_metal_args_dsv4_indexed_attention &args, device const char *q, device const char *raw_kv, @@ -5372,18 +5373,27 @@ kernel void kernel_dsv4_indexed_mixed_attention_heads16_dual( const uint first = max(first_raw_pos, window_first); const uint last = min(qpos, raw_last_pos); if (first <= last) { - for (uint pos = first; pos <= last; pos++) { - const uint logical = pos - 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); - if (tid < 128) kv_shared[tid] = (half4)src[tid]; + for (uint base = first; base <= last; base += row_block) { + const uint rows = min(row_block, last - base + 1u); + for (uint off = tid; off < rows * 128u; off += 256u) { + const uint rr = off / 128u; + const uint col = off - rr * 128u; + const uint logical = base + rr - 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[col]; + } threadgroup_barrier(mem_flags::mem_threadgroup); - dsv4_attend_shared_h4_row(kv_shared, qa0, qa1, qa2, qa3, - args.scale, lane, Ma, Sa, ao0, ao1, ao2, ao3); - if (head1 < args.n_head) { - dsv4_attend_shared_h4_row(kv_shared, qb0, qb1, qb2, qb3, - args.scale, lane, Mb, Sb, bo0, bo1, bo2, bo3); + for (uint rr = 0u; rr < rows; rr++) { + dsv4_attend_shared_h4_row_at(kv_shared, rr, + 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, rr, + qb0, qb1, qb2, qb3, + args.scale, lane, Mb, Sb, bo0, bo1, bo2, bo3); + } } threadgroup_barrier(mem_flags::mem_threadgroup); } @@ -5392,20 +5402,29 @@ kernel void kernel_dsv4_indexed_mixed_attention_heads16_dual( 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); - for (uint i = 0; i < args.top_k; i++) { - const int32_t idx = row_topk[i]; - if (idx < 0) continue; - if ((uint)idx >= visible) break; - if (tid < 128) { - kv_shared[tid] = dsv4_load_cache_h4(comp_kv, - args.comp_row_stride, (uint)idx, tid, args.comp_kv_f16 != 0u); + for (uint base = 0u; base < args.top_k; base += row_block) { + const uint rows = min(row_block, args.top_k - base); + for (uint off = tid; off < rows * 128u; off += 256u) { + const uint rr = off / 128u; + const uint col = off - rr * 128u; + const int32_t idx = row_topk[base + rr]; + kv_shared[off] = (idx >= 0 && (uint)idx < visible) ? + dsv4_load_cache_h4(comp_kv, args.comp_row_stride, + (uint)idx, col, args.comp_kv_f16 != 0u) : + half4(0.0h); } threadgroup_barrier(mem_flags::mem_threadgroup); - dsv4_attend_shared_h4_row(kv_shared, qa0, qa1, qa2, qa3, - args.scale, lane, Ma, Sa, ao0, ao1, ao2, ao3); - if (head1 < args.n_head) { - dsv4_attend_shared_h4_row(kv_shared, qb0, qb1, qb2, qb3, - args.scale, lane, Mb, Sb, bo0, bo1, bo2, bo3); + for (uint rr = 0u; rr < rows; rr++) { + const int32_t idx = row_topk[base + rr]; + if (idx < 0 || (uint)idx >= visible) continue; + dsv4_attend_shared_h4_row_at(kv_shared, rr, + 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, rr, + qb0, qb1, qb2, qb3, + args.scale, lane, Mb, Sb, bo0, bo1, bo2, bo3); + } } threadgroup_barrier(mem_flags::mem_threadgroup); } @@ -5430,6 +5449,20 @@ kernel void kernel_dsv4_indexed_mixed_attention_heads16_dual( } } +typedef decltype(kernel_dsv4_indexed_mixed_attention_heads16_dual_impl<1u>) + dsv4_indexed_mixed_attention_heads16_dual_t; + +template [[host_name("kernel_dsv4_indexed_mixed_attention_heads16_dual")]] +kernel dsv4_indexed_mixed_attention_heads16_dual_t +kernel_dsv4_indexed_mixed_attention_heads16_dual_impl<1u>; + +// The 16-row specialization keeps attention arithmetic in chronological row +// order but amortizes each pair of threadgroup barriers across a full cache +// block. It mirrors the proven decode row-blocking scheme for dual-head prefill. +template [[host_name("kernel_dsv4_indexed_mixed_attention_heads16_dual_rb16")]] +kernel dsv4_indexed_mixed_attention_heads16_dual_t +kernel_dsv4_indexed_mixed_attention_heads16_dual_impl<16u>; + // 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/tests/ds4_test.c b/tests/ds4_test.c index c745269ea7..5e2f1d0424 100644 --- a/tests/ds4_test.c +++ b/tests/ds4_test.c @@ -4353,6 +4353,237 @@ static void test_metal_hc_rms_scale_project_f16_exact(void) { test_restore_env(disable_env, saved_disable); } +static void test_metal_topk_fused512_exact(void) { + static const uint32_t widths[] = {1537u, 2048u}; + const uint32_t max_n_comp = 2048u; + const uint32_t n_tokens = 33u; + const uint32_t top_k = 512u; + const uint32_t repeats = 4u; + const uint64_t selected_count = (uint64_t)top_k * n_tokens; + const uint64_t max_score_bytes = + (uint64_t)max_n_comp * n_tokens * sizeof(float); + const uint64_t selected_bytes = selected_count * sizeof(uint32_t); + const char *enable_env = "DS4_METAL_ENABLE_TOPK_FUSED512"; + const char *disable_env = "DS4_METAL_DISABLE_TOPK_FUSED512"; + + ds4_gpu_tensor *scores = ds4_gpu_tensor_alloc(max_score_bytes); + ds4_gpu_tensor *reference = ds4_gpu_tensor_alloc(selected_bytes); + ds4_gpu_tensor *candidate = ds4_gpu_tensor_alloc(selected_bytes); + float *scores_host = malloc((size_t)max_score_bytes); + uint32_t *reference_host = malloc((size_t)selected_bytes); + uint32_t *candidate_host = malloc((size_t)selected_bytes); + TEST_ASSERT(scores != NULL); + TEST_ASSERT(reference != NULL); + TEST_ASSERT(candidate != NULL); + TEST_ASSERT(scores_host != NULL); + TEST_ASSERT(reference_host != NULL); + TEST_ASSERT(candidate_host != NULL); + + char *saved_enable = test_save_env(enable_env); + char *saved_disable = test_save_env(disable_env); + size_t total_mismatches = 0; + uint64_t total_compared = 0u; + const bool allocated = scores && reference && candidate && scores_host && + reference_host && candidate_host; + if (allocated) { + TEST_ASSERT(setenv(enable_env, "1", 1) == 0); + for (size_t width_i = 0u; + width_i < sizeof(widths) / sizeof(widths[0]); + width_i++) { + const uint32_t n_comp = widths[width_i]; + const uint64_t score_bytes = + (uint64_t)n_comp * n_tokens * sizeof(float); + // Cover both a unique permutation and heavy ties: the latter + // guards canonical bitonic/merge order, not just set membership. + for (uint32_t pattern = 0u; pattern < 2u; pattern++) { + for (uint32_t token = 0u; token < n_tokens; token++) { + for (uint32_t col = 0u; col < n_comp; col++) { + uint32_t rank = + (col * 2053u + token * 977u) % n_comp; + if (pattern != 0u) rank %= 31u; + scores_host[(uint64_t)token * n_comp + col] = + (float)((int32_t)rank - + (int32_t)(pattern ? 15u : n_comp / 2u)); + } + } + TEST_ASSERT(ds4_gpu_tensor_write( + scores, 0, scores_host, score_bytes) != 0); + + TEST_ASSERT(setenv(disable_env, "1", 1) == 0); + TEST_ASSERT(ds4_gpu_indexer_topk_tensor( + reference, scores, n_comp, + n_tokens, top_k) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + reference, 0, + reference_host, selected_bytes) != 0); + + TEST_ASSERT(unsetenv(disable_env) == 0); + for (uint32_t repeat = 0u; repeat < repeats; repeat++) { + TEST_ASSERT(ds4_gpu_indexer_topk_tensor( + candidate, scores, n_comp, + n_tokens, top_k) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + candidate, 0, + candidate_host, selected_bytes) != 0); + for (uint64_t i = 0u; i < selected_count; i++) { + const bool mismatch = + candidate_host[i] != reference_host[i]; + total_mismatches += mismatch; + } + total_compared += selected_count; + } + } + } + } + + fprintf(stderr, + "ds4-test: fused top-512 exact " + "rows=%u widths=%zu max_width=%u repeats=%u " + "mismatches=%zu/%llu\n", + n_tokens, sizeof(widths) / sizeof(widths[0]), max_n_comp, + repeats, total_mismatches, + (unsigned long long)total_compared); + TEST_ASSERT(total_mismatches == 0); + + test_restore_env(enable_env, saved_enable); + test_restore_env(disable_env, saved_disable); + free(candidate_host); + free(reference_host); + free(scores_host); + ds4_gpu_tensor_free(candidate); + ds4_gpu_tensor_free(reference); + ds4_gpu_tensor_free(scores); +} + +static void test_metal_indexed_attention_dual_rb16_exact(void) { + const uint32_t n_tokens = 33u; + const uint32_t n_head = 64u; + const uint32_t head_dim = 512u; + const uint32_t n_raw = 128u; + const uint32_t raw_cap = 128u; + const uint32_t n_comp = 1024u; + const uint32_t top_k = 512u; + const uint32_t pos0 = n_comp * 4u; + const uint32_t repeats = 2u; + 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 *enable_env = "DS4_METAL_ENABLE_INDEXED_ATTN_DUAL_RB16"; + const char *disable_env = "DS4_METAL_DISABLE_INDEXED_ATTN_DUAL_RB16"; + + 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); + void *model_raw = NULL; + TEST_ASSERT(posix_memalign(&model_raw, (size_t)page, (size_t)page) == 0); + TEST_ASSERT(q && raw && comp && topk && reference && candidate); + TEST_ASSERT(q_host && raw_host && comp_host && topk_host && + reference_host && candidate_host && model_raw); + + char *saved_enable = test_save_env(enable_env); + char *saved_disable = test_save_env(disable_env); + size_t total_mismatches = 0u; + uint64_t total_compared = 0u; + const bool allocated = q && raw && comp && topk && reference && candidate && + q_host && raw_host && comp_host && topk_host && reference_host && + candidate_host && model_raw; + if (allocated) { + memset(model_raw, 0, (size_t)page); + for (uint32_t h = 0u; h < n_head; h++) { + ((float *)model_raw)[h] = -0.25f + (float)(h % 7u) * 0.03125f; + } + for (uint64_t i = 0u; i < q_count; i++) { + const int32_t v = (int32_t)((i * 17u + i / 97u) % 19u) - 9; + q_host[i] = (float)v * 0.015625f; + } + for (uint64_t i = 0u; i < raw_count; i++) { + const int32_t v = (int32_t)((i * 29u + i / 53u) % 23u) - 11; + raw_host[i] = (float)v * 0.015625f; + } + for (uint64_t i = 0u; i < comp_count; i++) { + const int32_t v = (int32_t)((i * 13u + i / 43u) % 29u) - 14; + comp_host[i] = test_float_to_f16((float)v * 0.015625f); + } + for (uint32_t token = 0u; token < n_tokens; token++) { + for (uint32_t s = 0u; s < top_k; s++) { + topk_host[(uint64_t)token * top_k + s] = + (int32_t)((s * 197u + token * 17u) % n_comp); + } + } + 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_set_model_map(model_raw, page) != 0); + ds4_gpu_set_quality(false); + + TEST_ASSERT(setenv(enable_env, "1", 1) == 0); + TEST_ASSERT(setenv(disable_env, "1", 1) == 0); + TEST_ASSERT(ds4_gpu_attention_indexed_mixed_batch_heads_tensor( + reference, model_raw, page, 0u, + q, raw, comp, 1u, topk, + n_tokens, pos0, n_raw, raw_cap, 0u, + n_comp, top_k, 128u, 4u, n_head, head_dim) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + reference, 0, reference_host, q_bytes) != 0); + + TEST_ASSERT(unsetenv(disable_env) == 0); + for (uint32_t repeat = 0u; repeat < repeats; repeat++) { + TEST_ASSERT(ds4_gpu_attention_indexed_mixed_batch_heads_tensor( + candidate, model_raw, page, 0u, + q, raw, comp, 1u, topk, + n_tokens, pos0, n_raw, raw_cap, 0u, + n_comp, top_k, 128u, 4u, n_head, head_dim) != 0); + TEST_ASSERT(ds4_gpu_tensor_read( + candidate, 0, candidate_host, q_bytes) != 0); + for (uint64_t i = 0u; i < q_count; i++) { + total_mismatches += + memcmp(&candidate_host[i], &reference_host[i], + sizeof(float)) != 0; + } + total_compared += q_count; + } + } + + fprintf(stderr, + "ds4-test: indexed attention dual rb16 exact " + "tokens=%u selected=%u repeats=%u mismatches=%zu/%llu\n", + n_tokens, top_k, repeats, total_mismatches, + (unsigned long long)total_compared); + TEST_ASSERT(total_mismatches == 0u); + + test_restore_env(enable_env, saved_enable); + test_restore_env(disable_env, saved_disable); + free(model_raw); + 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); +} + static void test_metal_router_simd_finalize_exact(void) { typedef struct { const char *name; @@ -4893,6 +5124,8 @@ static void test_metal_kernel_group(void) { test_metal_output_hc_weights4_exact(); test_metal_output_hc_sum_norm_exact(); test_metal_hc_rms_scale_project_f16_exact(); + test_metal_topk_fused512_exact(); + test_metal_indexed_attention_dual_rb16_exact(); test_metal_router_simd_finalize_exact(); test_metal_router_weights_batch_exact(); #endif