diff --git a/ds4.c b/ds4.c index f19ddcad..fb306f5a 100644 --- a/ds4.c +++ b/ds4.c @@ -40995,6 +40995,10 @@ typedef struct { float rows[2][DS4_ENGRAM_COLS * DS4_ENGRAM_DIM]; ds4_gpu_tensor *window[40]; ds4_gpu_tensor *compressed[4], *index_cache[4]; + /* Metal: index_cache's bf16 view for the decode scorer, valid unless a restore brought keys that are not + * bf16-exact (then the f32 cache is scored) */ + ds4_gpu_tensor *index_bf16[4], *index_bf16_flag; + bool index_bf16_ok; ds4_gpu_tensor *previous_kv[4], *previous_score[4]; ds4_gpu_tensor *verify_previous_kv[4], *verify_previous_score[4]; bool verify_prefix; @@ -41016,6 +41020,10 @@ static bool ds41_read_array(const ds4_model *m, const char *key, uint32_t type, return cursor_read(&c, out, bytes); } +#ifdef __APPLE__ +static bool ds41_index_bf16_on(void) { return !(ds4_gpu_v41_off() & DS4_GPU_V41_INDEX_BF16); } +#endif + static void ds41_graph_free(ds41_gpu_graph *g) { if (!g) return; ds41_dspark_free(g->dspark); @@ -41047,9 +41055,11 @@ static void ds41_graph_free(ds41_gpu_graph *g) { for (uint32_t i = 0; i < 4; i++) { ds4_gpu_tensor_free(g->compressed[i]); ds4_gpu_tensor_free(g->index_cache[i]); + ds4_gpu_tensor_free(g->index_bf16[i]); ds4_gpu_tensor_free(g->previous_kv[i]); ds4_gpu_tensor_free(g->previous_score[i]); } + ds4_gpu_tensor_free(g->index_bf16_flag); #define DS41_FREE(name, count) ds4_gpu_tensor_free(g->name); DS41_SCRATCH(DS41_FREE) #undef DS41_FREE @@ -41068,8 +41078,13 @@ static uint64_t ds41_graph_bytes(uint32_t ctx) { .prefill_alias = !getenv("DS4_METAL_DISABLE_V41_PREFILL_ALIAS"), .compact_carry = !getenv("DS4_METAL_DISABLE_V41_COMPACT_CARRY")}, *g = &shape; uint64_t floats = (uint64_t)40 * 128 * 512; +#ifdef __APPLE__ + const uint32_t view = ds41_index_bf16_on() ? 64u : 0u; /* the index keys' bf16 view */ +#else + const uint32_t view = 0; +#endif for (uint32_t i = 0; i < 4; i++) - floats += ((uint64_t)ctx / (i < 3 ? 2u : 1u) + 1u) * (512u + 128u) + 2u * 512u; + floats += ((uint64_t)ctx / (i < 3 ? 2u : 1u) + 1u) * (512u + 128u + view) + 2u * 512u; floats += (uint64_t)2 * 2 * DS4_N_EMBD * DS4_N_HC; #define DS41_COUNT(name, count) floats += (count); DS41_SCRATCH(DS41_COUNT) @@ -41109,6 +41124,7 @@ static ds4_context_memory ds41_graph_memory(uint32_t ctx) { static void ds41_graph_reset(ds41_gpu_graph *g) { g->pos = 0; g->valid = true; + g->index_bf16_ok = true; /* every row is rewritten through its producer */ ds4_engram_history_reset(&g->history); ds41_dspark_reset(g->dspark); /* Valid lengths, not zeroed storage, determine cache visibility. Each @@ -41218,7 +41234,14 @@ static DS4_MAYBE_UNUSED bool ds41_graph_alloc(ds41_gpu_graph *g, const ds4_model g->previous_kv[i] = ds4_gpu_tensor_alloc(512u * sizeof(float)); g->previous_score[i] = ds4_gpu_tensor_alloc(512u * sizeof(float)); if (!g->compressed[i] || !g->index_cache[i] || !g->previous_kv[i] || !g->previous_score[i]) goto fail; +#ifdef __APPLE__ + if (ds41_index_bf16_on() && !(g->index_bf16[i] = ds4_gpu_tensor_alloc(cap * 128u * sizeof(uint16_t)))) + goto fail; +#endif } +#ifdef __APPLE__ + if (ds41_index_bf16_on() && !(g->index_bf16_flag = ds4_gpu_tensor_alloc(sizeof(uint32_t)))) goto fail; +#endif #define DS41_ALLOC(name, count) \ if (!(g->name = (!strcmp(#name, "engram_prefetch") || !strcmp(#name, "image_text_mask") ? \ ds4_gpu_tensor_alloc_managed : ds4_gpu_tensor_alloc)( \ @@ -41607,6 +41630,39 @@ static bool ds41_short_sweep_off(void) { return metal_graph_env_flag("DS4_METAL_DISABLE_V41_SHORT_SWEEP", &cache); } +/* Rows [first, first + rows) of an owner's index keys into its bf16 view. + * Producers round every key to bf16, so the view holds the keys exactly. */ +static bool ds41_index_bf16_store(ds41_gpu_graph *g, uint32_t owner, uint32_t first, uint32_t rows) { +#ifdef __APPLE__ + return !g->index_bf16[owner] || + ds4_gpu_dsv41_index_keys_bf16(g->index_bf16[owner], g->index_cache[owner], first, rows, NULL); +#else + (void)g; (void)owner; (void)first; (void)rows; + return true; +#endif +} +/* A restore brings f32 keys from a save: rebuild the views and keep them + * only if every live key is bf16-exact; otherwise decode scores the f32 + * cache until the next reset. */ +static void ds41_index_bf16_restore(ds41_gpu_graph *g, uint32_t pos) { + g->index_bf16_ok = false; +#ifdef __APPLE__ + const uint32_t zero = 0; + if (!g->index_bf16_flag || !ds4_gpu_tensor_write(g->index_bf16_flag, 0, &zero, sizeof(zero))) return; + for (uint32_t i = 0; i < 4; i++) { + const uint32_t live = pos / (i < 3 ? 2u : 1u); + if (live && !ds4_gpu_dsv41_index_keys_bf16(g->index_bf16[i], g->index_cache[i], 0, live, g->index_bf16_flag)) + return; + } + uint32_t flag = 1; + if (!ds4_gpu_synchronize() || !ds4_gpu_tensor_read(g->index_bf16_flag, 0, &flag, sizeof(flag))) return; + g->index_bf16_ok = flag == 0; + if (flag) fprintf(stderr, "ds4: V4.1 restored index keys are not bf16-exact; decode scores the f32 keys\n"); +#else + (void)pos; +#endif +} + /* One block's HC input: the mixer projection, then split, weighted sum with * the previous block's pre weights and norm; each half fused where the fused * kernel reproduces the separate ones. */ @@ -41778,6 +41834,7 @@ static bool ds41_attention_publish(ds41_gpu_graph *g, const ds4_model *m, !ds41_norm(g->index_k, g->index_k, m, l->indexer_k_norm) || !ds41_rope_store(g->index_k, g->index_cache[owner], (uint64_t)(n_comp - 1u) * 128u * 4u, DS4_N_INDEXER_HEAD_DIM, il, rope_pos, DS4_V41_FP4_E8M0) || + !ds41_index_bf16_store(g, owner, n_comp - 1u, 1) || !ds41_rope_store(g->latent, g->compressed[owner], (uint64_t)(n_comp - 1u) * 512u * 4u, 512, il, rope_pos, DS4_V41_FP4_E4M3)) return false; } @@ -41854,8 +41911,8 @@ static bool ds41_attention_select_published(ds41_gpu_graph *g, const ds4_model * 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; + if (!ds4_gpu_dsv41_indexer_score_row(g->index_scores, g->index_q, g->index_weights, g->index_cache[owner], + g->index_bf16_ok ? g->index_bf16[owner] : NULL, n_comp)) return false; #else 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; @@ -42689,6 +42746,7 @@ static bool ds41_attention_publish_batch(ds41_gpu_graph *g, ds41_prefill_row *b, ds4_gpu_dsv41_quantize(b->index_k, DS4_N_INDEXER_HEAD_DIM, rows, DS4_V41_FP4_E8M0) && ds4_gpu_tensor_copy(g->index_cache[owner], (uint64_t)first * DS4_N_INDEXER_HEAD_DIM * sizeof(float), b->index_k, 0, (uint64_t)rows * DS4_N_INDEXER_HEAD_DIM * sizeof(float)) && + ds41_index_bf16_store(g, owner, first, rows) && ds4_gpu_dsv41_rope_stride(b->latent, DS4_N_HEAD_DIM, 1, rows, first * ratio, ratio, true, false) && ds4_gpu_dsv41_quantize(b->latent, DS4_N_HEAD_DIM, rows, DS4_V41_FP4_E4M3) && @@ -64597,6 +64655,7 @@ static int ds41_load_payload(ds4_session *s, FILE *fp, const uint32_t *h, for (uint32_t i = 0; i < 3 && i < pos; i++) g->history.tail[i] = (int32_t)g->token_map[tokens.v[pos - 1u - i]]; g->pos = pos; + ds41_index_bf16_restore(g, pos); ds4_tokens_copy(&s->checkpoint, &tokens); s->checkpoint_valid = true; } diff --git a/ds4_deepseek41_gpu.h b/ds4_deepseek41_gpu.h index d1801d85..f50b7afd 100644 --- a/ds4_deepseek41_gpu.h +++ b/ds4_deepseek41_gpu.h @@ -143,6 +143,13 @@ 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); +/* Rows [first, first + rows) of the index-key cache into its bf16 view; with flag, a key that is not bf16-exact + * sets it nonzero. score_row: one decode index row (32 heads x 128, scale 1/64) over the cache, reading view, when + * given, where the wide scorer runs (the same bits). */ +int ds4_gpu_dsv41_index_keys_bf16(ds4_gpu_tensor *view, const ds4_gpu_tensor *cache, + uint32_t first, uint32_t rows, ds4_gpu_tensor *flag); +int ds4_gpu_dsv41_indexer_score_row(ds4_gpu_tensor *scores, const ds4_gpu_tensor *q, const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *cache, const ds4_gpu_tensor *view, uint32_t rows); /* 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. */ diff --git a/ds4_gpu.h b/ds4_gpu.h index 8e71d8b6..44fded89 100644 --- a/ds4_gpu.h +++ b/ds4_gpu.h @@ -55,6 +55,7 @@ enum { DS4_GPU_V41_MODEL_TRACKED = 1u << 3, /* opt-out: DS4_METAL_MODEL_TRACKED */ DS4_GPU_V41_ROPE_TABLE = 1u << 4, DS4_GPU_V41_SPLIT_REDUCE = 1u << 5, + DS4_GPU_V41_INDEX_BF16 = 1u << 6, DS4_GPU_V41_PLAIN_READERS = 1u << 29, /* test only: the narrow readers' own kernels */ }; uint32_t ds4_gpu_v41_off(void); diff --git a/ds4_metal.m b/ds4_metal.m index 18fb7cae..a6982817 100644 --- a/ds4_metal.m +++ b/ds4_metal.m @@ -447,6 +447,7 @@ uint32_t ds4_gpu_v41_off(void) { "DS4_METAL_MODEL_TRACKED", "DS4_METAL_DISABLE_V41_ROPE_TABLE", "DS4_METAL_DISABLE_V41_SPLIT_REDUCE", + "DS4_METAL_DISABLE_V41_INDEX_BF16", }; for (uint32_t i = 0; i < sizeof(names) / sizeof(*names); i++) { const char *v = getenv(names[i]); @@ -587,6 +588,8 @@ static id g_glm_indexer_rope_tail_pipeline; static id g_glm_indexer_score_one_pipeline; static id g_glm_indexer_score_one_direct_pipeline; static id g_glm_indexer_score_one_wide_pipeline; +static id g_dsv41_indexer_score_bf16_pipeline; +static id g_dsv41_index_keys_bf16_pipeline; static id g_glm_indexer_scores_batch_pipeline; static id g_glm_indexer_scores_tiled_pipeline; static id g_glm_indexer_scores_tiled_f32_pipeline; @@ -9092,6 +9095,10 @@ int ds4_gpu_init(void) { ds4_gpu_get_pipeline("kernel_glm_indexer_score_one_direct"); g_glm_indexer_score_one_wide_pipeline = ds4_gpu_get_pipeline("kernel_glm_indexer_score_one_wide"); + g_dsv41_indexer_score_bf16_pipeline = + ds4_gpu_get_pipeline("kernel_dsv41_indexer_score_bf16"); + g_dsv41_index_keys_bf16_pipeline = + ds4_gpu_get_pipeline("kernel_dsv41_index_keys_bf16"); g_glm_indexer_scores_batch_pipeline = ds4_gpu_get_pipeline("kernel_glm_indexer_scores_batch"); g_glm_indexer_scores_tiled_pipeline = @@ -9214,6 +9221,8 @@ int ds4_gpu_init(void) { !g_glm_indexer_score_one_pipeline || !g_glm_indexer_score_one_direct_pipeline || !g_glm_indexer_score_one_wide_pipeline || + !g_dsv41_indexer_score_bf16_pipeline || + !g_dsv41_index_keys_bf16_pipeline || !g_glm_indexer_scores_batch_pipeline || !g_glm_indexer_scores_tiled_pipeline || !g_glm_indexer_scores_tiled_f32_pipeline || @@ -37238,6 +37247,75 @@ int ds4_gpu_dsv41_indexer_score_one_tensor( n_rows, n_head, head_dim, scale, cache_f16, true); } +int ds4_gpu_dsv41_index_keys_bf16(ds4_gpu_tensor *view, const ds4_gpu_tensor *cache, + uint32_t first, uint32_t rows, ds4_gpu_tensor *flag) { + if (!view || !cache || !rows) return rows == 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + const uint64_t end = (uint64_t)first + rows; + if (ds4_gpu_tensor_bytes(cache) < end * 128u * 4u || ds4_gpu_tensor_bytes(view) < end * 128u * 2u || + (flag && ds4_gpu_tensor_bytes(flag) < 4u) || end * 32u > UINT32_MAX) return 0; + @autoreleasepool { + id pipeline = ds4_gpu_hot_pipeline( + g_dsv41_index_keys_bf16_pipeline, "kernel_dsv41_index_keys_bf16"); + if (!pipeline) return 0; + const struct { uint32_t n_vec, check; } args = {rows * 32u, flag != NULL}; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(cache) offset:ds4_gpu_tensor_offset(cache) + (uint64_t)first * 128u * 4u atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(view) offset:ds4_gpu_tensor_offset(view) + (uint64_t)first * 128u * 2u atIndex:2]; + const ds4_gpu_tensor *f = flag ? flag : view; /* unread without check */ + [enc setBuffer:ds4_gpu_tensor_buffer(f) offset:ds4_gpu_tensor_offset(f) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)rows * 32u + 255u) / 256u, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 bf16 index keys"); + } +} + +int ds4_gpu_dsv41_indexer_score_row(ds4_gpu_tensor *scores, const ds4_gpu_tensor *q, const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *cache, const ds4_gpu_tensor *view, uint32_t rows) { + /* the view where the wide scorer runs (rows >= 8192, as ds4_gpu_glm_indexer_score_one_impl decides) */ + const bool wide = (g_test_flags & DS4_GPU_TEST_INDEX_SCORE_WIDE) != 0u || + ((g_test_flags & DS4_GPU_TEST_INDEX_SCORE_DIRECT) == 0u && rows >= 8192u && !ds4_gpu_v41_index_score_wide_off()); + if (!view || !wide || !ds4_gpu_v41_on(DS4_GPU_V41_INDEX_BF16) || !scores || !q || !weights || + ds4_gpu_tensor_bytes(view) < (uint64_t)rows * 128u * 2u) + return ds4_gpu_dsv41_indexer_score_one_tensor(scores, q, weights, cache, rows, 32, 128, 1.0f / 64.0f, false); + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!rows || ds4_gpu_tensor_bytes(scores) < (uint64_t)rows * 4u || ds4_gpu_tensor_bytes(q) < 4096u * 4u || + ds4_gpu_tensor_bytes(weights) < 32u * 4u) return 0; + @autoreleasepool { + /* below 65536 rows, S lanes share each key (four below 10240 rows, else two); above, two keys per thread, + * 512 threads below 98304 rows and 128 from there (measured) */ + const char *name = rows < 10240u ? "kernel_dsv41_indexer_score_rows_split4_bf16" : + rows < 65536u ? "kernel_dsv41_indexer_score_rows_split2_bf16" : "kernel_dsv41_indexer_score_bf16"; + const NSUInteger tg = rows < 65536u ? 256 : rows < 98304u ? 512 : 128; + const NSUInteger keys_per_tg = rows < 10240u ? tg / 2u : rows < 65536u ? tg : 2u * tg; + id pipeline = rows < 65536u ? ds4_gpu_get_pipeline(name) : + ds4_gpu_hot_pipeline(g_dsv41_indexer_score_bf16_pipeline, name); + if (!pipeline) return 0; + DS4_TEST_NOTE(name); + ds4_gpu_glm_indexer_score_one_args args = {rows, 32u, 128u, 0u, 1.0f / 64.0f}; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + const ds4_gpu_tensor *t[] = {q, weights, view, scores}; + for (unsigned i = 0; i < 4; i++) + [enc setBuffer:ds4_gpu_tensor_buffer(t[i]) offset:ds4_gpu_tensor_offset(t[i]) atIndex:i + 1]; + [enc setThreadgroupMemoryLength:(4096u + 32u) * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)rows + keys_per_tg - 1u) / keys_per_tg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(tg, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 bf16 index score"); + } +} + static int ds4_gpu_glm_indexer_scores_batch_grouped_tensor( ds4_gpu_tensor *scores, const ds4_gpu_tensor *q, diff --git a/metal/dsv4_misc.metal b/metal/dsv4_misc.metal index db522a26..9b8e7e0f 100644 --- a/metal/dsv4_misc.metal +++ b/metal/dsv4_misc.metal @@ -273,6 +273,11 @@ struct ds4_metal_args_glm_indexer_score_one { float scale; }; +struct ds4_metal_args_dsv41_index_keys_bf16 { + uint32_t n_vec; + uint32_t check; +}; + struct ds4_metal_args_glm_indexer_scores_batch { uint32_t n_rows; uint32_t n_tokens; @@ -2064,6 +2069,67 @@ kernel void kernel_glm_indexer_score_one_wide( scores[row] = score; } +/* Index keys as bf16 (FP4 values on a power-of-two scale are bf16-exact); with check, a key that is not raises *flag. */ +kernel void kernel_dsv41_index_keys_bf16( + constant ds4_metal_args_dsv41_index_keys_bf16 & args, + device const uint4 *src, + device ushort4 *dst, + device atomic_uint *flag, + uint gid [[thread_position_in_grid]]) { + if (gid >= args.n_vec) return; + const uint4 b = src[gid]; + dst[gid] = ushort4(b >> 16); + if (args.check != 0u && any((b & 0xffffu) != 0u)) + atomic_fetch_or_explicit(flag, 1u, memory_order_relaxed); +} + +/* The wide scorer over the bf16 keys, two keys per thread: each dot and the head sum (0..31) in its order, its bits. */ +kernel void kernel_dsv41_indexer_score_bf16( + constant ds4_metal_args_glm_indexer_score_one & args, + device const float *qf, + device const float *weights, + device const ushort4 *keys, + device float *scores, + threadgroup float *shared [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + uint3 ntg_u [[threads_per_threadgroup]]) { + const uint ntg = ntg_u.x; + threadgroup float *qs = shared; /* [32][128] */ + threadgroup float *ws = shared + 4096u; /* [32] */ + 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); + if (args.n_head != 32u || args.head_dim != 128u) return; + + const uint first = tgpig.x * ntg * 2u + tid; + /* the tail's second key rereads the last row and is not written */ + const uint row[2] = {min(first, args.n_rows - 1u), min(first + ntg, args.n_rows - 1u)}; + float score[2] = {0.0f, 0.0f}; + for (uint g = 0; g < 32u; g += 16u) { + float acc[2][16]; + FOR_UNROLL (uint r = 0; r < 2u; r++) FOR_UNROLL (uint h = 0; h < 16u; h++) acc[r][h] = 0.0f; + for (uint d0 = 0; d0 < 128u; d0 += 16u) { + float4 k[2][4]; + FOR_UNROLL (uint r = 0; r < 2u; r++) { + device const ushort4 *kb = keys + (((uint64_t)row[r] * 128u + d0) >> 2); + FOR_UNROLL (uint i = 0; i < 4u; i++) k[r][i] = as_type(uint4(kb[i]) << 16); + } + FOR_UNROLL (uint h = 0; h < 16u; h++) { + threadgroup const float4 *q4 = (threadgroup const float4 *)(qs + (g + h) * 128u + d0); + FOR_UNROLL (uint r = 0; r < 2u; r++) + acc[r][h] += dot(q4[0], k[r][0]) + dot(q4[1], k[r][1]) + dot(q4[2], k[r][2]) + dot(q4[3], k[r][3]); + } + } + FOR_UNROLL (uint r = 0; r < 2u; r++) FOR_UNROLL (uint h = 0; h < 16u; h++) { +#pragma clang fp contract(off) + score[r] += max(acc[r][h] * args.scale, 0.0f) * ws[g + h]; + } + } + if (first < args.n_rows) scores[first] = score[0]; + if (first + ntg < args.n_rows) scores[first + ntg] = score[1]; +} + /* 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( @@ -2114,6 +2180,67 @@ kernel void kernel_dsv41_indexer_score_candidates( scores[slot] = score; } +/* The same with S lanes per key, the head sum passed lane to lane in order (2T / S keys per T threads). */ +template +kernel void kernel_dsv41_indexer_score_rows_split_bf16( + constant ds4_metal_args_glm_indexer_score_one & args, + device const float *qf, + device const float *weights, + device const ushort *keys, + device float *scores, + threadgroup float *shared [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + uint3 ntg_u [[threads_per_threadgroup]]) { + const uint ntg = ntg_u.x; + threadgroup float *qs = shared; /* [32][128] */ + threadgroup float *ws = shared + 4096u; /* [32] */ + 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); + if (args.n_head != 32u || args.head_dim != 128u) return; + constexpr uint HS = 32u / S; + const uint part = tid % S, lanes = ntg / S; + uint slot[2], row[2]; + FOR_UNROLL (uint r = 0; r < 2u; r++) { + slot[r] = tgpig.x * lanes * 2u + r * lanes + tid / S; + row[r] = min(slot[r], args.n_rows - 1u); + } + float acc[2][HS]; + FOR_UNROLL (uint r = 0; r < 2u; r++) FOR_UNROLL (uint i = 0; i < HS; i++) acc[r][i] = 0.0f; + for (uint d0 = 0; d0 < 128u; d0 += 16u) { + float4 k[2][4]; + FOR_UNROLL (uint r = 0; r < 2u; r++) { + device const ushort4 *kb = (device const ushort4 *)keys + (((uint64_t)row[r] * 128u + d0) >> 2); + FOR_UNROLL (uint i = 0; i < 4u; i++) k[r][i] = as_type(uint4(kb[i]) << 16); + } + FOR_UNROLL (uint i = 0; i < HS; i++) { + threadgroup const float4 *q4 = (threadgroup const float4 *)(qs + (part * HS + i) * 128u + d0); + FOR_UNROLL (uint r = 0; r < 2u; r++) + acc[r][i] += dot(q4[0], k[r][0]) + dot(q4[1], k[r][1]) + dot(q4[2], k[r][2]) + dot(q4[3], k[r][3]); + } + } + FOR_UNROLL (uint r = 0; r < 2u; r++) { + float score = 0.0f; + FOR_UNROLL (uint p = 0; p < S; p++) { + if (p) score = simd_shuffle_up(score, (ushort)1); + if (part == p) { + FOR_UNROLL (uint i = 0; i < HS; i++) { +#pragma clang fp contract(off) + score += max(acc[r][i] * args.scale, 0.0f) * ws[p * HS + i]; + } + } + } + if (part == S - 1u && slot[r] < args.n_rows) scores[slot[r]] = score; + } +} +#define DSV41_ROWS_SPLIT(NAME, S) template [[host_name(NAME)]] \ + kernel void kernel_dsv41_indexer_score_rows_split_bf16( \ + constant ds4_metal_args_glm_indexer_score_one &, device const float *, device const float *, device const ushort *, \ + device float *, threadgroup float *, uint3, ushort, uint3); +DSV41_ROWS_SPLIT("kernel_dsv41_indexer_score_rows_split2_bf16", 2) +DSV41_ROWS_SPLIT("kernel_dsv41_indexer_score_rows_split4_bf16", 4) + 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 703cbbf8..c5a3e539 100644 --- a/tests/test_deepseek41_metal.c +++ b/tests/test_deepseek41_metal.c @@ -1780,6 +1780,84 @@ static int check_index_candidates(void) { return 1; } +/* The bf16 key view and its scorers against the wide scorer on the f32 cache: the same bits for every key. Keys pass + * the FP4_E8M0 quantizer with block scales spread over 41 exponents and over 12 (where summation order shows in the + * bits); the view converts in two pieces at an odd boundary. The decode row entry reads the view where the wide + * scorer runs and scores the f32 keys without it. */ +static int check_index_score_bf16(void) { + enum { HEADS = 32, DIM = 128 }; + const uint32_t counts[] = {8200, 65535, 65536, 98303, 131077}; + const uint32_t max_keys = 131077; + const char *kernels[] = {"kernel_dsv41_indexer_score_rows_split4_bf16", "kernel_dsv41_indexer_score_rows_split2_bf16", + "kernel_dsv41_indexer_score_bf16"}; + ds4_gpu_tensor *qt = upload(NULL, (size_t)HEADS * DIM * sizeof(float)); + ds4_gpu_tensor *wt = upload(NULL, (size_t)HEADS * sizeof(float)); + ds4_gpu_tensor *kt = upload(NULL, (size_t)max_keys * DIM * sizeof(float)); + ds4_gpu_tensor *vt = upload(NULL, (size_t)max_keys * DIM * sizeof(uint16_t)); + ds4_gpu_tensor *sw = upload(NULL, ((size_t)max_keys + 1) * sizeof(float)); + ds4_gpu_tensor *sb = upload(NULL, ((size_t)max_keys + 1) * sizeof(float)); + ds4_gpu_tensor *ft = upload(NULL, sizeof(uint32_t)); + CHECK(qt && wt && kt && vt && sw && sb && ft); + float *q = ds4_gpu_tensor_contents(qt), *w = ds4_gpu_tensor_contents(wt); + float *k = ds4_gpu_tensor_contents(kt); + float *wide = ds4_gpu_tensor_contents(sw), *fast = ds4_gpu_tensor_contents(sb); + uint32_t *flag = ds4_gpu_tensor_contents(ft); + for (int spread = 0; spread < 2; spread++) { + for (size_t i = 0; i < (size_t)HEADS * DIM; i++) + q[i] = spread == 0 ? ldexpf(random_value(), (int)((i / 32) % 41) - 20) : + ldexpf(random_value(), (int)((i / 32 * 7) % 13) - 6); + for (size_t i = 0; i < HEADS; i++) w[i] = bf16(random_value()); + for (size_t i = 0; i < (size_t)max_keys * DIM; i++) + k[i] = spread == 0 ? ldexpf(random_value(), (int)((i / 32) % 41) - 20) : + ldexpf(random_value(), (int)((i / 32 * 5 + i / DIM) % 13) - 6); + CHECK(ds4_gpu_dsv41_quantize(qt, DIM, HEADS, DS4_V41_FP4_E8M0)); + CHECK(ds4_gpu_dsv41_quantize(kt, DIM, max_keys, DS4_V41_FP4_E8M0)); + *flag = 0; + CHECK(ds4_gpu_dsv41_index_keys_bf16(vt, kt, 0, 77777, ft)); + CHECK(ds4_gpu_dsv41_index_keys_bf16(vt, kt, 77777, max_keys - 77777, ft)); + CHECK(ds4_gpu_synchronize() && *flag == 0); + for (size_t ci = 0; ci < sizeof(counts) / sizeof(*counts); ci++) { + const uint32_t keys = counts[ci]; + const char *kernel = kernels[keys < 10240u ? 0 : keys < 65536u ? 1 : 2]; + wide[keys] = fast[keys] = 12345; + ds4_gpu_test_set_flags(DS4_GPU_TEST_INDEX_SCORE_WIDE); + CHECK(ds4_gpu_glm_indexer_score_one_tensor(sw, qt, wt, kt, keys, HEADS, DIM, 1.0f / 64.0f, false)); + ds4_gpu_test_set_flags(DS4_GPU_TEST_COUNT_KERNELS); + for (int view = 1; view >= 0; view--) { + const uint64_t runs = ds4_gpu_test_kernel_runs(kernel); + memset(fast, 0xff, (size_t)keys * sizeof(float)); + CHECK(ds4_gpu_dsv41_indexer_score_row(sb, qt, wt, kt, view ? vt : NULL, keys) && ds4_gpu_synchronize()); + CHECK(ds4_gpu_test_kernel_runs(kernel) == runs + (uint64_t)view); + CHECK(!memcmp(wide, fast, (size_t)keys * sizeof(float)) && fast[keys] == 12345); + } + uint32_t positive = 0, rounded = 0; + for (uint32_t j = 0; j < keys; j++) positive += wide[j] > 0; + /* sensitivity: keys whose f32 score is not the exact sum, so order shows */ + for (uint32_t j = 0; ci == 0 && j < 2048; j++) { + double expected = 0; + for (uint32_t h = 0; h < HEADS; h++) { + double dot = 0; + for (uint32_t d = 0; d < DIM; d++) + dot += (double)q[(size_t)h * DIM + d] * k[(size_t)j * DIM + d]; + expected += fmax(dot / 64.0, 0) * w[h]; + } + rounded += (double)wide[j] != expected; + } + CHECK(wide[keys] == 12345 && positive > keys / 16 && (ci != 0 || rounded > 0)); + } + } + /* a key that is not bf16-exact raises the flag */ + uint32_t bits; + memcpy(&bits, &k[(size_t)123456 * DIM + 77], 4); + bits |= 1u; + memcpy(&k[(size_t)123456 * DIM + 77], &bits, 4); + CHECK(ds4_gpu_dsv41_index_keys_bf16(vt, kt, 123456, 1, ft) && ds4_gpu_synchronize() && *flag != 0); + ds4_gpu_tensor_free(ft); ds4_gpu_tensor_free(sb); ds4_gpu_tensor_free(sw); + ds4_gpu_tensor_free(vt); ds4_gpu_tensor_free(kt); ds4_gpu_tensor_free(wt); ds4_gpu_tensor_free(qt); + fprintf(stderr, "V4.1 bf16 index score (split4, split2, full) == wide scorer bits PASS\n"); + return 1; +} + /* The FP8 copy reader (ds4_gpu_matmul_e4m3_bf16_tensor), with its blend kernels and with the narrow readers, gives * the BF16-rounded F16 projection's bits at the Engram (6144 -> 256, two-row blend) and indexer q_b (1280 -> 4096, * block-owner blend) shapes: power-of-two 32x32 block scales 2^-15..2^-8 (every weight exact in F16), normal and @@ -1861,6 +1939,7 @@ static int check_e4m3_copies(void) { #else static int check_index_score_wide(void) { return 1; } static int check_index_candidates(void) { return 1; } +static int check_index_score_bf16(void) { return 1; } static int check_e4m3_copies(void) { return 1; } #endif @@ -2115,6 +2194,11 @@ int main(int argc, char **argv) { ds4_gpu_cleanup(); return ok ? 0 : 1; } + if (argc == 2 && !strcmp(argv[1], "--index-score-bf16")) { + const int ok = ds4_gpu_init() && check_index_score_bf16(); + ds4_gpu_cleanup(); + return ok ? 0 : 1; + } if (argc == 2 && !strcmp(argv[1], "--e4m3-copies")) { const int ok = ds4_gpu_init() && check_e4m3_copies(); ds4_gpu_cleanup(); @@ -2151,6 +2235,7 @@ int main(int argc, char **argv) { ok = ok && check_e4m3_copies(); ok = ok && check_rope_tab(); ok = ok && check_split_collapse(); + ok = ok && check_index_score_bf16(); ds4_gpu_cleanup(); return ok ? 0 : 1; }