diff --git a/ds4.c b/ds4.c index 49be4d39..8eb2d6cb 100644 --- a/ds4.c +++ b/ds4.c @@ -40927,7 +40927,7 @@ static bool ds41_attention_select_published(ds41_gpu_graph *g, const ds4_model * !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) || #ifdef __APPLE__ - !ds4_gpu_glm_indexer_score_one_tensor(g->index_scores, g->index_q, g->index_weights, + !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, diff --git a/ds4_deepseek41_gpu.h b/ds4_deepseek41_gpu.h index 6ca540ad..635d40bb 100644 --- a/ds4_deepseek41_gpu.h +++ b/ds4_deepseek41_gpu.h @@ -127,6 +127,15 @@ int ds4_gpu_dsv41_indexer_scores_batch(ds4_gpu_tensor *scores, const ds4_gpu_tensor *keys, uint32_t source_rows, uint32_t rows, uint32_t start, uint32_t ratio); +/* One-token index scores of FP4 queries/keys: the GLM scorer, plus its wide + * kernel for long rows (Metal only). */ +int ds4_gpu_dsv41_indexer_score_one_tensor(ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *indexer_key_cache, + uint32_t n_rows, uint32_t n_head, + uint32_t head_dim, float scale, + bool cache_f16); 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 824540f3..47d0c58e 100644 --- a/ds4_metal.m +++ b/ds4_metal.m @@ -36769,7 +36769,7 @@ static int ds4_gpu_v41_index_score_wide_off(void) { return off; } -int ds4_gpu_glm_indexer_score_one_tensor( +static int ds4_gpu_glm_indexer_score_one_impl( ds4_gpu_tensor *scores, const ds4_gpu_tensor *q, const ds4_gpu_tensor *weights, @@ -36778,7 +36778,8 @@ int ds4_gpu_glm_indexer_score_one_tensor( uint32_t n_head, uint32_t head_dim, float scale, - bool cache_f16) { + bool cache_f16, + bool fp4_inputs) { if (!g_initialized && !ds4_gpu_init()) return 0; if (!scores || !q || !weights || !indexer_key_cache || n_rows == 0 || n_head == 0 || head_dim == 0 || @@ -36814,9 +36815,10 @@ int ds4_gpu_glm_indexer_score_one_tensor( }; if (n_head == 32u && head_dim == 128u) { - /* long rows stream one key per thread; short ones keep the per-key kernel */ + /* long rows stream one key per thread when FP4 query and keys make + * its summation order exact; the rest keep the per-key kernel */ const bool wide = (g_test_flags & DS4_GPU_TEST_INDEX_SCORE_WIDE) != 0u || - ((g_test_flags & DS4_GPU_TEST_INDEX_SCORE_DIRECT) == 0u && + ((g_test_flags & DS4_GPU_TEST_INDEX_SCORE_DIRECT) == 0u && fp4_inputs && n_rows >= 8192u && !ds4_gpu_v41_index_score_wide_off()); id direct_pipeline = wide ? ds4_gpu_hot_pipeline(g_glm_indexer_score_one_wide_pipeline, @@ -36880,6 +36882,34 @@ int ds4_gpu_glm_indexer_score_one_tensor( return 1; } +int ds4_gpu_glm_indexer_score_one_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *indexer_key_cache, + uint32_t n_rows, + uint32_t n_head, + uint32_t head_dim, + float scale, + bool cache_f16) { + return ds4_gpu_glm_indexer_score_one_impl(scores, q, weights, indexer_key_cache, + n_rows, n_head, head_dim, scale, cache_f16, false); +} + +int ds4_gpu_dsv41_indexer_score_one_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *indexer_key_cache, + uint32_t n_rows, + uint32_t n_head, + uint32_t head_dim, + float scale, + bool cache_f16) { + return ds4_gpu_glm_indexer_score_one_impl(scores, q, weights, indexer_key_cache, + n_rows, n_head, head_dim, scale, cache_f16, true); +} + 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 513fb7e1..1f554378 100644 --- a/metal/dsv4_misc.metal +++ b/metal/dsv4_misc.metal @@ -2057,7 +2057,10 @@ kernel void kernel_glm_indexer_score_one_wide( } } float score = 0.0f; - FOR_UNROLL (uint h = 0; h < 32u; h++) score += max(acc[h] * args.scale, 0.0f) * ws[h]; + FOR_UNROLL (uint h = 0; h < 32u; h++) { +#pragma clang fp contract(off) + score += max(acc[h] * args.scale, 0.0f) * ws[h]; + } scores[row] = score; }