diff --git a/ds4.c b/ds4.c index 050ea6cc..27eaca94 100644 --- a/ds4.c +++ b/ds4.c @@ -40925,7 +40925,7 @@ static uint32_t ds41_carry_cap(uint32_t ctx) { X(selected_kv, DS4_N_INDEXER_TOP_K * DS4_N_HEAD_DIM) \ X(block_scores, DS41_INDEX_BATCH * ((g->ctx + 7u) / 8u)) \ X(block_selected, DS41_INDEX_BATCH * 2048u) \ - X(block_mask, (g->ctx + 7u) / 8u) \ + X(block_mask, (g->ctx + 7u) / 8u) X(block_list, 2048u) \ X(heads, DS4_N_HEAD * DS4_N_HEAD_DIM) X(low, DS4_N_OUT_GROUP * DS4_N_LORA_O) \ X(route_logits, DS4_N_EXPERT) X(route_probs, DS4_N_EXPERT) \ X(selected, DS4_N_EXPERT_USED) X(route_weights, DS4_N_EXPERT_USED) \ @@ -41673,12 +41673,17 @@ static bool ds41_attention_select_published(ds41_gpu_graph *g, const ds4_model * if (!ds41_matmul(g->index_q, m, l->indexer_attn_q_b, g->qr, true) || !ds41_rope(g->index_q, DS4_N_INDEXER_HEAD, 128, il, pos, false) || !ds4_gpu_dsv41_quantize(g->index_q, 128, DS4_N_INDEXER_HEAD, DS4_V41_FP4_E8M0) || - !ds41_matmul(g->index_weights, m, l->indexer_proj, g->norm, true) || + !ds41_matmul(g->index_weights, m, l->indexer_proj, g->norm, true)) return false; #ifdef __APPLE__ - !ds4_gpu_dsv41_indexer_score_one_tensor(g->index_scores, g->index_q, g->index_weights, + /* past layer 20 only its candidate blocks can be selected */ + if (il > 20u && (n_comp + 7u) / 8u > 2048u && + ds4_gpu_dsv41_indexer_candidates_tensor(g->selected_comp, g->index_scores, g->block_list, + g->index_q, g->index_weights, g->index_cache[owner], g->block_mask, n_comp, + 1.0f / 64.0f, DS4_N_INDEXER_TOP_K)) return true; + if (!ds4_gpu_dsv41_indexer_score_one_tensor(g->index_scores, g->index_q, g->index_weights, g->index_cache[owner], n_comp, DS4_N_INDEXER_HEAD, 128, 1.0f / 64.0f, false)) return false; #else - !ds4_gpu_dsv41_indexer_scores_batch(g->index_scores, g->index_q, g->index_weights, + if (!ds4_gpu_dsv41_indexer_scores_batch(g->index_scores, g->index_q, g->index_weights, g->index_cache[owner], n_comp, 1, pos, ratio)) return false; #endif } diff --git a/ds4_deepseek41_gpu.h b/ds4_deepseek41_gpu.h index 6d6a5572..ebc2eb3a 100644 --- a/ds4_deepseek41_gpu.h +++ b/ds4_deepseek41_gpu.h @@ -136,6 +136,14 @@ int ds4_gpu_dsv41_indexer_score_one_tensor(ds4_gpu_tensor *scores, uint32_t n_rows, uint32_t n_head, uint32_t head_dim, float scale, bool cache_f16); +/* The top-k of a row filtered by a block mask of more than 2048 blocks, from + * the scores of its 2048 candidate blocks only: the same ids in the same order + * (Metal only). Returns 0, leaving the full row to the caller, where unavailable. */ +int ds4_gpu_dsv41_indexer_candidates_tensor(ds4_gpu_tensor *selected, ds4_gpu_tensor *scores, + ds4_gpu_tensor *list, const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, const ds4_gpu_tensor *keys, + const ds4_gpu_tensor *block_mask, uint32_t n_comp, + float scale, uint32_t top_k); int ds4_gpu_dsv41_tensor_ops_available(void); /* Reuse exact BF16 views of FP4 queries/keys across score tiles. Invalid * casts retain the F32 arithmetic for the affected tile. */ diff --git a/ds4_metal.m b/ds4_metal.m index aa260667..c84a3e85 100644 --- a/ds4_metal.m +++ b/ds4_metal.m @@ -49536,6 +49536,79 @@ int ds4_gpu_dsv41_candidate_mask_batch(ds4_gpu_tensor *mask, return ds4_gpu_dsv4_topk_mask_tensor_impl(mask, selected, width, rows, 2048u, mask_stride); } +/* The wide scorer and radix select of a filtered row, over its candidates + * only: in key order the winners, their order and ties are the same, and the + * newest block, always a candidate, ends the list. Returns 0 to keep the + * full row where those kernels would not have run. */ +int ds4_gpu_dsv41_indexer_candidates_tensor(ds4_gpu_tensor *selected, ds4_gpu_tensor *scores, + ds4_gpu_tensor *list, const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, const ds4_gpu_tensor *keys, + const ds4_gpu_tensor *block_mask, uint32_t n_comp, + float scale, uint32_t top_k) { + static id list_pipeline, score_pipeline, ids_pipeline; + static int env = -1; + if (env < 0) env = !getenv("DS4_METAL_DISABLE_V41_INDEX_CANDIDATES"); + const uint32_t blocks = n_comp <= (1u << 23) ? (n_comp + 7u) / 8u : 0u; + const uint32_t count = 16384u - (blocks * 8u - n_comp); + if (!env || blocks <= 2048u || !top_k || top_k > 2048u || ds4_gpu_v41_index_score_wide_off() || + getenv("DS4_METAL_DISABLE_V41_TOPK_SELECT") || + (g_test_flags & DS4_GPU_TEST_INDEX_SCORE_DIRECT) != 0u || + !dsv41_tensor_has_floats(selected, top_k) || !dsv41_tensor_has_floats(scores, count) || + !dsv41_tensor_has_floats(list, 2048u) || !dsv41_tensor_has_floats(block_mask, blocks) || + !dsv41_tensor_has_floats(q, 4096u) || !dsv41_tensor_has_floats(weights, 32u) || + !dsv41_tensor_has_floats(keys, (uint64_t)n_comp * 128u)) return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + @autoreleasepool { + if (!list_pipeline) { + id hist = ds4_gpu_get_pipeline("kernel_topk_select_hist"); + list_pipeline = ds4_gpu_get_pipeline("kernel_dsv41_candidate_list"); + score_pipeline = ds4_gpu_get_pipeline("kernel_dsv41_indexer_score_candidates"); + ids_pipeline = ds4_gpu_get_pipeline("kernel_dsv41_candidate_ids"); + if (!hist || hist.maxTotalThreadsPerThreadgroup < 1024u || !score_pipeline || !ids_pipeline || + list_pipeline.maxTotalThreadsPerThreadgroup < 1024u) list_pipeline = nil; + if (!list_pipeline) env = 0; + } + if (!list_pipeline) return 0; + const ds4_gpu_glm_indexer_score_one_args args = { + .n_rows = count, .n_head = 32u, .head_dim = 128u, .cache_f16 = 0u, .scale = scale, + }; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return 0; + [enc setComputePipelineState:list_pipeline]; + [enc setBytes:&blocks length:sizeof(blocks) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(block_mask) offset:ds4_gpu_tensor_offset(block_mask) atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(list) offset:ds4_gpu_tensor_offset(list) atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) threadsPerThreadgroup:MTLSizeMake(1024, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:score_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + const ds4_gpu_tensor *inputs[] = {q, weights, keys, scores, list}; + for (NSUInteger i = 0; i < 5; i++) + [enc setBuffer:ds4_gpu_tensor_buffer(inputs[i]) + offset:ds4_gpu_tensor_offset(inputs[i]) atIndex:i + 1]; + [enc setThreadgroupMemoryLength:(4096u + 32u) * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((count + 255u) / 256u, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + if (!ds4_gpu_finish_command_buffer(cb, owned, "V4.1 candidate index scores") || + !ds4_gpu_indexer_topk_select(ds4_gpu_tensor_buffer(selected), ds4_gpu_tensor_offset(selected), + ds4_gpu_tensor_buffer(scores), ds4_gpu_tensor_offset(scores), count, 1, top_k)) return 0; + cb = ds4_gpu_command_buffer(&owned); + enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return 0; + [enc setComputePipelineState:ids_pipeline]; + [enc setBytes:&top_k length:sizeof(top_k) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(list) offset:ds4_gpu_tensor_offset(list) atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(selected) offset:ds4_gpu_tensor_offset(selected) atIndex:2]; + [enc dispatchThreads:MTLSizeMake(top_k, 1, 1) threadsPerThreadgroup:MTLSizeMake(MIN(top_k, 256u), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 candidate ids"); + } +} + int ds4_gpu_dsv41_hc_mean(uint32_t rows, uint32_t dim, uint32_t hc, const ds4_gpu_tensor *stream, ds4_gpu_tensor *out, uint32_t out_stride, uint32_t out_off) { diff --git a/metal/dsv41.metal b/metal/dsv41.metal index ff536175..c9f7682a 100644 --- a/metal/dsv41.metal +++ b/metal/dsv41.metal @@ -375,6 +375,55 @@ kernel void kernel_dsv41_candidate_filter( ? scores[offset] : -INFINITY; } +/* The 2048 candidate blocks of a 0/-inf mask in ascending order; one + * threadgroup of 1024 threads, four mask entries per thread and step. */ +kernel void kernel_dsv41_candidate_list( + constant uint &blocks, + device const float *mask, + device uint *list, + uint tid [[thread_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], + uint simd [[simdgroup_index_in_threadgroup]]) { + threadgroup uint part[33]; + uint base = 0u; + for (uint b0 = 0u; b0 < blocks; b0 += 4096u) { + const uint b = b0 + tid * 4u; + bool keep[4]; + uint count = 0u; + for (uint j = 0u; j < 4u; j++) { + keep[j] = b + j < blocks && mask[b + j] == 0.0f; + count += keep[j] ? 1u : 0u; + } + const uint before = simd_prefix_exclusive_sum(count); + if (lane == 31u) part[simd] = before + count; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (simd == 0u) { + const uint c = part[lane], c_before = simd_prefix_exclusive_sum(c); + part[lane] = c_before; + if (lane == 31u) part[32] = c_before + c; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + uint at = base + part[simd] + before; + for (uint j = 0u; j < 4u; j++) { + if (keep[j] && at < 2048u) list[at] = b + j; + at += keep[j] ? 1u : 0u; + } + base += part[32]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} + +/* Positions in the candidate list back to key ids, in place and in order. */ +kernel void kernel_dsv41_candidate_ids( + constant uint &count, + device const uint *list, + device int *selected, + uint i [[thread_position_in_grid]]) { + if (i >= count) return; + const uint p = (uint)selected[i]; + selected[i] = (int)(list[p >> 3] * 8u + (p & 7u)); +} + struct ds4_metal_args_dsv41_carry { uint width, rows, words, format, pack; }; diff --git a/metal/dsv4_misc.metal b/metal/dsv4_misc.metal index 1f554378..a5cee126 100644 --- a/metal/dsv4_misc.metal +++ b/metal/dsv4_misc.metal @@ -2064,6 +2064,56 @@ kernel void kernel_glm_indexer_score_one_wide( scores[row] = score; } +/* The wide kernel over V4.1 candidate keys: key i is position i % 8 of block + * list[i / 8], so the scores keep ascending key order. */ +kernel void kernel_dsv41_indexer_score_candidates( + constant ds4_metal_args_glm_indexer_score_one & args, + device const char *q, + device const float *weights, + device const char *indexer_key_cache, + device float *scores, + device const uint *list, + threadgroup float *shared [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + uint3 ntg_u [[threads_per_threadgroup]]) { + if (args.n_head != 32u || args.head_dim != 128u) return; + const uint ntg = ntg_u.x; + threadgroup float *qs = shared; /* [32][128] */ + threadgroup float *ws = shared + 4096u; /* [32] */ + device const float *qf = (device const float *)q; + for (uint i = tid; i < 4096u; i += ntg) qs[i] = qf[i]; + if (tid < 32u) ws[tid] = weights[tid]; + threadgroup_barrier(mem_flags::mem_threadgroup); + + const uint slot = tgpig.x * ntg + tid; + if (slot >= args.n_rows) return; + const uint row = list[slot >> 3] * 8u + (slot & 7u); + float acc[32]; + FOR_UNROLL (uint h = 0; h < 32u; h++) acc[h] = 0.0f; + const uint64_t base = (uint64_t)row * 128u; + for (uint d0 = 0; d0 < 128u; d0 += 16u) { + float4 k[4]; + if (args.cache_f16 != 0u) { + device const half4 *kh = (device const half4 *)indexer_key_cache + ((base + d0) >> 2); + FOR_UNROLL (uint i = 0; i < 4u; i++) k[i] = float4(kh[i]); + } else { + device const float4 *kf = (device const float4 *)indexer_key_cache + ((base + d0) >> 2); + FOR_UNROLL (uint i = 0; i < 4u; i++) k[i] = kf[i]; + } + FOR_UNROLL (uint h = 0; h < 32u; h++) { + threadgroup const float4 *q4 = (threadgroup const float4 *)(qs + h * 128u + d0); + acc[h] += dot(q4[0], k[0]) + dot(q4[1], k[1]) + dot(q4[2], k[2]) + dot(q4[3], k[3]); + } + } + float score = 0.0f; + FOR_UNROLL (uint h = 0; h < 32u; h++) { +#pragma clang fp contract(off) + score += max(acc[h] * args.scale, 0.0f) * ws[h]; + } + scores[slot] = score; +} + kernel void kernel_glm_indexer_scores_batch( constant ds4_metal_args_glm_indexer_scores_batch & args, device const char *q, diff --git a/tests/test_deepseek41_metal.c b/tests/test_deepseek41_metal.c index 533b1ebb..cd5ecdc3 100644 --- a/tests/test_deepseek41_metal.c +++ b/tests/test_deepseek41_metal.c @@ -1360,8 +1360,102 @@ static int check_index_score_wide(void) { fprintf(stderr, "V4.1 wide index score: double-precision oracle PASS\n"); return 1; } + +/* Candidate-only index rows against the wide row they replace (wide scores, + * candidate filter, radix top-512): the same ids in the same order, with + * realistic and wide FP4 block scales, partial newest blocks and ties. */ +static int check_index_candidates(void) { + enum { HEADS = 32, DIM = 128, TOP = 512, ROUNDS = 6 }; + const uint32_t widths[] = {16385, 16391, 20007, 65536, 131072, 131075, 262147}; + const uint32_t max_keys = 262147, max_blocks = (max_keys + 7) / 8; + ds4_gpu_tensor *qt = upload(NULL, (size_t)HEADS * DIM * 4), *wt = upload(NULL, HEADS * 4); + ds4_gpu_tensor *kt = upload(NULL, (size_t)max_keys * DIM * 4); + ds4_gpu_tensor *full = upload(NULL, (size_t)max_keys * 4), *plain = upload(NULL, (size_t)max_keys * 4); + ds4_gpu_tensor *source = upload(NULL, (size_t)max_keys * 4), *cand = upload(NULL, 16384 * 4); + ds4_gpu_tensor *list = upload(NULL, 2048 * 4), *blocks_t = upload(NULL, (size_t)max_blocks * 4); + ds4_gpu_tensor *picked = upload(NULL, 2048 * 4), *mask = upload(NULL, (size_t)max_blocks * 4); + ds4_gpu_tensor *old_sel = upload(NULL, TOP * 4), *new_sel = upload(NULL, TOP * 4); + CHECK(qt && wt && kt && full && plain && source && cand && list && blocks_t && picked && mask && + old_sel && new_sel); + float *q = ds4_gpu_tensor_contents(qt), *w = ds4_gpu_tensor_contents(wt); + float *k = ds4_gpu_tensor_contents(kt), *s = ds4_gpu_tensor_contents(source); + const float *p = ds4_gpu_tensor_contents(plain), *c = ds4_gpu_tensor_contents(cand); + const float *m = ds4_gpu_tensor_contents(mask); + const uint32_t *l = ds4_gpu_tensor_contents(list); + const int32_t *a = ds4_gpu_tensor_contents(old_sel), *b = ds4_gpu_tensor_contents(new_sel); + for (size_t wi = 0; wi < sizeof(widths) / sizeof(*widths); wi++) { + const uint32_t n = widths[wi], blocks = (n + 7) / 8, count = 16384 - (blocks * 8 - n); + double seconds[2] = {0, 0}; + uint32_t ties = 0; + for (int mode = 0; mode < 4; mode++) { + /* bit 0: block exponents over +-20 instead of +-2; bit 1: ties */ + const int range = mode & 1 ? 20 : 2; + for (size_t i = 0; i < (size_t)HEADS * DIM; i++) + q[i] = ldexpf(random_value(), (int)((i / 32) % (2 * range + 1)) - range); + for (size_t i = 0; i < (size_t)n * DIM; i++) + k[i] = ldexpf(random_value(), (int)((i / 32 + n) % (2 * range + 1)) - range); + for (size_t i = 0; i < HEADS; i++) w[i] = random_value() * 0.37139f; + for (uint32_t i = 0; i < n; i++) s[i] = (float)((i * 7919u + mode) % 1048573u); + CHECK(ds4_gpu_dsv41_quantize(qt, DIM, HEADS, DS4_V41_FP4_E8M0)); + CHECK(ds4_gpu_dsv41_quantize(kt, DIM, n, DS4_V41_FP4_E8M0)); + CHECK(ds4_gpu_synchronize()); + if (mode & 2) { + /* copy the best key over 1/16 of all keys: the 512th score ties */ + CHECK(ds4_gpu_dsv41_indexer_score_one_tensor(plain, qt, wt, kt, n, HEADS, DIM, + 1.0f / 64.0f, false) && ds4_gpu_synchronize()); + uint32_t best = 0; + for (uint32_t i = 1; i < n; i++) if (p[i] > p[best]) best = i; + for (uint32_t i = 0; i < n; i += 16) + memcpy(k + (size_t)(i ^ 5) % n * DIM, k + (size_t)best * DIM, DIM * 4); + } + CHECK(ds4_gpu_begin_commands()); + CHECK(ds4_gpu_dsv41_candidate_blocks(blocks_t, source, n, 1, n - 1, 1)); + CHECK(ds4_gpu_dsv41_indexer_topk_tensor(picked, blocks_t, blocks, 1, 2048)); + CHECK(ds4_gpu_dsv4_topk_mask_tensor(mask, picked, blocks, 1, 2048)); + CHECK(ds4_gpu_dsv41_indexer_score_one_tensor(plain, qt, wt, kt, n, HEADS, DIM, 1.0f / 64.0f, false)); + CHECK(ds4_gpu_end_commands()); + for (int round = 0; round < ROUNDS; round++) { + for (int path = 0; path < 2; path++) { + const double t0 = monotonic_seconds(); + CHECK(ds4_gpu_begin_commands()); + if (path) { + CHECK(ds4_gpu_dsv41_indexer_candidates_tensor(new_sel, cand, list, qt, wt, kt, mask, + n, 1.0f / 64.0f, TOP)); + } else { + CHECK(ds4_gpu_dsv41_indexer_score_one_tensor(full, qt, wt, kt, n, HEADS, DIM, + 1.0f / 64.0f, false)); + CHECK(ds4_gpu_dsv41_candidate_filter(full, mask, n, 1, n - 1, 1, 0)); + CHECK(ds4_gpu_dsv41_indexer_topk_tensor(old_sel, full, n, 1, TOP)); + } + CHECK(ds4_gpu_end_commands() && ds4_gpu_synchronize()); + if (round) seconds[path] += monotonic_seconds() - t0; + } + CHECK(!memcmp(a, b, TOP * sizeof(*a))); + } + /* the list is the mask's candidates in order; each score is the wide one */ + uint32_t at = 0; + for (uint32_t j = 0; j < blocks; j++) if (m[j] == 0.0f) CHECK(at < 2048 && l[at++] == j); + CHECK(at == 2048 && l[2047] == blocks - 1); + for (uint32_t i = 0; i < count; i++) { + const uint32_t key = l[i / 8] * 8 + i % 8; + CHECK(key < n && !memcmp(c + i, p + key, sizeof(*c))); + } + for (uint32_t i = 0; i < TOP; i++) ties += i && p[a[i]] == p[a[i - 1]]; + } + fprintf(stderr, "V4.1 candidate index rows keys=%u candidates=%u: ids, order and scores " + "identical (%u tied neighbours) us per row full=%.0f candidates=%.0f\n", + n, count, ties, seconds[0] / (4 * (ROUNDS - 1)) * 1e6, seconds[1] / (4 * (ROUNDS - 1)) * 1e6); + } + CHECK(!ds4_gpu_dsv41_indexer_candidates_tensor(new_sel, cand, list, qt, wt, kt, mask, + 16384, 1.0f / 64.0f, TOP)); + ds4_gpu_tensor *tensors[] = {qt, wt, kt, full, plain, source, cand, list, blocks_t, picked, mask, + old_sel, new_sel}; + for (size_t i = 0; i < sizeof(tensors) / sizeof(*tensors); i++) ds4_gpu_tensor_free(tensors[i]); + return 1; +} #else static int check_index_score_wide(void) { return 1; } +static int check_index_candidates(void) { return 1; } #endif static int check_compact_carry(void) { @@ -1585,6 +1679,11 @@ int main(int argc, char **argv) { ds4_gpu_cleanup(); return ok ? 0 : 1; } + if (argc == 2 && !strcmp(argv[1], "--index-candidates")) { + const int ok = ds4_gpu_init() && check_index_candidates(); + ds4_gpu_cleanup(); + return ok ? 0 : 1; + } if (argc == 2 && !strcmp(argv[1], "--index-topk")) { const int ok = ds4_gpu_init() && check_causal_topk(); ds4_gpu_cleanup(); @@ -1609,7 +1708,7 @@ int main(int argc, char **argv) { int ok = ds4_gpu_init() && check_router() && check_quantization() && check_engram() && check_rope_stride() && check_rope_quantize() && check_pool() && check_candidates() && check_candidates_batch(16449, 17, 16400) && check_candidates_batch(17017, 17, 17000) && check_sparse_gather() && check_indexer_batch() && - check_index_score_wide() && check_embedding() && check_index_projection() && check_general_topk() && + check_index_score_wide() && check_index_candidates() && check_embedding() && check_index_projection() && check_general_topk() && check_topk_select() && check_causal_topk() && check_indexer_all() && check_compact_carry() && check_attention_output(false) && check_tp_attention(); ds4_gpu_cleanup();