From 663db6a56483613a15208717d7eaf86064470892 Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Sun, 30 Aug 2026 12:00:48 +0000 Subject: [PATCH 01/16] perf(ds4): extend sparse prefill scheduling to 8k --- server/src/deepseek4/deepseek4_graph.cpp | 104 ++++++++++++---------- server/src/deepseek4/deepseek4_internal.h | 4 +- 2 files changed, 58 insertions(+), 50 deletions(-) diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 8ac1c06a9..f76be6bc0 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -2007,12 +2007,12 @@ static ggml_tensor * build_mla_attention( // [n_kv,n_query] F16; the explicit path broadcasts the same values over // heads in F32. ggml_tensor * score_mask = nullptr; - const bool exact_two_band = + const bool exact_numerical_bands = attention_impl == DeepSeek4AttentionImpl::DenseFlash && causal_batch && n_tokens > DS4_NUMERICAL_PREFILL_BAND && - n_tokens <= 2 * DS4_NUMERICAL_PREFILL_BAND; - if (!exact_two_band) { + n_tokens <= DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS; + if (!exact_numerical_bands) { if (masked_kv && n_tokens > 1) { score_mask = ggml_reshape_2d(ctx, cached_inputs->attn_row_mask, n_attn, n_tokens); @@ -2107,21 +2107,12 @@ static ggml_tensor * build_mla_attention( const bool use_flash = attention_impl != DeepSeek4AttentionImpl::Explicit && (n_tokens > 1 || indexer_topk != nullptr); if (use_flash) { - if (exact_two_band) { - // A larger scheduling band must retain the numerical topology of - // two 2K requests. Prefix queries use the first band's F32 raw KV; - // suffix queries see its final SWA tail after the same F16 cache - // round-trip. HC, projections and MoE still run once over the full - // token batch, avoiding a second expert-weight sweep. - const int first_count = DS4_NUMERICAL_PREFILL_BAND; - const int second_count = n_tokens - first_count; - const int first_comp = ratio > 0 - ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, - kv_start + first_count - 1) - : 0; - const int second_comp = n_comp_live; - const int second_prior_count = std::min(first_count, w.n_swa); - + if (exact_numerical_bands) { + // A larger scheduling batch retains the numerical topology of + // sequential 2K requests. Each later band sees the previous + // band's final SWA tail after the same F16 cache round-trip. HC, + // projections and MoE still run once over the full token batch, + // avoiding another expert-weight sweep. auto view_kv = [&](int first, int count) { return ggml_view_2d( ctx, kv, head_dim, count, kv->nb[1], @@ -2203,36 +2194,53 @@ static ggml_tensor * build_mla_attention( (size_t) first * q_fa->nb[1]); }; - ggml_tensor * first_raw = ds4_cast_if_needed( - ctx, view_kv(0, first_count), GGML_TYPE_F32); - if (prior_rows_scratch) { - first_raw = ggml_concat( - ctx, prior_rows_scratch, first_raw, 1); - } - ggml_tensor * first_kv = append_comp(first_raw, first_comp); - ggml_tensor * first_mask = make_band_mask( - kv_start, first_count, n_prior_rows, first_comp); - ggml_tensor * first_context = make_flash( - view_q(0, first_count), first_kv, first_mask, - n_prior_rows + first_count, kv_start); - - ggml_tensor * rounded_prior = ggml_cast( - ctx, view_kv(first_count - second_prior_count, - second_prior_count), - GGML_TYPE_F16); - rounded_prior = ggml_cast(ctx, rounded_prior, GGML_TYPE_F32); - ggml_tensor * second_raw = ggml_concat( - ctx, rounded_prior, view_kv(first_count, second_count), 1); - ggml_tensor * second_kv = append_comp(second_raw, second_comp); - ggml_tensor * second_mask = make_band_mask( - kv_start + first_count, second_count, - second_prior_count, second_comp); - ggml_tensor * second_context = make_flash( - view_q(first_count, second_count), second_kv, second_mask, - second_prior_count + second_count, - kv_start + first_count); - - context = ggml_concat(ctx, first_context, second_context, 2); + for (int band_start = 0; band_start < n_tokens; + band_start += DS4_NUMERICAL_PREFILL_BAND) { + const int band_count = std::min( + DS4_NUMERICAL_PREFILL_BAND, n_tokens - band_start); + const int band_pos = kv_start + band_start; + const int band_prior_count = band_start == 0 + ? n_prior_rows + : std::min(band_start, w.n_swa); + const int band_comp_count = ratio > 0 + ? ds4_comp_rows_used( + lc.comp_kv, lc.n_comp, ratio, + band_pos + band_count - 1) + : 0; + + ggml_tensor * band_raw = nullptr; + if (band_start == 0) { + band_raw = ds4_cast_if_needed( + ctx, view_kv(0, band_count), GGML_TYPE_F32); + if (prior_rows_scratch) { + band_raw = ggml_concat( + ctx, prior_rows_scratch, band_raw, 1); + } + } else { + ggml_tensor * rounded_prior = ggml_cast( + ctx, + view_kv(band_start - band_prior_count, + band_prior_count), + GGML_TYPE_F16); + rounded_prior = ggml_cast( + ctx, rounded_prior, GGML_TYPE_F32); + band_raw = ggml_concat( + ctx, rounded_prior, + view_kv(band_start, band_count), 1); + } + + ggml_tensor * band_kv = append_comp( + band_raw, band_comp_count); + ggml_tensor * band_mask = make_band_mask( + band_pos, band_count, band_prior_count, + band_comp_count); + ggml_tensor * band_context = make_flash( + view_q(band_start, band_count), band_kv, band_mask, + band_prior_count + band_count, band_pos); + context = context + ? ggml_concat(ctx, context, band_context, 2) + : band_context; + } inverse_rope_fused = true; } else { // ggml FA convention: Q[D,T,H], K/V[D,K,Hkv]. DS4 MLA has one shared diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index 66d80e417..1be9f7d60 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -28,10 +28,10 @@ namespace dflash::common { -// Layer-major prefill may schedule two 2K numerical bands while preserving +// Layer-major prefill may schedule four 2K numerical bands while preserving // the raw-cache rounding boundary between them. inline constexpr int DS4_NUMERICAL_PREFILL_BAND = 2048; -inline constexpr int DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS = 4096; +inline constexpr int DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS = 8192; // Normal verification stays within one ratio-4 compressor window. Q5 is an // explicit opt-in whose fused graph models a second boundary. inline constexpr int DS4_CONSERVATIVE_VERIFY_MAX_TOKENS = 4; From 1709ba5407a158739352dae821691efb32c491c1 Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Sun, 30 Aug 2026 18:04:41 +0000 Subject: [PATCH 02/16] perf(ds4): pass long-prefill top-k directly --- server/src/deepseek4/deepseek4_graph.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index f76be6bc0..82183e521 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -2082,8 +2082,9 @@ static ggml_tensor * build_mla_attention( score_mask = ggml_reshape_2d(ctx, cmask, n_attn, n_tokens); } } - const bool direct_indexer_topk = indexer_topk && - ds4_env_flag("DFLASH_DS4_DIRECT_INDEXER_TOPK"); + // Long sparse prefill already has the authoritative selected rows. Pass + // them through directly instead of materializing and rescanning a mask. + const bool direct_indexer_topk = indexer_topk && n_tokens > w.n_swa; if (indexer_topk) { if (!score_mask) { score_mask = ggml_new_tensor_2d( From 82973a86e0047a771a3ab1b5229a8c70fb108aa8 Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Mon, 31 Aug 2026 12:34:04 +0200 Subject: [PATCH 03/16] perf(ds4): keep sparse selected KV in F16 --- server/src/deepseek4/deepseek4_graph.cpp | 23 ++++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 82183e521..f7f9b1877 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -1709,6 +1709,7 @@ static ggml_tensor * build_mla_attention( ggml_tensor * old_rows_scratch_f16 = nullptr; int n_old_rows = 0; ggml_tensor * prior_rows_scratch = nullptr; + ggml_tensor * prior_rows_scratch_f16 = nullptr; int n_prior_rows = 0; const bool fused_causal = cached_inputs && cached_inputs->attn_row_mask && n_tokens > 1; if (fused_causal) { @@ -1765,6 +1766,7 @@ static ggml_tensor * build_mla_attention( prior_rows_scratch = ggml_cont(ctx, prior_rows_scratch); } ggml_build_forward_expand(gf, prior_rows_scratch); + prior_rows_scratch_f16 = prior_rows_scratch; prior_rows_scratch = ds4_cast_if_needed( ctx, prior_rows_scratch, GGML_TYPE_F32); } @@ -1890,6 +1892,9 @@ static ggml_tensor * build_mla_attention( index_visibility_mask, i32_array_inputs); } + const bool f16_sparse_prefill = + attention_impl == DeepSeek4AttentionImpl::SparseFlash && + indexer_topk && n_tokens > w.n_swa; // Stable path reads the full physical ring (masking not-yet-written slots) // and a padded compressed-row span; the plain path reads only valid rows. const int n_raw = masked_kv ? w.n_swa @@ -1920,9 +1925,13 @@ static ggml_tensor * build_mla_attention( ctx, raw_kv_source, head_dim, w.n_swa, raw_kv_source->nb[1], 0); kv_attn = ds4_cast_if_needed(ctx, ring, GGML_TYPE_F32); } else if (layer_major_batch) { - ggml_tensor * current = ds4_cast_if_needed(ctx, kv, GGML_TYPE_F32); - kv_attn = prior_rows_scratch - ? ggml_concat(ctx, prior_rows_scratch, current, 1) + ggml_tensor * current = ds4_cast_if_needed( + ctx, kv, + f16_sparse_prefill ? GGML_TYPE_F16 : GGML_TYPE_F32); + ggml_tensor * prior = f16_sparse_prefill + ? prior_rows_scratch_f16 : prior_rows_scratch; + kv_attn = prior + ? ggml_concat(ctx, prior, current, 1) : current; } else if (n_tokens == 1) { ggml_tensor * cur_kv = ds4_cast_if_needed(ctx, kv, GGML_TYPE_F32); @@ -1993,7 +2002,9 @@ static ggml_tensor * build_mla_attention( ggml_tensor * comp = ggml_view_2d( ctx, comp_kv_source, head_dim, n_comp_attn, comp_kv_source->nb[1], 0); - comp = ds4_cast_if_needed(ctx, comp, GGML_TYPE_F32); + comp = ds4_cast_if_needed( + ctx, comp, + f16_sparse_prefill ? GGML_TYPE_F16 : GGML_TYPE_F32); kv_attn = ggml_concat(ctx, kv_attn, comp, 1); } if (old_rows_scratch) { @@ -2249,7 +2260,9 @@ static ggml_tensor * build_mla_attention( // The DS4 D=512 kernel consumes Q strides directly, avoiding a full // [D,H,T] -> [D,T,H] materialization for every layer. ggml_tensor * q_fa = ggml_permute(ctx, q, 0, 2, 1, 3); - ggml_tensor * kv_fa = ds4_cast_if_needed(ctx, kv_attn, GGML_TYPE_F32); + ggml_tensor * kv_fa = f16_sparse_prefill + ? kv_attn + : ds4_cast_if_needed(ctx, kv_attn, GGML_TYPE_F32); ggml_tensor * k_fa = ggml_reshape_3d(ctx, kv_fa, head_dim, n_attn, 1); ggml_tensor * v_fa = k_fa; ggml_tensor * mask_fa = score_mask From 5a055d77def69960729759a0c5ef3a3ee9e81062 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sun, 30 Aug 2026 23:19:41 +0200 Subject: [PATCH 04/16] perf(hip): add opt-in streaming MLA top-k --- .../llama.cpp/ggml/src/ggml-cuda/fattn.cu | 240 ++++++++++++++++++ 1 file changed, 240 insertions(+) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu index 78215b4d3..4195a210f 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu @@ -1555,6 +1555,202 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( } } +// Streaming indexed MLA for long prefill. The compact grouped kernel above +// stores every score and then reloads V. Once the trained indexer has reduced +// the compressed history to a bounded top-k set, that extra traffic is no +// longer necessary: stage one latent row in LDS, share it across the heads in +// a block, and update online-softmax state while the row is resident. Sixteen +// wave32 heads amortize each LDS load best on gfx1151. The contract is +// backend-generic (D512 MQA with K == V and direct indexed rows); model policy +// remains in the graph/backend layer. +template +__global__ static void ds4_flash_attn_d512_streaming_topk_kernel( + float * dst, + const float * q, + size_t q_stride_token, + size_t q_stride_head, + const KV * kv, + const Mask * mask, + const float * sinks, + int n_tokens, + int n_heads, + int n_kv, + float scale, + const int * visibility_bounds, + const int * indexed_rows, + const int * indexed_counts, + int indexed_capacity, + ds4_inverse_rope_params inverse_rope, + const float * inverse_rope_coefficients, + const float * forward_rope_coefficients) { + constexpr int D = 512; + constexpr int WAVE = 32; + constexpr int N_THREADS = HEADS_PER_BLOCK * WAVE; + constexpr int VALUES_PER_LANE = D / WAVE; + static_assert(HEADS_PER_BLOCK == 16); + static_assert(KEYS_PER_STAGE == 16); + static_assert(N_THREADS == 512); + + const int token = (int) blockIdx.x; + const int head_begin = (int) blockIdx.y * HEADS_PER_BLOCK; + const int tid = (int) threadIdx.x; + const int wave = tid / WAVE; + const int lane = tid & (WAVE - 1); + const int head = head_begin + wave; + + // The launch gate makes the grid exact. Keeping every thread live is + // required because each stage has workgroup-wide barriers. + if (token >= n_tokens || head_begin + HEADS_PER_BLOCK > n_heads) return; + + __shared__ KV staged_kv[KEYS_PER_STAGE * D]; + __shared__ int staged_rows[KEYS_PER_STAGE]; + __shared__ float staged_masks[KEYS_PER_STAGE]; + + const int * token_visibility = visibility_bounds + (size_t) token * 4; + const int raw_first = token_visibility[0]; + const int raw_last = token_visibility[1]; + const int raw_count = raw_last >= raw_first + ? raw_last - raw_first + 1 : 0; + const int indexed_count = indexed_counts[token]; + const int total_rows = raw_count + indexed_count; + const int * token_rows = indexed_rows + + (size_t) token * indexed_capacity; + const Mask * token_mask = mask + (size_t) token * n_kv; + + const float * qh = q + (size_t) token * q_stride_token + + (size_t) head * q_stride_head; + float q_values[VALUES_PER_LANE]; + float accum[VALUES_PER_LANE] = {}; +#pragma unroll + for (int i = 0; i < VALUES_PER_LANE; ++i) { + const int dim = lane + i * WAVE; + float qv = qh[dim]; + if (inverse_rope.forward_q_enabled && dim >= D - 64) { + const int pair = (dim - (D - 64)) / 2; + const float x0 = qh[D - 64 + 2 * pair + 0]; + const float x1 = qh[D - 64 + 2 * pair + 1]; + const size_t coefficient_index = + ((size_t) token * 32 + (size_t) pair) * 2; + float y0; + float y1; + ds4_apply_inverse_rope_pair( + x0, x1, + forward_rope_coefficients[coefficient_index + 0], + forward_rope_coefficients[coefficient_index + 1], + y0, y1); + qv = (dim & 1) == 0 ? y0 : y1; + } + q_values[i] = qv; + } + + float row_max = -3.402823466e38f; + float row_sum = 0.0f; + for (int row_base = 0; row_base < total_rows; + row_base += KEYS_PER_STAGE) { + if (tid < KEYS_PER_STAGE) { + const int selected = row_base + tid; + int row = -1; + if (selected < raw_count) { + row = raw_first + selected; + } else if (selected < total_rows) { + row = token_rows[selected - raw_count]; + } + staged_rows[tid] = row; + staged_masks[tid] = row >= 0 && row < n_kv + ? ds4_fa_load(token_mask + row) + : -3.402823466e38f; + } + __syncthreads(); + + for (int index = tid; index < KEYS_PER_STAGE * D; + index += N_THREADS) { + const int slot = index / D; + const int dim = index - slot * D; + const int row = staged_rows[slot]; + staged_kv[index] = row >= 0 && row < n_kv + ? kv[(size_t) row * D + dim] : KV{}; + } + __syncthreads(); + +#pragma unroll + for (int slot = 0; slot < KEYS_PER_STAGE; ++slot) { + const int selected = row_base + slot; + if (selected >= total_rows) continue; + const float mask_value = staged_masks[slot]; + if (mask_value <= -1.0e20f) continue; + + float partial = 0.0f; +#pragma unroll + for (int i = 0; i < VALUES_PER_LANE; ++i) { + const int dim = lane + i * WAVE; + partial += q_values[i] * + ds4_fa_load(staged_kv + slot * D + dim); + } + partial = warp_reduce_sum(partial); + // XOR reduction can associate operands differently in each lane. + // Broadcast lane zero so every output dimension advances one + // identical softmax state. + partial = __shfl_sync(0xffffffffu, partial, 0, WAVE); + + const float score = partial * scale + mask_value; + const float next_max = fmaxf(row_max, score); + const float old_scale = row_sum == 0.0f + ? 0.0f : expf(row_max - next_max); + const float value_scale = expf(score - next_max); + row_sum = row_sum * old_scale + value_scale; + row_max = next_max; +#pragma unroll + for (int i = 0; i < VALUES_PER_LANE; ++i) { + const int dim = lane + i * WAVE; + const float value = ds4_fa_load( + staged_kv + slot * D + dim); + accum[i] = accum[i] * old_scale + value_scale * value; + } + } + __syncthreads(); + } + + if (sinks) { + const float sink = sinks[head]; + const float next_max = fmaxf(row_max, sink); + const float old_scale = row_sum == 0.0f + ? 0.0f : expf(row_max - next_max); + row_sum = row_sum * old_scale + expf(sink - next_max); +#pragma unroll + for (int i = 0; i < VALUES_PER_LANE; ++i) { + accum[i] *= old_scale; + } + } + + const float inv_sum = row_sum == 0.0f ? 0.0f : 1.0f / row_sum; + float * out = dst + + ((size_t) token * (size_t) n_heads + (size_t) head) * D; +#pragma unroll + for (int i = 0; i < VALUES_PER_LANE; ++i) { + const int dim = lane + i * WAVE; + float value = accum[i] * inv_sum; + if (inverse_rope.enabled && dim >= D - 64) { + const float partner = __shfl_xor_sync( + 0xffffffffu, value, 1, WAVE); + const float x0 = (dim & 1) == 0 ? value : partner; + const float x1 = (dim & 1) == 0 ? partner : value; + const int pair = (dim - (D - 64)) / 2; + const size_t coefficient_index = + ((size_t) token * 32 + (size_t) pair) * 2; + float y0; + float y1; + ds4_apply_inverse_rope_pair( + x0, x1, + inverse_rope_coefficients[coefficient_index + 0], + inverse_rope_coefficients[coefficient_index + 1], + y0, y1); + value = (dim & 1) == 0 ? y0 : y1; + } + out[dim] = value; + } +} + template static bool ds4_launch_flash_attn_d512_grouped( ggml_tensor * dst, @@ -2047,6 +2243,50 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( n_tokens, n_kv, raw_rows); } CUDA_CHECK(cudaGetLastError()); + // Vulkan shares each selected latent row across eight subgroup64 + // heads. HIP tuning selected sixteen wave32 heads for the same K == V + // reuse. Keep the native HIP path opt-in until a model-backed A/B + // qualifies its online-softmax association. The shape gate is + // expressed in terms of the D512 indexed-attention + // contract so another model can reuse the kernel without DS4 policy + // leaking into it. + const char * streaming_topk_env = + getenv("GGML_CUDA_MLA_STREAM_TOPK"); + if (!streaming_topk_env) { + streaming_topk_env = getenv("GGML_DS4_FA_STREAM_TOPK"); + } + const bool streaming_topk_enabled = streaming_topk_env && + streaming_topk_env[0] != '\0' && + strcmp(streaming_topk_env, "0") != 0; + constexpr int streaming_min_tokens = 64; + const int active_row_upper_bound = raw_window + indexed_capacity; + const int device_warp_size = + ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; + if (streaming_topk_enabled && indexed_mask && indexer_topk && + kv_f16 && mask->type == GGML_TYPE_F16 && + K->data == V->data && n_heads % 16 == 0 && + device_warp_size == 32 && n_tokens >= streaming_min_tokens && + active_row_upper_bound > 0 && + n_kv >= 3 * active_row_upper_bound) { + constexpr int streaming_heads = 16; + const dim3 streaming_grid( + (unsigned) n_tokens, + (unsigned) (n_heads / streaming_heads), 1); + ds4_flash_attn_d512_streaming_topk_kernel + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const half *) K->data, + (const half *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, + visibility_bounds, indexed_rows, indexed_counts, + indexed_capacity, inverse_rope, + inverse_rope_coefficients, + forward_rope_coefficients); + CUDA_CHECK(cudaGetLastError()); + return true; + } if (indexed_mask) { return ds4_launch_flash_attn_d512_grouped_compact< group4, true, 4>( From c7dffc0d612054712b15baa85733843ef5a0ef58 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 31 Aug 2026 02:11:45 +0200 Subject: [PATCH 05/16] perf(hip): vectorize sparse attention staging --- .../llama.cpp/ggml/src/ggml-cuda/fattn.cu | 92 ++++++++++++++----- 1 file changed, 69 insertions(+), 23 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu index 4195a210f..dd273caf0 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu @@ -7,6 +7,8 @@ #include "fattn-chunked.cuh" #include "fattn.cuh" +#include + #if defined(GGML_USE_HIP) __device__ static float ds4_fa_block_sum(float v) { @@ -1564,7 +1566,7 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( // backend-generic (D512 MQA with K == V and direct indexed rows); model policy // remains in the graph/backend layer. template + int KEYS_PER_STAGE = 16, bool STAGE_F32 = false> __global__ static void ds4_flash_attn_d512_streaming_topk_kernel( float * dst, const float * q, @@ -1603,7 +1605,8 @@ __global__ static void ds4_flash_attn_d512_streaming_topk_kernel( // required because each stage has workgroup-wide barriers. if (token >= n_tokens || head_begin + HEADS_PER_BLOCK > n_heads) return; - __shared__ KV staged_kv[KEYS_PER_STAGE * D]; + using stage_type = std::conditional_t; + __shared__ __align__(16) stage_type staged_kv[KEYS_PER_STAGE * D]; __shared__ int staged_rows[KEYS_PER_STAGE]; __shared__ float staged_masks[KEYS_PER_STAGE]; @@ -1663,13 +1666,35 @@ __global__ static void ds4_flash_attn_d512_streaming_topk_kernel( } __syncthreads(); - for (int index = tid; index < KEYS_PER_STAGE * D; - index += N_THREADS) { - const int slot = index / D; - const int dim = index - slot * D; - const int row = staged_rows[slot]; - staged_kv[index] = row >= 0 && row < n_kv - ? kv[(size_t) row * D + dim] : KV{}; + // Convert aligned half2 pairs once while loading them. Every head in + // the block then consumes the same F32 LDS values without repeating + // half conversion in both the score and value passes. + if constexpr (STAGE_F32 && std::is_same_v) { + constexpr int PAIRS_PER_ROW = D / 2; + for (int pair_index = tid; + pair_index < KEYS_PER_STAGE * PAIRS_PER_ROW; + pair_index += N_THREADS) { + const int slot = pair_index / PAIRS_PER_ROW; + const int pair = pair_index - slot * PAIRS_PER_ROW; + const int row = staged_rows[slot]; + float2 unpacked = make_float2(0.0f, 0.0f); + if (row >= 0 && row < n_kv) { + ds4_fa_load_pair( + kv + (size_t) row * D + 2 * pair, + unpacked.x, unpacked.y); + } + reinterpret_cast(staged_kv)[pair_index] = unpacked; + } + } else { + for (int index = tid; index < KEYS_PER_STAGE * D; + index += N_THREADS) { + const int slot = index / D; + const int dim = index - slot * D; + const int row = staged_rows[slot]; + staged_kv[index] = row >= 0 && row < n_kv + ? static_cast(kv[(size_t) row * D + dim]) + : stage_type{}; + } } __syncthreads(); @@ -1685,7 +1710,8 @@ __global__ static void ds4_flash_attn_d512_streaming_topk_kernel( for (int i = 0; i < VALUES_PER_LANE; ++i) { const int dim = lane + i * WAVE; partial += q_values[i] * - ds4_fa_load(staged_kv + slot * D + dim); + ds4_fa_load( + staged_kv + slot * D + dim); } partial = warp_reduce_sum(partial); // XOR reduction can associate operands differently in each lane. @@ -1703,7 +1729,7 @@ __global__ static void ds4_flash_attn_d512_streaming_topk_kernel( #pragma unroll for (int i = 0; i < VALUES_PER_LANE; ++i) { const int dim = lane + i * WAVE; - const float value = ds4_fa_load( + const float value = ds4_fa_load( staged_kv + slot * D + dim); accum[i] = accum[i] * old_scale + value_scale * value; } @@ -2268,22 +2294,42 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( device_warp_size == 32 && n_tokens >= streaming_min_tokens && active_row_upper_bound > 0 && n_kv >= 3 * active_row_upper_bound) { + const char * f32_stage_env = + getenv("GGML_CUDA_MLA_STREAM_F32_STAGE"); + const bool f32_stage = f32_stage_env && f32_stage_env[0] != '\0' && + strcmp(f32_stage_env, "0") != 0; constexpr int streaming_heads = 16; const dim3 streaming_grid( (unsigned) n_tokens, (unsigned) (n_heads / streaming_heads), 1); - ds4_flash_attn_d512_streaming_topk_kernel - <<>>( - (float *) dst->data, (const float *) Q->data, - q_stride_token, q_stride_head, - (const half *) K->data, - (const half *) mask->data, - sinks ? (const float *) sinks->data : nullptr, - n_tokens, n_heads, n_kv, scale, - visibility_bounds, indexed_rows, indexed_counts, - indexed_capacity, inverse_rope, - inverse_rope_coefficients, - forward_rope_coefficients); + if (f32_stage) { + ds4_flash_attn_d512_streaming_topk_kernel< + half, half, streaming_heads, 16, true> + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const half *) K->data, + (const half *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, + visibility_bounds, indexed_rows, indexed_counts, + indexed_capacity, inverse_rope, + inverse_rope_coefficients, + forward_rope_coefficients); + } else { + ds4_flash_attn_d512_streaming_topk_kernel + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const half *) K->data, + (const half *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, + visibility_bounds, indexed_rows, indexed_counts, + indexed_capacity, inverse_rope, + inverse_rope_coefficients, + forward_rope_coefficients); + } CUDA_CHECK(cudaGetLastError()); return true; } From caa632b72281727d110df53166a19615fd5907a1 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 31 Aug 2026 03:06:39 +0200 Subject: [PATCH 06/16] perf(hip): accelerate streaming MLA softmax --- .../llama.cpp/ggml/src/ggml-cuda/fattn.cu | 40 ++++++++++++++++--- 1 file changed, 34 insertions(+), 6 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu index dd273caf0..84110a23b 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu @@ -1566,7 +1566,8 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( // backend-generic (D512 MQA with K == V and direct indexed rows); model policy // remains in the graph/backend layer. template + int KEYS_PER_STAGE = 16, bool STAGE_F32 = false, + bool FAST_EXP = false> __global__ static void ds4_flash_attn_d512_streaming_topk_kernel( float * dst, const float * q, @@ -1722,8 +1723,13 @@ __global__ static void ds4_flash_attn_d512_streaming_topk_kernel( const float score = partial * scale + mask_value; const float next_max = fmaxf(row_max, score); const float old_scale = row_sum == 0.0f - ? 0.0f : expf(row_max - next_max); - const float value_scale = expf(score - next_max); + ? 0.0f + : (FAST_EXP + ? __expf(row_max - next_max) + : expf(row_max - next_max)); + const float value_scale = FAST_EXP + ? __expf(score - next_max) + : expf(score - next_max); row_sum = row_sum * old_scale + value_scale; row_max = next_max; #pragma unroll @@ -1741,8 +1747,12 @@ __global__ static void ds4_flash_attn_d512_streaming_topk_kernel( const float sink = sinks[head]; const float next_max = fmaxf(row_max, sink); const float old_scale = row_sum == 0.0f - ? 0.0f : expf(row_max - next_max); - row_sum = row_sum * old_scale + expf(sink - next_max); + ? 0.0f + : (FAST_EXP + ? __expf(row_max - next_max) + : expf(row_max - next_max)); + row_sum = row_sum * old_scale + + (FAST_EXP ? __expf(sink - next_max) : expf(sink - next_max)); #pragma unroll for (int i = 0; i < VALUES_PER_LANE; ++i) { accum[i] *= old_scale; @@ -2298,11 +2308,29 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( getenv("GGML_CUDA_MLA_STREAM_F32_STAGE"); const bool f32_stage = f32_stage_env && f32_stage_env[0] != '\0' && strcmp(f32_stage_env, "0") != 0; + const char * fast_exp_env = + getenv("GGML_CUDA_MLA_STREAM_FAST_EXP"); + const bool fast_exp = fast_exp_env && fast_exp_env[0] != '\0' && + strcmp(fast_exp_env, "0") != 0; constexpr int streaming_heads = 16; const dim3 streaming_grid( (unsigned) n_tokens, (unsigned) (n_heads / streaming_heads), 1); - if (f32_stage) { + if (f32_stage && fast_exp) { + ds4_flash_attn_d512_streaming_topk_kernel< + half, half, streaming_heads, 16, true, true> + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const half *) K->data, + (const half *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, + visibility_bounds, indexed_rows, indexed_counts, + indexed_capacity, inverse_rope, + inverse_rope_coefficients, + forward_rope_coefficients); + } else if (f32_stage) { ds4_flash_attn_d512_streaming_topk_kernel< half, half, streaming_heads, 16, true> <<>>( From 5d1ac6aebef5f714c6da899419ee53863a192d75 Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Mon, 31 Aug 2026 12:59:53 +0200 Subject: [PATCH 07/16] perf(hip): qualify streaming selected KV --- .../llama.cpp/ggml/src/ggml-cuda/fattn.cu | 127 ++++++--------- server/tests/test_deepseek4_unit.cpp | 150 ++++++++++++++++++ 2 files changed, 196 insertions(+), 81 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu index 84110a23b..3b90a44f9 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu @@ -1722,14 +1722,20 @@ __global__ static void ds4_flash_attn_d512_streaming_topk_kernel( const float score = partial * scale + mask_value; const float next_max = fmaxf(row_max, score); - const float old_scale = row_sum == 0.0f - ? 0.0f - : (FAST_EXP - ? __expf(row_max - next_max) - : expf(row_max - next_max)); - const float value_scale = FAST_EXP - ? __expf(score - next_max) - : expf(score - next_max); + float old_scale = 0.0f; + float value_scale = 0.0f; + if (lane == 0) { + old_scale = row_sum == 0.0f + ? 0.0f + : (FAST_EXP + ? __expf(row_max - next_max) + : expf(row_max - next_max)); + value_scale = FAST_EXP + ? __expf(score - next_max) + : expf(score - next_max); + } + old_scale = __shfl_sync(0xffffffffu, old_scale, 0, WAVE); + value_scale = __shfl_sync(0xffffffffu, value_scale, 0, WAVE); row_sum = row_sum * old_scale + value_scale; row_max = next_max; #pragma unroll @@ -1746,13 +1752,21 @@ __global__ static void ds4_flash_attn_d512_streaming_topk_kernel( if (sinks) { const float sink = sinks[head]; const float next_max = fmaxf(row_max, sink); - const float old_scale = row_sum == 0.0f - ? 0.0f - : (FAST_EXP - ? __expf(row_max - next_max) - : expf(row_max - next_max)); - row_sum = row_sum * old_scale + - (FAST_EXP ? __expf(sink - next_max) : expf(sink - next_max)); + float old_scale = 0.0f; + float sink_scale = 0.0f; + if (lane == 0) { + old_scale = row_sum == 0.0f + ? 0.0f + : (FAST_EXP + ? __expf(row_max - next_max) + : expf(row_max - next_max)); + sink_scale = FAST_EXP + ? __expf(sink - next_max) + : expf(sink - next_max); + } + old_scale = __shfl_sync(0xffffffffu, old_scale, 0, WAVE); + sink_scale = __shfl_sync(0xffffffffu, sink_scale, 0, WAVE); + row_sum = row_sum * old_scale + sink_scale; #pragma unroll for (int i = 0; i < VALUES_PER_LANE; ++i) { accum[i] *= old_scale; @@ -2279,85 +2293,36 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( n_tokens, n_kv, raw_rows); } CUDA_CHECK(cudaGetLastError()); - // Vulkan shares each selected latent row across eight subgroup64 - // heads. HIP tuning selected sixteen wave32 heads for the same K == V - // reuse. Keep the native HIP path opt-in until a model-backed A/B - // qualifies its online-softmax association. The shape gate is - // expressed in terms of the D512 indexed-attention - // contract so another model can reuse the kernel without DS4 policy - // leaking into it. - const char * streaming_topk_env = - getenv("GGML_CUDA_MLA_STREAM_TOPK"); - if (!streaming_topk_env) { - streaming_topk_env = getenv("GGML_DS4_FA_STREAM_TOPK"); - } - const bool streaming_topk_enabled = streaming_topk_env && - streaming_topk_env[0] != '\0' && - strcmp(streaming_topk_env, "0") != 0; + // Long sparse F16 K == V attention shares each selected latent row + // across sixteen wave32 heads. The shape gate is expressed in terms + // of the reusable D512 indexed-attention contract. constexpr int streaming_min_tokens = 64; const int active_row_upper_bound = raw_window + indexed_capacity; const int device_warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; - if (streaming_topk_enabled && indexed_mask && indexer_topk && + if (indexed_mask && indexer_topk && kv_f16 && mask->type == GGML_TYPE_F16 && K->data == V->data && n_heads % 16 == 0 && device_warp_size == 32 && n_tokens >= streaming_min_tokens && active_row_upper_bound > 0 && n_kv >= 3 * active_row_upper_bound) { - const char * f32_stage_env = - getenv("GGML_CUDA_MLA_STREAM_F32_STAGE"); - const bool f32_stage = f32_stage_env && f32_stage_env[0] != '\0' && - strcmp(f32_stage_env, "0") != 0; - const char * fast_exp_env = - getenv("GGML_CUDA_MLA_STREAM_FAST_EXP"); - const bool fast_exp = fast_exp_env && fast_exp_env[0] != '\0' && - strcmp(fast_exp_env, "0") != 0; constexpr int streaming_heads = 16; const dim3 streaming_grid( (unsigned) n_tokens, (unsigned) (n_heads / streaming_heads), 1); - if (f32_stage && fast_exp) { - ds4_flash_attn_d512_streaming_topk_kernel< - half, half, streaming_heads, 16, true, true> - <<>>( - (float *) dst->data, (const float *) Q->data, - q_stride_token, q_stride_head, - (const half *) K->data, - (const half *) mask->data, - sinks ? (const float *) sinks->data : nullptr, - n_tokens, n_heads, n_kv, scale, - visibility_bounds, indexed_rows, indexed_counts, - indexed_capacity, inverse_rope, - inverse_rope_coefficients, - forward_rope_coefficients); - } else if (f32_stage) { - ds4_flash_attn_d512_streaming_topk_kernel< - half, half, streaming_heads, 16, true> - <<>>( - (float *) dst->data, (const float *) Q->data, - q_stride_token, q_stride_head, - (const half *) K->data, - (const half *) mask->data, - sinks ? (const float *) sinks->data : nullptr, - n_tokens, n_heads, n_kv, scale, - visibility_bounds, indexed_rows, indexed_counts, - indexed_capacity, inverse_rope, - inverse_rope_coefficients, - forward_rope_coefficients); - } else { - ds4_flash_attn_d512_streaming_topk_kernel - <<>>( - (float *) dst->data, (const float *) Q->data, - q_stride_token, q_stride_head, - (const half *) K->data, - (const half *) mask->data, - sinks ? (const float *) sinks->data : nullptr, - n_tokens, n_heads, n_kv, scale, - visibility_bounds, indexed_rows, indexed_counts, - indexed_capacity, inverse_rope, - inverse_rope_coefficients, - forward_rope_coefficients); - } + ds4_flash_attn_d512_streaming_topk_kernel< + half, half, streaming_heads, 16, true, true> + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const half *) K->data, + (const half *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, + visibility_bounds, indexed_rows, indexed_counts, + indexed_capacity, inverse_rope, + inverse_rope_coefficients, + forward_rope_coefficients); CUDA_CHECK(cudaGetLastError()); return true; } diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index 0d3d8d361..55d7a0440 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -2831,6 +2831,155 @@ static void test_ds4_flash_attention_parallel_index_scan_gpu() { std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); } +static void test_ds4_flash_attention_streaming_topk_gpu() { + std::fprintf(stderr, + " test_ds4_flash_attention_streaming_topk_gpu ..."); +#if !defined(GGML_USE_HIP) + std::fprintf(stderr, " skipped (HIP-only candidate)\n"); + return; +#endif + ggml_backend_t backend = ggml_backend_cuda_init(0); + if (!backend) { + std::fprintf(stderr, " skipped (no GPU backend)\n"); + return; + } + + constexpr int head_dim = 512; + constexpr int n_heads = 64; + constexpr int n_tokens = 64; + constexpr int raw_rows = 128; + constexpr int raw_window = 128; + constexpr int selected_rows = 512; + constexpr int n_comp_rows = 1920; + constexpr int n_kv = raw_rows + n_comp_rows; + + ggml_context * ctx = make_test_context(4u << 20); + TEST_ASSERT_MSG(ctx != nullptr, "ggml_init failed"); + if (!ctx) { + ggml_backend_free(backend); + std::fprintf(stderr, " FAIL\n"); + return; + } + + ggml_tensor * q = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, head_dim, n_tokens, n_heads); + ggml_tensor * kv = ggml_new_tensor_3d( + ctx, GGML_TYPE_F16, head_dim, n_kv, 1); + ggml_tensor * mask = ggml_new_tensor_2d( + ctx, GGML_TYPE_F16, n_kv, n_tokens); + ggml_tensor * topk = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, selected_rows, n_tokens); + ggml_tensor * reference = ggml_flash_attn_ext( + ctx, q, kv, kv, mask, 1.0f / std::sqrt((float) head_dim), + 0.0f, 0.0f); + ggml_flash_attn_ext_set_ds4_sparse( + reference, raw_rows, raw_window, -selected_rows, 1); + ggml_tensor * candidate = ggml_flash_attn_ext( + ctx, q, kv, kv, mask, 1.0f / std::sqrt((float) head_dim), + 0.0f, 0.0f); + ggml_flash_attn_ext_set_ds4_sparse( + candidate, raw_rows, raw_window, -selected_rows, 1); + ggml_flash_attn_ext_set_ds4_indexer_topk(candidate, topk); + ggml_set_output(reference); + ggml_set_output(candidate); + TEST_ASSERT_MSG(ggml_backend_supports_op(backend, reference), + "GPU rejected compact F16 attention reference"); + TEST_ASSERT_MSG(ggml_backend_supports_op(backend, candidate), + "GPU rejected streaming F16 attention candidate"); + + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 64, false); + ggml_build_forward_expand(graph, reference); + ggml_build_forward_expand(graph, candidate); + ggml_gallocr_t alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + const bool allocated = ggml_gallocr_alloc_graph(alloc, graph); + TEST_ASSERT_MSG(allocated, "streaming top-k graph allocation failed"); + if (allocated) { + std::vector q_data( + (size_t) head_dim * n_tokens * n_heads); + std::vector kv_data((size_t) head_dim * n_kv); + std::vector mask_data( + (size_t) n_kv * n_tokens, ggml_fp32_to_fp16(-1.0e30f)); + std::vector topk_data( + (size_t) selected_rows * n_tokens); + uint32_t rng = 0x91e10da5u; + const auto sample = [&rng]() { + rng = rng * 1664525u + 1013904223u; + return ((int32_t) (rng >> 8) - 8388608) / 8388608.0f; + }; + for (float & value : q_data) { + value = 1.5f * sample(); + } + for (ggml_fp16_t & value : kv_data) { + value = ggml_fp32_to_fp16(1.5f * sample()); + } + for (int token = 0; token < n_tokens; ++token) { + ggml_fp16_t * token_mask = + mask_data.data() + (size_t) token * n_kv; + for (int row = 0; row < raw_rows; ++row) { + token_mask[row] = ggml_fp32_to_fp16(0.0f); + } + for (int rank = 0; rank < selected_rows; ++rank) { + const int row = + (token * 17 + selected_rows - 1 - rank) % n_comp_rows; + topk_data[(size_t) token * selected_rows + rank] = row; + token_mask[raw_rows + row] = ggml_fp32_to_fp16(0.0f); + } + } + ggml_backend_tensor_set(q, q_data.data(), 0, + q_data.size() * sizeof(float)); + ggml_backend_tensor_set(kv, kv_data.data(), 0, + kv_data.size() * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(mask, mask_data.data(), 0, + mask_data.size() * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(topk, topk_data.data(), 0, + topk_data.size() * sizeof(int32_t)); + + ScopedCudaGraphOverrides eager( + /*disable_graphs=*/true, + /*mmvq_max_ncols=*/0, + /*skip_property_check=*/false); + TEST_ASSERT_MSG( + ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS, + "streaming top-k attention graph failed"); + std::vector reference_data( + (size_t) ggml_nelements(reference)); + std::vector candidate_data(reference_data.size()); + ggml_backend_tensor_get( + reference, reference_data.data(), 0, + reference_data.size() * sizeof(float)); + ggml_backend_tensor_get( + candidate, candidate_data.data(), 0, + candidate_data.size() * sizeof(float)); + + bool finite = true; + bool bounded = true; + double mean_abs = 0.0; + float max_abs = 0.0f; + for (size_t i = 0; i < reference_data.size(); ++i) { + finite = finite && std::isfinite(candidate_data[i]); + const float error = std::abs( + reference_data[i] - candidate_data[i]); + max_abs = std::max(max_abs, error); + mean_abs += error; + bounded = bounded && nearly_equal( + reference_data[i], candidate_data[i], 5.0e-4f, 5.0e-4f); + } + mean_abs /= reference_data.size(); + std::fprintf(stderr, " max_abs=%.3g mean_abs=%.3g", + max_abs, mean_abs); + TEST_ASSERT_MSG(finite, + "streaming top-k output must be finite"); + TEST_ASSERT_MSG(bounded, + "streaming top-k exceeded numeric smoke tolerance"); + } + + ggml_gallocr_free(alloc); + ggml_free(ctx); + ggml_backend_free(backend); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + static void test_ds4_indexer_score_packed_q4_gpu() { std::fprintf(stderr, " test_ds4_indexer_score_packed_q4_gpu ..."); #if !defined(GGML_USE_HIP) @@ -4168,6 +4317,7 @@ int main() { #if defined(GGML_USE_CUDA) || defined(GGML_USE_HIP) test_ds4_flash_attention_keep_cap_gpu(); test_ds4_flash_attention_parallel_index_scan_gpu(); + test_ds4_flash_attention_streaming_topk_gpu(); test_ds4_indexer_score_packed_q4_gpu(); test_ds4_topk_block_radix_gpu(); test_ds4_flash_attention_inverse_rope_fallback_gpu(); From a7c2f23174ac9284d0c8ae44d60380986f901c9d Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Mon, 31 Aug 2026 14:32:55 +0200 Subject: [PATCH 08/16] perf(ds4): derive ratio-4 prefill visibility --- .../llama.cpp/ggml/src/ggml-cuda/fattn.cu | 192 +++++++++++++----- server/src/deepseek4/deepseek4_backend.cpp | 3 + server/src/deepseek4/deepseek4_graph.cpp | 32 ++- server/tests/test_deepseek4_unit.cpp | 86 +++++++- 4 files changed, 259 insertions(+), 54 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu index 3b90a44f9..baf897329 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu @@ -300,6 +300,31 @@ __global__ static void ds4_fa_visibility_bounds_kernel( } } +// Exact DS4 ratio-4 layer-major visibility without a materialized mask. +// Physical raw rows are [prior chronological SWA | current chunk]. A +// compressed row becomes visible only after its four source tokens complete. +__global__ static void ds4_fa_ratio4_causal_bounds_kernel( + int * bounds, + int n_tokens, + int n_kv, + int raw_rows, + int raw_window, + int kv_start) { + const int t = (int) blockIdx.x * (int) blockDim.x + + (int) threadIdx.x; + if (t >= n_tokens) return; + + const int prior_rows = raw_rows - n_tokens; + const int n_comp_rows = n_kv - raw_rows; + const int visible_comp = min(n_comp_rows, (kv_start + t + 1) / 4); + int * token_bounds = bounds + (size_t) t * 4; + token_bounds[0] = max(0, prior_rows + t - raw_window + 1); + token_bounds[1] = prior_rows + t; + token_bounds[2] = visible_comp > 0 ? raw_rows : n_kv; + token_bounds[3] = visible_comp > 0 + ? raw_rows + visible_comp - 1 : -1; +} + // Convert an externally selected compressed-row mask into exact lookup tables. // selected_rows preserves ascending physical-row order for the value pass. // owner_offsets/owner_ranks group those ascending ranks by the thread that @@ -480,7 +505,7 @@ __global__ static void ds4_fa_indexed_rows_parallel_kernel( // A shared-memory bitonic sort restores ascending physical-row order, matching // the old top-k -> mask -> physical scan path and therefore preserving each // reduction lane's accumulation order exactly. -template +template __global__ static void ds4_fa_indexed_rows_topk_kernel( const Mask * mask, const int32_t * topk, @@ -491,7 +516,8 @@ __global__ static void ds4_fa_indexed_rows_topk_kernel( int n_tokens, int n_kv, int raw_rows, - int capacity) { + int capacity, + int kv_start = 0) { const int t = (int) blockIdx.x; const int tid = (int) threadIdx.x; if (t >= n_tokens) return; @@ -504,7 +530,8 @@ __global__ static void ds4_fa_indexed_rows_topk_kernel( __shared__ int count; const int n_comp_rows = n_kv - raw_rows; - const Mask * token_mask = mask + (size_t) t * n_kv; + const Mask * token_mask = RATIO4_CAUSAL + ? nullptr : mask + (size_t) t * n_kv; const int32_t * token_topk = topk + (size_t) t * capacity; int * token_rows = selected_rows + (size_t) t * capacity; int * token_owner_offsets = owner_offsets + (size_t) t * (N_OWNERS + 1); @@ -514,8 +541,14 @@ __global__ static void ds4_fa_indexed_rows_topk_kernel( if (tid < capacity) { const int comp = token_topk[tid]; const int physical = raw_rows + comp; - if (comp >= 0 && comp < n_comp_rows && - ds4_fa_load(token_mask + physical) > -1.0e20f) { + bool visible = comp >= 0 && comp < n_comp_rows; + if constexpr (RATIO4_CAUSAL) { + visible = visible && comp < (kv_start + t + 1) / 4; + } else { + visible = visible && + ds4_fa_load(token_mask + physical) > -1.0e20f; + } + if (visible) { row = physical; } } @@ -1092,7 +1125,7 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_kernel( // every visible row keeps its original owner thread, dot-product order, // reduction tree, softmax order, and value-accumulation position. template + bool MASKLESS_CAUSAL, int VALUES_PER_THREAD> __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( float * dst, const float * q, @@ -1246,8 +1279,11 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( } } - const float mask_v = ds4_fa_load( - mask + (size_t) t * n_kv + r); + float mask_v = 0.0f; + if constexpr (!MASKLESS_CAUSAL) { + mask_v = ds4_fa_load( + mask + (size_t) t * n_kv + r); + } const bool visible = mask_v > -1.0e20f; float dot[HEADS_PER_BLOCK] = {}; if (visible) { @@ -1567,7 +1603,7 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( // remains in the graph/backend layer. template + bool FAST_EXP = false, bool MASKLESS_CAUSAL = false> __global__ static void ds4_flash_attn_d512_streaming_topk_kernel( float * dst, const float * q, @@ -1620,7 +1656,8 @@ __global__ static void ds4_flash_attn_d512_streaming_topk_kernel( const int total_rows = raw_count + indexed_count; const int * token_rows = indexed_rows + (size_t) token * indexed_capacity; - const Mask * token_mask = mask + (size_t) token * n_kv; + const Mask * token_mask = MASKLESS_CAUSAL + ? nullptr : mask + (size_t) token * n_kv; const float * qh = q + (size_t) token * q_stride_token + (size_t) head * q_stride_head; @@ -1661,9 +1698,14 @@ __global__ static void ds4_flash_attn_d512_streaming_topk_kernel( row = token_rows[selected - raw_count]; } staged_rows[tid] = row; - staged_masks[tid] = row >= 0 && row < n_kv - ? ds4_fa_load(token_mask + row) - : -3.402823466e38f; + if constexpr (MASKLESS_CAUSAL) { + staged_masks[tid] = row >= 0 && row < n_kv + ? 0.0f : -3.402823466e38f; + } else { + staged_masks[tid] = row >= 0 && row < n_kv + ? ds4_fa_load(token_mask + row) + : -3.402823466e38f; + } } __syncthreads(); @@ -1865,7 +1907,8 @@ static bool ds4_launch_flash_attn_d512_grouped( return true; } -template +template static bool ds4_launch_flash_attn_d512_grouped_compact( ggml_tensor * dst, const ggml_tensor * Q, @@ -1895,7 +1938,7 @@ static bool ds4_launch_flash_attn_d512_grouped_compact( const float * forward_rope_coefficients, size_t shmem, cudaStream_t stream) { - GGML_ASSERT(mask && visibility_bounds); + GGML_ASSERT((MASKLESS_CAUSAL || mask) && visibility_bounds); if constexpr (INDEXED_MASK) { GGML_ASSERT(indexed_rows && indexed_counts && indexed_owner_offsets && indexed_owner_ranks); @@ -1903,14 +1946,15 @@ static bool ds4_launch_flash_attn_d512_grouped_compact( dim3 grid( (unsigned) n_tokens, (unsigned) (n_heads / HEADS_PER_BLOCK), 1); - if (kv_f16 && mask->type == GGML_TYPE_F16) { + if (kv_f16 && (MASKLESS_CAUSAL || mask->type == GGML_TYPE_F16)) { ds4_flash_attn_d512_shared_kv_grouped_compact_kernel< - half, half, HEADS_PER_BLOCK, INDEXED_MASK, VALUES_PER_THREAD> + half, half, HEADS_PER_BLOCK, INDEXED_MASK, MASKLESS_CAUSAL, + VALUES_PER_THREAD> <<>>( (float *) dst->data, (const float *) Q->data, q_stride_token, q_stride_head, (const half *) K->data, (const half *) V->data, - (const half *) mask->data, + mask ? (const half *) mask->data : nullptr, sinks ? (const float *) sinks->data : nullptr, n_tokens, n_heads, n_kv, scale, raw_rows, raw_score_capacity, score_stride, visibility_bounds, @@ -1918,14 +1962,16 @@ static bool ds4_launch_flash_attn_d512_grouped_compact( indexed_owner_offsets, indexed_owner_ranks, indexed_capacity, inverse_rope, inverse_rope_coefficients, forward_rope_coefficients); - } else if (kv_f32 && mask->type == GGML_TYPE_F32) { + } else if (kv_f32 && + (MASKLESS_CAUSAL || mask->type == GGML_TYPE_F32)) { ds4_flash_attn_d512_shared_kv_grouped_compact_kernel< - float, float, HEADS_PER_BLOCK, INDEXED_MASK, VALUES_PER_THREAD> + float, float, HEADS_PER_BLOCK, INDEXED_MASK, MASKLESS_CAUSAL, + VALUES_PER_THREAD> <<>>( (float *) dst->data, (const float *) Q->data, q_stride_token, q_stride_head, (const float *) K->data, (const float *) V->data, - (const float *) mask->data, + mask ? (const float *) mask->data : nullptr, sinks ? (const float *) sinks->data : nullptr, n_tokens, n_heads, n_kv, scale, raw_rows, raw_score_capacity, score_stride, visibility_bounds, @@ -1933,9 +1979,11 @@ static bool ds4_launch_flash_attn_d512_grouped_compact( indexed_owner_offsets, indexed_owner_ranks, indexed_capacity, inverse_rope, inverse_rope_coefficients, forward_rope_coefficients); - } else if (kv_f32 && mask->type == GGML_TYPE_F16) { + } else if (!MASKLESS_CAUSAL && kv_f32 && + mask->type == GGML_TYPE_F16) { ds4_flash_attn_d512_shared_kv_grouped_compact_kernel< - float, half, HEADS_PER_BLOCK, INDEXED_MASK, VALUES_PER_THREAD> + float, half, HEADS_PER_BLOCK, INDEXED_MASK, MASKLESS_CAUSAL, + VALUES_PER_THREAD> <<>>( (float *) dst->data, (const float *) Q->data, q_stride_token, q_stride_head, @@ -1965,6 +2013,7 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32_supported(const ggml_tensor * dst) const ggml_tensor * mask = dst->src[3]; const ggml_tensor * sinks = dst->src[4]; const ggml_tensor * indexer_topk = dst->src[5]; + const bool ratio4_causal = indexer_topk && !mask; const bool kv_f32 = K && V && K->type == GGML_TYPE_F32 && V->type == GGML_TYPE_F32; const bool kv_f16 = K && V && K->type == GGML_TYPE_F16 && @@ -2028,7 +2077,7 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32_supported(const ggml_tensor * dst) } if (raw_rows < 0 || raw_rows > n_kv || (ds4_layout != 0 && (raw_window <= 0 || sparse_block_size <= 0)) || - (sparse_keep_rows != 0 && !mask) || + (sparse_keep_rows != 0 && !mask && !ratio4_causal) || (rope_flags & ~3) != 0 || ((rope_flags & 2) != 0 && (rope_flags & 1) == 0)) { return false; @@ -2059,10 +2108,19 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32_supported(const ggml_tensor * dst) (size_t) group4 * 4 * sizeof(int) + ((rope_flags & 2) != 0 ? (size_t) group4 * 64 * sizeof(float) : 0); - if (!mask || raw_rows <= 0 || n_comp_rows <= 0 || + if (raw_rows <= 0 || n_comp_rows <= 0 || n_heads % group4 != 0 || compact_group4_shmem > 24 * 1024) { return false; } + if (ratio4_causal) { + const int kv_start = ggml_get_op_params_i32(dst, 8); + const int prior_rows = raw_rows - n_tokens; + if ((rope_flags & 1) == 0 || n_tokens <= raw_window || + prior_rows != std::min(kv_start, raw_window) || + n_comp_rows != (kv_start + n_tokens) / 4) { + return false; + } + } } return true; @@ -2080,6 +2138,7 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( const ggml_tensor * mask = dst->src[3]; const ggml_tensor * sinks = dst->src[4]; const ggml_tensor * indexer_topk = dst->src[5]; + const bool ratio4_causal = indexer_topk && !mask; const bool kv_f32 = K->type == GGML_TYPE_F32; const bool kv_f16 = K->type == GGML_TYPE_F16; const int n_tokens = (int) Q->ne[1]; @@ -2212,7 +2271,7 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( // Compacting score storage lets both shapes keep the four-head kernel at // two-block occupancy. Ordinary dense shapes avoid the extra bounds scan. const bool compact_group4 = - !sparse && mask && n_heads % group4 == 0 && + !sparse && (mask || ratio4_causal) && n_heads % group4 == 0 && (raw_rows > raw_window || indexed_mask) && (indexed_mask || group4_shmem > 24 * 1024) && compact_group4_shmem <= 24 * 1024; @@ -2238,9 +2297,17 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( (size_t) n_tokens * indexed_capacity); const bool parallel_index_scan = n_comp_rows > 512 && getenv("GGML_DS4_FA_SERIAL_INDEX_SCAN") == nullptr; - if (mask->type == GGML_TYPE_F16) { + if (ratio4_causal) { + ds4_fa_indexed_rows_topk_kernel + <<>>( + nullptr, (const int32_t *) indexer_topk->data, + indexed_rows, indexed_counts, + indexed_owner_offsets, indexed_owner_ranks, + n_tokens, n_kv, raw_rows, indexed_capacity, + inverse_rope.kv_start); + } else if (mask->type == GGML_TYPE_F16) { if (indexer_topk) { - ds4_fa_indexed_rows_topk_kernel<<>>( + ds4_fa_indexed_rows_topk_kernel<<>>( (const half *) mask->data, (const int32_t *) indexer_topk->data, indexed_rows, indexed_counts, @@ -2261,7 +2328,7 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( } } else { if (indexer_topk) { - ds4_fa_indexed_rows_topk_kernel<<>>( + ds4_fa_indexed_rows_topk_kernel<<>>( (const float *) mask->data, (const int32_t *) indexer_topk->data, indexed_rows, indexed_counts, @@ -2283,7 +2350,12 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( } CUDA_CHECK(cudaGetLastError()); } - if (mask->type == GGML_TYPE_F16) { + if (ratio4_causal) { + ds4_fa_ratio4_causal_bounds_kernel<<< + (n_tokens + 255) / 256, 256, 0, stream>>>( + visibility_bounds, n_tokens, n_kv, raw_rows, + raw_window, inverse_rope.kv_start); + } else if (mask->type == GGML_TYPE_F16) { ds4_fa_visibility_bounds_kernel<<>>( (const half *) mask->data, visibility_bounds, n_tokens, n_kv, raw_rows); @@ -2301,7 +2373,7 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( const int device_warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; if (indexed_mask && indexer_topk && - kv_f16 && mask->type == GGML_TYPE_F16 && + kv_f16 && (ratio4_causal || mask->type == GGML_TYPE_F16) && K->data == V->data && n_heads % 16 == 0 && device_warp_size == 32 && n_tokens >= streaming_min_tokens && active_row_upper_bound > 0 && @@ -2310,25 +2382,53 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( const dim3 streaming_grid( (unsigned) n_tokens, (unsigned) (n_heads / streaming_heads), 1); - ds4_flash_attn_d512_streaming_topk_kernel< - half, half, streaming_heads, 16, true, true> - <<>>( - (float *) dst->data, (const float *) Q->data, - q_stride_token, q_stride_head, - (const half *) K->data, - (const half *) mask->data, - sinks ? (const float *) sinks->data : nullptr, - n_tokens, n_heads, n_kv, scale, - visibility_bounds, indexed_rows, indexed_counts, - indexed_capacity, inverse_rope, - inverse_rope_coefficients, - forward_rope_coefficients); + if (ratio4_causal) { + ds4_flash_attn_d512_streaming_topk_kernel< + half, half, streaming_heads, 16, true, true, true> + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const half *) K->data, nullptr, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, + visibility_bounds, indexed_rows, indexed_counts, + indexed_capacity, inverse_rope, + inverse_rope_coefficients, + forward_rope_coefficients); + } else { + ds4_flash_attn_d512_streaming_topk_kernel< + half, half, streaming_heads, 16, true, true, false> + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const half *) K->data, + (const half *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, + visibility_bounds, indexed_rows, indexed_counts, + indexed_capacity, inverse_rope, + inverse_rope_coefficients, + forward_rope_coefficients); + } CUDA_CHECK(cudaGetLastError()); return true; } if (indexed_mask) { + if (ratio4_causal) { + return ds4_launch_flash_attn_d512_grouped_compact< + group4, true, true, 4>( + dst, Q, K, V, mask, sinks, kv_f16, kv_f32, + n_tokens, n_heads, n_kv, scale, raw_rows, + raw_window, compact_score_stride, visibility_bounds, + indexed_rows, indexed_counts, + indexed_owner_offsets, indexed_owner_ranks, + indexed_capacity, q_stride_token, q_stride_head, + inverse_rope, inverse_rope_coefficients, + forward_rope_coefficients, + compact_group4_shmem, stream); + } return ds4_launch_flash_attn_d512_grouped_compact< - group4, true, 4>( + group4, true, false, 4>( dst, Q, K, V, mask, sinks, kv_f16, kv_f32, n_tokens, n_heads, n_kv, scale, raw_rows, raw_window, compact_score_stride, visibility_bounds, @@ -2341,7 +2441,7 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( compact_group4_shmem, stream); } return ds4_launch_flash_attn_d512_grouped_compact< - group4, false, 4>( + group4, false, false, 4>( dst, Q, K, V, mask, sinks, kv_f16, kv_f32, n_tokens, n_heads, n_kv, scale, raw_rows, raw_window, compact_score_stride, visibility_bounds, diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index c8750bb28..6c0882df3 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -2019,6 +2019,7 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, save_snapshot && !snapshot_saved, spec_snap_from, spec_snap_to); } + const auto chunk_t0 = Clock::now(); // Bulk prompt graphs and the final DSpark feature-capture graph have // different HC/owner arena shapes. Once all earlier chunks are @@ -2131,6 +2132,8 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, if (timing) { add_step_tel(tel_acc, step_tel); steps++; + log_step_tel("prefill-chunk", n_tok, 1, + elapsed_s(chunk_t0), step_tel); } last_logits_ = std::move(logits); pos += n_tok; diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index f7f9b1877..b1ed54f26 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -2018,6 +2018,11 @@ static ggml_tensor * build_mla_attention( // [n_kv,n_query] F16; the explicit path broadcasts the same values over // heads in F32. ggml_tensor * score_mask = nullptr; + // Ratio-4 sparse prefill already carries the authoritative compressed + // row IDs. The CUDA/HIP kernel can derive the raw causal window and the + // completed compressed-row frontier from kv_start and the query index. + // Keep every other attention shape on the explicit mask contract. + const bool direct_indexer_topk = indexer_topk && n_tokens > w.n_swa; const bool exact_numerical_bands = attention_impl == DeepSeek4AttentionImpl::DenseFlash && causal_batch && @@ -2030,7 +2035,7 @@ static ggml_tensor * build_mla_attention( } else if (masked_kv) { score_mask = ggml_reshape_2d(ctx, cached_inputs->attn_row_mask, n_attn, 1); - } else if (layer_major_batch) { + } else if (layer_major_batch && !direct_indexer_topk) { // Per-token causal mask over [prior rows | current rows | comp rows]. ggml_tensor * cmask = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_attn, 1, n_tokens); ggml_set_input(cmask); @@ -2095,9 +2100,8 @@ static ggml_tensor * build_mla_attention( } // Long sparse prefill already has the authoritative selected rows. Pass // them through directly instead of materializing and rescanning a mask. - const bool direct_indexer_topk = indexer_topk && n_tokens > w.n_swa; if (indexer_topk) { - if (!score_mask) { + if (!score_mask && !direct_indexer_topk) { score_mask = ggml_new_tensor_2d( ctx, GGML_TYPE_F32, n_attn, n_tokens); ggml_set_input(score_mask); @@ -6297,6 +6301,8 @@ static int ds4_try_layer_major_prefill( const int64_t hc_dim = (int64_t) n_embd * n_hc; const int64_t mix_dim = 2 * (int64_t) n_hc + (int64_t) n_hc * n_hc; const int next_pos = kv_start + n_tokens; + size_t retained_f32_bytes = 0; + size_t scratch_bytes = 0; const std::vector * capture_layer_ids = verify_hooks ? verify_hooks->capture_layer_ids : nullptr; @@ -6481,6 +6487,8 @@ static int ds4_try_layer_major_prefill( !ggml_gallocr_alloc_graph(alloc, layer.gf)) { return fail("cached scratch allocation failed", il); } + scratch_bytes = std::max( + scratch_bytes, ggml_gallocr_get_buffer_size(alloc, 0)); if (telemetry) { telemetry->full_graph_build_us += ds4_elapsed_us( alloc_t0, Ds4TimingClock::now()); @@ -6498,6 +6506,7 @@ static int ds4_try_layer_major_prefill( sizeof(int64_t) * b.values.size()); } for (const auto & b : layer.f32_array_inputs) { + retained_f32_bytes += sizeof(float) * b.values.size(); ggml_backend_tensor_set(b.tensor, b.values.data(), 0, sizeof(float) * b.values.size()); } @@ -6544,6 +6553,13 @@ static int ds4_try_layer_major_prefill( } } cache.cur_pos = next_pos; + if (telemetry) { + std::fprintf(stderr, + "[deepseek4-timing] layer-major kv_start=%d " + "tokens=%d f32_bindings=%zu scratch=%zu cache=hit\n", + kv_start, n_tokens, retained_f32_bytes, + scratch_bytes); + } return out_logits.empty() ? -1 : 1; } @@ -6710,6 +6726,8 @@ static int ds4_try_layer_major_prefill( if (!cached_layer) ggml_free(ctx); return fail("scratch allocation failed", il); } + scratch_bytes = std::max( + scratch_bytes, ggml_gallocr_get_buffer_size(alloc, 0)); for (const auto & b : i32_inputs) { ggml_backend_tensor_set(b.tensor, &b.value, 0, sizeof(b.value)); } @@ -6722,6 +6740,7 @@ static int ds4_try_layer_major_prefill( sizeof(int64_t) * b.values.size()); } for (const auto & b : f32_array_inputs) { + retained_f32_bytes += sizeof(float) * b.values.size(); ggml_backend_tensor_set(b.tensor, b.values.data(), 0, sizeof(float) * b.values.size()); } @@ -6781,6 +6800,13 @@ static int ds4_try_layer_major_prefill( std::swap(state_in, state_out); } + if (telemetry) { + std::fprintf(stderr, + "[deepseek4-timing] layer-major kv_start=%d tokens=%d " + "f32_bindings=%zu scratch=%zu cache=%s\n", + kv_start, n_tokens, retained_f32_bytes, scratch_bytes, + cache_build ? "built" : "uncached"); + } if (cache_build) { graph_cache->ready = true; } else { diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index 55d7a0440..6450e3cdd 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -64,6 +64,51 @@ static bool nearly_equal(float a, float b, float atol = 1.0e-5f, float rtol = 1. return diff <= atol + rtol * scale; } +static void test_ds4_ratio4_causal_visibility_formula() { + std::fprintf(stderr, " test_ds4_ratio4_causal_visibility_formula ..."); + constexpr int raw_window = 128; + const auto check_chunk = [&](int kv_start, int n_tokens) { + const int prior_rows = std::min(kv_start, raw_window); + const int raw_rows = prior_rows + n_tokens; + const int n_comp_rows = (kv_start + n_tokens) / 4; + const int probes[] = { + 0, 1, 2, 3, 127, 128, 2050, 2051, 8191, 8192, 8193, + kv_start + n_tokens - 1, + }; + for (int position : probes) { + if (position < kv_start || position >= kv_start + n_tokens) { + continue; + } + const int token = position - kv_start; + int reference_first = raw_rows; + int reference_last = -1; + for (int row = 0; row < raw_rows; ++row) { + const int row_position = kv_start - prior_rows + row; + if (row_position >= position - raw_window + 1 && + row_position <= position) { + reference_first = std::min(reference_first, row); + reference_last = row; + } + } + const int analytic_first = std::max( + 0, prior_rows + token - raw_window + 1); + const int analytic_last = prior_rows + token; + const int analytic_comp = std::min( + n_comp_rows, (position + 1) / 4); + int reference_comp = 0; + for (int row = 0; row < n_comp_rows; ++row) { + reference_comp += row < (position + 1) / 4; + } + TEST_ASSERT(reference_first == analytic_first); + TEST_ASSERT(reference_last == analytic_last); + TEST_ASSERT(reference_comp == analytic_comp); + } + }; + check_chunk(0, 8192); + check_chunk(8192, 941); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + static ggml_tensor * test_hc_row_normalize(ggml_context * ctx, ggml_tensor * x) { ggml_tensor * sums = ggml_sum_rows(ctx, x); return ggml_div(ctx, x, ggml_repeat(ctx, sums, x)); @@ -2846,11 +2891,13 @@ static void test_ds4_flash_attention_streaming_topk_gpu() { constexpr int head_dim = 512; constexpr int n_heads = 64; - constexpr int n_tokens = 64; - constexpr int raw_rows = 128; + constexpr int kv_start = 8192; + constexpr int n_tokens = 129; constexpr int raw_window = 128; + constexpr int prior_rows = raw_window; + constexpr int raw_rows = prior_rows + n_tokens; constexpr int selected_rows = 512; - constexpr int n_comp_rows = 1920; + constexpr int n_comp_rows = (kv_start + n_tokens) / 4; constexpr int n_kv = raw_rows + n_comp_rows; ggml_context * ctx = make_test_context(4u << 20); @@ -2880,16 +2927,31 @@ static void test_ds4_flash_attention_streaming_topk_gpu() { ggml_flash_attn_ext_set_ds4_sparse( candidate, raw_rows, raw_window, -selected_rows, 1); ggml_flash_attn_ext_set_ds4_indexer_topk(candidate, topk); + ggml_tensor * maskless = ggml_flash_attn_ext( + ctx, q, kv, kv, nullptr, 1.0f / std::sqrt((float) head_dim), + 0.0f, 0.0f); + ggml_flash_attn_ext_set_ds4_sparse( + maskless, raw_rows, raw_window, -selected_rows, 1); + ggml_flash_attn_ext_set_ds4_indexer_topk(maskless, topk); + for (ggml_tensor * output : {reference, candidate, maskless}) { + ggml_flash_attn_ext_set_ds4_inverse_rope( + output, kv_start, 10000.0f, 1.0f, 0.0f, 1.0f, + 32.0f, 1.0f, 163840, false); + } ggml_set_output(reference); ggml_set_output(candidate); + ggml_set_output(maskless); TEST_ASSERT_MSG(ggml_backend_supports_op(backend, reference), "GPU rejected compact F16 attention reference"); TEST_ASSERT_MSG(ggml_backend_supports_op(backend, candidate), "GPU rejected streaming F16 attention candidate"); + TEST_ASSERT_MSG(ggml_backend_supports_op(backend, maskless), + "GPU rejected analytic ratio-4 attention candidate"); ggml_cgraph * graph = ggml_new_graph_custom(ctx, 64, false); ggml_build_forward_expand(graph, reference); ggml_build_forward_expand(graph, candidate); + ggml_build_forward_expand(graph, maskless); ggml_gallocr_t alloc = ggml_gallocr_new( ggml_backend_get_default_buffer_type(backend)); const bool allocated = ggml_gallocr_alloc_graph(alloc, graph); @@ -2916,14 +2978,20 @@ static void test_ds4_flash_attention_streaming_topk_gpu() { for (int token = 0; token < n_tokens; ++token) { ggml_fp16_t * token_mask = mask_data.data() + (size_t) token * n_kv; - for (int row = 0; row < raw_rows; ++row) { + const int raw_first = std::max( + 0, prior_rows + token - raw_window + 1); + const int raw_last = prior_rows + token; + for (int row = raw_first; row <= raw_last; ++row) { token_mask[row] = ggml_fp32_to_fp16(0.0f); } for (int rank = 0; rank < selected_rows; ++rank) { const int row = (token * 17 + selected_rows - 1 - rank) % n_comp_rows; topk_data[(size_t) token * selected_rows + rank] = row; - token_mask[raw_rows + row] = ggml_fp32_to_fp16(0.0f); + if (row < (kv_start + token + 1) / 4) { + token_mask[raw_rows + row] = + ggml_fp32_to_fp16(0.0f); + } } } ggml_backend_tensor_set(q, q_data.data(), 0, @@ -2945,12 +3013,16 @@ static void test_ds4_flash_attention_streaming_topk_gpu() { std::vector reference_data( (size_t) ggml_nelements(reference)); std::vector candidate_data(reference_data.size()); + std::vector maskless_data(reference_data.size()); ggml_backend_tensor_get( reference, reference_data.data(), 0, reference_data.size() * sizeof(float)); ggml_backend_tensor_get( candidate, candidate_data.data(), 0, candidate_data.size() * sizeof(float)); + ggml_backend_tensor_get( + maskless, maskless_data.data(), 0, + maskless_data.size() * sizeof(float)); bool finite = true; bool bounded = true; @@ -2964,6 +3036,9 @@ static void test_ds4_flash_attention_streaming_topk_gpu() { mean_abs += error; bounded = bounded && nearly_equal( reference_data[i], candidate_data[i], 5.0e-4f, 5.0e-4f); + TEST_ASSERT_MSG( + candidate_data[i] == maskless_data[i], + "analytic ratio-4 visibility changed attention output"); } mean_abs /= reference_data.size(); std::fprintf(stderr, " max_abs=%.3g mean_abs=%.3g", @@ -4282,6 +4357,7 @@ int main() { test_indexer_qat_cpu(backend); test_indexer_score_cpu(backend); test_indexer_mask_cpu(backend); + test_ds4_ratio4_causal_visibility_formula(); test_hash_routing_lookup(); test_raw_ring_spans_after_wrap(); test_auto_split_computation(); From ccbff860979cf9479cb7b75dcb7bf57f763aad65 Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Mon, 31 Aug 2026 14:33:13 +0200 Subject: [PATCH 09/16] perf(ds4): schedule five prefill bands --- server/src/deepseek4/deepseek4_internal.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index 1be9f7d60..a31084bd4 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -28,10 +28,10 @@ namespace dflash::common { -// Layer-major prefill may schedule four 2K numerical bands while preserving +// Layer-major prefill may schedule five 2K numerical bands while preserving // the raw-cache rounding boundary between them. inline constexpr int DS4_NUMERICAL_PREFILL_BAND = 2048; -inline constexpr int DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS = 8192; +inline constexpr int DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS = 10240; // Normal verification stays within one ratio-4 compressor window. Q5 is an // explicit opt-in whose fused graph models a second boundary. inline constexpr int DS4_CONSERVATIVE_VERIFY_MAX_TOKENS = 4; From 654d505f33c83703aa73d1cb9a5333c55d8ae9e0 Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Mon, 31 Aug 2026 16:10:05 +0200 Subject: [PATCH 10/16] fix(ds4): keep maskless prefill out of ring mask --- server/src/deepseek4/deepseek4_graph.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index b1ed54f26..2b8026b15 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -2061,7 +2061,7 @@ static ggml_tensor * build_mla_attention( } f32_array_inputs->push_back({cmask, std::move(mvals)}); score_mask = ggml_reshape_2d(ctx, cmask, n_attn, n_tokens); - } else if (causal_batch) { + } else if (causal_batch && !layer_major_batch) { // Speculative verification keeps the physical ring order and // appends snapshots of rows overwritten by later batch tokens. ggml_tensor * cmask = ggml_new_tensor_3d( From 1ba93fe51022c5cf15a85e1369ddf1f728a28e70 Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Mon, 31 Aug 2026 16:46:22 +0200 Subject: [PATCH 11/16] chore(ds4): remove prefill experiment telemetry --- server/src/deepseek4/deepseek4_backend.cpp | 4 ---- server/src/deepseek4/deepseek4_graph.cpp | 23 ---------------------- 2 files changed, 27 deletions(-) diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 6c0882df3..7b68ef8e1 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -2019,8 +2019,6 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, save_snapshot && !snapshot_saved, spec_snap_from, spec_snap_to); } - const auto chunk_t0 = Clock::now(); - // Bulk prompt graphs and the final DSpark feature-capture graph have // different HC/owner arena shapes. Once all earlier chunks are // complete, retire their reusable prefill arenas before entering the @@ -2132,8 +2130,6 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, if (timing) { add_step_tel(tel_acc, step_tel); steps++; - log_step_tel("prefill-chunk", n_tok, 1, - elapsed_s(chunk_t0), step_tel); } last_logits_ = std::move(logits); pos += n_tok; diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 2b8026b15..4158a0ffd 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -6301,9 +6301,6 @@ static int ds4_try_layer_major_prefill( const int64_t hc_dim = (int64_t) n_embd * n_hc; const int64_t mix_dim = 2 * (int64_t) n_hc + (int64_t) n_hc * n_hc; const int next_pos = kv_start + n_tokens; - size_t retained_f32_bytes = 0; - size_t scratch_bytes = 0; - const std::vector * capture_layer_ids = verify_hooks ? verify_hooks->capture_layer_ids : nullptr; std::vector * capture_out = @@ -6487,8 +6484,6 @@ static int ds4_try_layer_major_prefill( !ggml_gallocr_alloc_graph(alloc, layer.gf)) { return fail("cached scratch allocation failed", il); } - scratch_bytes = std::max( - scratch_bytes, ggml_gallocr_get_buffer_size(alloc, 0)); if (telemetry) { telemetry->full_graph_build_us += ds4_elapsed_us( alloc_t0, Ds4TimingClock::now()); @@ -6506,7 +6501,6 @@ static int ds4_try_layer_major_prefill( sizeof(int64_t) * b.values.size()); } for (const auto & b : layer.f32_array_inputs) { - retained_f32_bytes += sizeof(float) * b.values.size(); ggml_backend_tensor_set(b.tensor, b.values.data(), 0, sizeof(float) * b.values.size()); } @@ -6553,13 +6547,6 @@ static int ds4_try_layer_major_prefill( } } cache.cur_pos = next_pos; - if (telemetry) { - std::fprintf(stderr, - "[deepseek4-timing] layer-major kv_start=%d " - "tokens=%d f32_bindings=%zu scratch=%zu cache=hit\n", - kv_start, n_tokens, retained_f32_bytes, - scratch_bytes); - } return out_logits.empty() ? -1 : 1; } @@ -6726,8 +6713,6 @@ static int ds4_try_layer_major_prefill( if (!cached_layer) ggml_free(ctx); return fail("scratch allocation failed", il); } - scratch_bytes = std::max( - scratch_bytes, ggml_gallocr_get_buffer_size(alloc, 0)); for (const auto & b : i32_inputs) { ggml_backend_tensor_set(b.tensor, &b.value, 0, sizeof(b.value)); } @@ -6740,7 +6725,6 @@ static int ds4_try_layer_major_prefill( sizeof(int64_t) * b.values.size()); } for (const auto & b : f32_array_inputs) { - retained_f32_bytes += sizeof(float) * b.values.size(); ggml_backend_tensor_set(b.tensor, b.values.data(), 0, sizeof(float) * b.values.size()); } @@ -6800,13 +6784,6 @@ static int ds4_try_layer_major_prefill( std::swap(state_in, state_out); } - if (telemetry) { - std::fprintf(stderr, - "[deepseek4-timing] layer-major kv_start=%d tokens=%d " - "f32_bindings=%zu scratch=%zu cache=%s\n", - kv_start, n_tokens, retained_f32_bytes, scratch_bytes, - cache_build ? "built" : "uncached"); - } if (cache_build) { graph_cache->ready = true; } else { From a93f017ff31f067cfef7bb683996f19d5753a6df Mon Sep 17 00:00:00 2001 From: cheese-cakee Date: Thu, 27 Aug 2026 00:43:12 +0530 Subject: [PATCH 12/16] perf(ds4): fuse expert-major MoE route combine --- server/CMakeLists.txt | 10 + server/deps/llama.cpp/ggml/include/ggml.h | 10 + .../llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c | 42 ++ .../llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu | 24 +- .../ggml/src/ggml-cuda/moe-fused-combine.cu | 147 +++++++ .../ggml/src/ggml-cuda/moe-fused-combine.cuh | 5 + server/deps/llama.cpp/ggml/src/ggml.c | 43 +- server/src/common/moe_hybrid_ffn_eval.cpp | 35 +- server/src/deepseek4/deepseek4_graph.cpp | 46 +- server/test/test_ds4_moe_combine_cuda.cpp | 399 ++++++++++++++++++ 10 files changed, 730 insertions(+), 31 deletions(-) create mode 100644 server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused-combine.cu create mode 100644 server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused-combine.cuh create mode 100644 server/test/test_ds4_moe_combine_cuda.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 0b237c94a..4905e4c8c 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -1061,6 +1061,16 @@ if(DFLASH27B_TESTS) ggml ${DFLASH27B_GGML_BACKEND_TARGET}) list(APPEND _raw_unit_test_targets test_deepseek4_mmid_grouped_cuda) endif() + if(DFLASH27B_GPU_BACKEND STREQUAL "hip" AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_ds4_moe_combine_cuda.cpp") + add_executable(test_ds4_moe_combine_cuda test/test_ds4_moe_combine_cuda.cpp) + set_source_files_properties(test/test_ds4_moe_combine_cuda.cpp PROPERTIES LANGUAGE HIP) + set_target_properties(test_ds4_moe_combine_cuda PROPERTIES HIP_ARCHITECTURES "${_dflash_archs}") + target_include_directories(test_ds4_moe_combine_cuda PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include) + target_link_libraries(test_ds4_moe_combine_cuda PRIVATE + ggml-cpu ggml ${DFLASH27B_GGML_BACKEND_TARGET}) + list(APPEND _raw_unit_test_targets test_ds4_moe_combine_cuda) + endif() # HIP-only standalone build; CUDA backend is covered by aggregated test_server_unit. if(DFLASH27B_GPU_BACKEND STREQUAL "hip" AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_draft_topk_cuda.cpp") # HIP build of the same GPU-vs-CPU parity test. The test source uses CUDA diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index 756193d7a..d0d0f46d6 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -617,6 +617,8 @@ extern "C" { GGML_OP_PAGED_ATTN, + GGML_OP_DS4_MOE_COMBINE, + GGML_OP_COUNT, }; @@ -2747,6 +2749,14 @@ extern "C" { struct ggml_tensor * selected, int raw_rows); + // Direct AST Fused MoE Combine Epilogue: down_e[n_embd, n_used, n_tokens] + + // weights[n_used, n_tokens] + shared_out[n_embd, n_tokens] -> dst[n_embd, n_tokens] + GGML_API struct ggml_tensor * ggml_ds4_moe_fused_combine_shared( + struct ggml_context * ctx, + struct ggml_tensor * down_e, + struct ggml_tensor * weights, + struct ggml_tensor * shared_out); + // TODO: needs to be adapted to ggml_flash_attn_ext GGML_API struct ggml_tensor * ggml_flash_attn_back( struct ggml_context * ctx, diff --git a/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c b/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c index 6756c8383..99a052b02 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c +++ b/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c @@ -249,6 +249,43 @@ static void ggml_compute_forward_ds4_indexer_mask( } } +static void ggml_compute_forward_ds4_moe_combine( + const struct ggml_compute_params * params, + struct ggml_tensor * dst) { + const struct ggml_tensor * down_e = dst->src[0]; + const struct ggml_tensor * weights = dst->src[1]; + const struct ggml_tensor * shared_out = dst->src[2]; + + GGML_ASSERT(down_e && weights); + GGML_ASSERT(down_e->type == GGML_TYPE_F32 && weights->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32); + + const int n_embd = (int) down_e->ne[0]; + const int n_used = (int) down_e->ne[1]; + const int n_tokens = (int) down_e->ne[2]; + + for (int t = params->ith; t < n_tokens; t += params->nth) { + const float * w_row = (const float *) ((const char *) weights->data + (size_t) t * weights->nb[1]); + const float * sh_row = shared_out ? (const float *) ((const char *) shared_out->data + (size_t) t * shared_out->nb[1]) : NULL; + float * dst_row = (float *) ((char *) dst->data + (size_t) t * dst->nb[1]); + + for (int i = 0; i < n_embd; ++i) { + float sum = 0.0f; + for (int e = 0; e < n_used; ++e) { + if (w_row[e] == 0.0f) { + continue; + } + const float * exp_row = (const float *) ((const char *) down_e->data + (size_t) t * down_e->nb[2] + (size_t) e * down_e->nb[1]); + const float prod = exp_row[i] * w_row[e]; + sum += prod; + } + if (sh_row) { + sum += sh_row[i]; + } + dst_row[i] = sum; + } + } +} + #if defined(__ARM_ARCH) struct ggml_arm_arch_features_type { int sve_cnt; @@ -2036,6 +2073,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { ggml_compute_forward_ds4_indexer_mask(params, tensor); } break; + case GGML_OP_DS4_MOE_COMBINE: + { + ggml_compute_forward_ds4_moe_combine(params, tensor); + } break; case GGML_OP_OUT_PROD: { ggml_compute_forward_out_prod(params, tensor); @@ -2588,6 +2629,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_DS4_INDEXER_QAT: case GGML_OP_DS4_INDEXER_SCORE: case GGML_OP_DS4_INDEXER_MASK: + case GGML_OP_DS4_MOE_COMBINE: case GGML_OP_FLASH_ATTN_EXT: case GGML_OP_FLASH_ATTN_SPARSE: case GGML_OP_PAGED_ATTN: diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu index a1c178fde..a7c8c0a66 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu @@ -32,6 +32,7 @@ #include "ggml-cuda/mmq.cuh" #include "ggml-cuda/mmvf.cuh" #include "ggml-cuda/mmvq.cuh" +#include "ggml-cuda/moe-fused-combine.cuh" #include "ggml-cuda/rocmfp3_mix.cuh" #include "ggml-cuda/rocmfp2_mix.cuh" #include "ggml-cuda/norm.cuh" @@ -768,6 +769,7 @@ ggml_backend_cuda_context::~ggml_backend_cuda_context() { if (copy_event != nullptr) { CUDA_CHECK(cudaEventDestroy(copy_event)); + copy_event = nullptr; } for (int i = 0; i < GGML_CUDA_MAX_DEVICES; ++i) { for (int j = 0; j < GGML_CUDA_MAX_STREAMS; ++j) { @@ -3465,6 +3467,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_DS4_INDEXER_MASK: ggml_cuda_op_ds4_indexer_mask(ctx, dst); break; + case GGML_OP_DS4_MOE_COMBINE: + ggml_cuda_op_ds4_moe_combine(ctx, dst); + break; case GGML_OP_GROUP_NORM: ggml_cuda_op_group_norm(ctx, dst); break; @@ -3852,14 +3857,13 @@ static bool ggml_backend_cuda_cpy_tensor_async(ggml_backend_t backend_src, ggml_ #endif // GGML_CUDA_NO_PEER_COPY } + ggml_cuda_set_device(cuda_ctx_src->device); if (!cuda_ctx_src->copy_event) { - ggml_cuda_set_device(cuda_ctx_src->device); - CUDA_CHECK(cudaEventCreateWithFlags(&cuda_ctx_src->copy_event, cudaEventDisableTiming)); - } - { - CUDA_CHECK(cudaEventRecord( - cuda_ctx_src->copy_event, cuda_ctx_src->stream())); + CUDA_CHECK(cudaEventCreateWithFlags( + &cuda_ctx_src->copy_event, cudaEventDisableTiming)); } + CUDA_CHECK(cudaEventRecord( + cuda_ctx_src->copy_event, cuda_ctx_src->stream())); // wait on dst stream for the copy to complete CUDA_CHECK(cudaStreamWaitEvent( @@ -6066,6 +6070,14 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g op->src[1]->type == GGML_TYPE_I32 && ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]); + case GGML_OP_DS4_MOE_COMBINE: + return op->src[0]->type == GGML_TYPE_F32 && + op->src[1]->type == GGML_TYPE_F32 && + op->src[0]->ne[0] % 4 == 0 && + op->src[0]->nb[1] % sizeof(float4) == 0 && + op->src[0]->nb[2] % sizeof(float4) == 0 && + op->nb[1] % sizeof(float4) == 0 && + (op->src[2] == nullptr || (op->src[2]->type == GGML_TYPE_F32 && op->src[2]->nb[1] % sizeof(float4) == 0)); case GGML_OP_MUL_MAT: case GGML_OP_MUL_MAT_GROUPED_SRC: case GGML_OP_MUL_MAT_ID: diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused-combine.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused-combine.cu new file mode 100644 index 000000000..19762791b --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused-combine.cu @@ -0,0 +1,147 @@ +#include "common.cuh" +#include "moe-fused-combine.cuh" + +#include +#include +#include +#include + +// Fused post-MMID route reduction and shared-expert add. The down projection +// remains a separate operation and materializes down_e before this kernel. +// Computes dst[t, d] = (shared_out ? shared_out[t, d] : 0) + sum_{e=0}^{n_used-1} (down_e[t, e, d] * weights[t, e]) +// Uses vectorized float4 128-bit memory transactions and sequential non-FMA FP32 accumulation +// to preserve the legacy route-order reduction. + +static __global__ void moe_fused_combine_shared_kernel_f32( + const float4 * __restrict__ down_e, + const float * __restrict__ weights, + const float4 * __restrict__ shared_out, + float4 * __restrict__ output, + const int n_embd_vec4, + const int n_used, + const int n_tokens, + const size_t down_nb1, + const size_t down_nb2, + const size_t weights_nb1, + const size_t shared_nb1, + const size_t out_nb1) { + + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + const int total = n_embd_vec4 * n_tokens; + if (idx >= total) return; + + const int h4 = idx % n_embd_vec4; + const int t = idx / n_embd_vec4; + + float sum0 = 0.0f; + float sum1 = 0.0f; + float sum2 = 0.0f; + float sum3 = 0.0f; + + for (int e = 0; e < n_used; ++e) { + const float w = weights[e + t * weights_nb1]; + // Expert-major owners encode routes assigned to the other device as + // weight zero and ID -1. Do not read those MMID output lanes: a masked + // lane is semantically zero, and stale NaN/Inf multiplied by zero would + // otherwise poison the reduction. + if (w == 0.0f) { + continue; + } + const float4 v = down_e[h4 + e * down_nb1 + t * down_nb2]; + + const float p0 = __fmul_rn(v.x, w); + const float p1 = __fmul_rn(v.y, w); + const float p2 = __fmul_rn(v.z, w); + const float p3 = __fmul_rn(v.w, w); + + // Start from +0 and add every active route, matching the legacy + // sum_rows reduction even for a first product of -0.0f. + sum0 = __fadd_rn(sum0, p0); + sum1 = __fadd_rn(sum1, p1); + sum2 = __fadd_rn(sum2, p2); + sum3 = __fadd_rn(sum3, p3); + } + + if (shared_out != nullptr) { + const float4 sh = shared_out[h4 + t * shared_nb1]; + sum0 = __fadd_rn(sh.x, sum0); + sum1 = __fadd_rn(sh.y, sum1); + sum2 = __fadd_rn(sh.z, sum2); + sum3 = __fadd_rn(sh.w, sum3); + } + + output[h4 + t * out_nb1] = make_float4(sum0, sum1, sum2, sum3); +} + +void ggml_cuda_op_ds4_moe_combine(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * down_e = dst->src[0]; + const ggml_tensor * weights = dst->src[1]; + const ggml_tensor * shared_out = dst->src[2]; + + GGML_ASSERT(down_e->type == GGML_TYPE_F32); + GGML_ASSERT(weights->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + const int n_embd = (int) down_e->ne[0]; + const int n_used = (int) down_e->ne[1]; + const int n_tokens = (int) down_e->ne[2]; + + static const bool trace_enabled = []() { + const char * raw = std::getenv("DFLASH_MOE_FUSED_COMBINE_TRACE"); + return raw && *raw && std::strcmp(raw, "0") != 0; + }(); + if (trace_enabled) { + static std::atomic launch_count{0}; + int device = -1; + CUDA_CHECK(cudaGetDevice(&device)); + const unsigned long long launch = + launch_count.fetch_add(1, std::memory_order_relaxed) + 1; + std::fprintf(stderr, + "[moe-fused-combine] launch=%llu device=%d tokens=%d routes=%d shared=%d\n", + launch, device, n_tokens, n_used, shared_out != nullptr ? 1 : 0); + } + + GGML_ASSERT(n_embd % 4 == 0 && "n_embd must be a multiple of 4 for float4 vectorization"); + GGML_ASSERT(reinterpret_cast(down_e->data) % 16 == 0 && "down_e->data must be 16-byte aligned"); + GGML_ASSERT(reinterpret_cast(dst->data) % 16 == 0 && "dst->data must be 16-byte aligned"); + GGML_ASSERT(down_e->nb[0] == sizeof(float) && "down_e must be contiguous in dimension 0"); + GGML_ASSERT(down_e->nb[1] % sizeof(float4) == 0 && "down_e->nb[1] must be divisible by sizeof(float4)"); + GGML_ASSERT(down_e->nb[2] % sizeof(float4) == 0 && "down_e->nb[2] must be divisible by sizeof(float4)"); + GGML_ASSERT(dst->nb[0] == sizeof(float) && "dst must be contiguous in dimension 0"); + GGML_ASSERT(dst->nb[1] % sizeof(float4) == 0 && "dst->nb[1] must be divisible by sizeof(float4)"); + if (shared_out != nullptr) { + GGML_ASSERT(reinterpret_cast(shared_out->data) % 16 == 0 && "shared_out->data must be 16-byte aligned"); + GGML_ASSERT(shared_out->nb[0] == sizeof(float) && "shared_out must be contiguous in dimension 0"); + GGML_ASSERT(shared_out->nb[1] % sizeof(float4) == 0 && "shared_out->nb[1] must be divisible by sizeof(float4)"); + } + + const int n_embd_vec4 = n_embd / 4; + const int total_threads = n_embd_vec4 * n_tokens; + + const int block_size = 256; + const int grid_size = (total_threads + block_size - 1) / block_size; + + const size_t down_nb1 = down_e->nb[1] / sizeof(float4); + const size_t down_nb2 = down_e->nb[2] / sizeof(float4); + const size_t weights_nb1 = weights->nb[1] / sizeof(float); + const size_t shared_nb1 = shared_out ? (shared_out->nb[1] / sizeof(float4)) : 0; + const size_t out_nb1 = dst->nb[1] / sizeof(float4); + + cudaStream_t stream = ctx.stream(); + + moe_fused_combine_shared_kernel_f32<<>>( + (const float4 *) down_e->data, + (const float *) weights->data, + shared_out ? (const float4 *) shared_out->data : nullptr, + (float4 *) dst->data, + n_embd_vec4, + n_used, + n_tokens, + down_nb1, + down_nb2, + weights_nb1, + shared_nb1, + out_nb1 + ); + CUDA_CHECK(cudaGetLastError()); +} diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused-combine.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused-combine.cuh new file mode 100644 index 000000000..173c3f23a --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused-combine.cuh @@ -0,0 +1,5 @@ +#pragma once + +#include "common.cuh" + +void ggml_cuda_op_ds4_moe_combine(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index 82e2d9a1c..736ec2add 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -1200,9 +1200,11 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "MUL_MAT_GROUPED_SRC", "PAGED_ATTN", + + "DS4_MOE_COMBINE", }; -static_assert(GGML_OP_COUNT == 105, "GGML_OP_COUNT != 105"); +static_assert(GGML_OP_COUNT == 106, "GGML_OP_COUNT != 106"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1327,9 +1329,11 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "X*grouped(Y)", "paged_attn(q,k,v)", + + "ds4_moe_combine(down,w,shared)", }; -static_assert(GGML_OP_COUNT == 105, "GGML_OP_COUNT != 105"); +static_assert(GGML_OP_COUNT == 106, "GGML_OP_COUNT != 106"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -9194,3 +9198,38 @@ struct ggml_tensor * ggml_ds4_indexer_mask( ggml_set_op_params_i32(result, 0, raw_rows); return result; } + +struct ggml_tensor * ggml_ds4_moe_fused_combine_shared( + struct ggml_context * ctx, + struct ggml_tensor * down_e, + struct ggml_tensor * weights, + struct ggml_tensor * shared_out) { + GGML_ASSERT(down_e != NULL); + GGML_ASSERT(weights != NULL); + GGML_ASSERT(down_e->type == GGML_TYPE_F32); + GGML_ASSERT(weights->type == GGML_TYPE_F32); + GGML_ASSERT(down_e->nb[0] == sizeof(float)); + GGML_ASSERT(down_e->ne[0] % 4 == 0); + GGML_ASSERT(down_e->ne[3] == 1); + GGML_ASSERT(weights->nb[0] == sizeof(float)); + GGML_ASSERT(down_e->ne[1] == weights->ne[0]); + GGML_ASSERT(down_e->ne[2] == weights->ne[1]); + GGML_ASSERT(weights->ne[2] == 1); + GGML_ASSERT(weights->ne[3] == 1); + if (shared_out != NULL) { + GGML_ASSERT(shared_out->type == GGML_TYPE_F32); + GGML_ASSERT(shared_out->nb[0] == sizeof(float)); + GGML_ASSERT(shared_out->ne[0] == down_e->ne[0]); + GGML_ASSERT(shared_out->ne[1] == down_e->ne[2]); + GGML_ASSERT(shared_out->ne[2] == 1); + GGML_ASSERT(shared_out->ne[3] == 1); + } + + struct ggml_tensor * result = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, down_e->ne[0], down_e->ne[2]); + result->op = GGML_OP_DS4_MOE_COMBINE; + result->src[0] = down_e; + result->src[1] = weights; + result->src[2] = shared_out; + return result; +} diff --git a/server/src/common/moe_hybrid_ffn_eval.cpp b/server/src/common/moe_hybrid_ffn_eval.cpp index d5e80cec0..6e024f32b 100644 --- a/server/src/common/moe_hybrid_ffn_eval.cpp +++ b/server/src/common/moe_hybrid_ffn_eval.cpp @@ -2877,18 +2877,31 @@ static bool eval_moe_owner_expert_major_batched( ggml_tensor * down_e = apply_scale2(ctx, ggml_mul_mat_id(ctx, down_tensor, gu, local_ids_tensor), desc.ffn_down_exps_s); - ggml_tensor * weights_3d = ggml_reshape_3d(ctx, owner_weights_tensor, 1, n_used, n_tokens); - ggml_tensor * routed_out = ggml_mul(ctx, down_e, weights_3d); - routed_out = ggml_cont(ctx, ggml_permute(ctx, routed_out, 1, 0, 2, 3)); - routed_out = ggml_sum_rows(ctx, routed_out); - routed_out = ggml_reshape_2d(ctx, routed_out, n_embd, n_tokens); - - ggml_tensor * combined_out = routed_out; + ggml_tensor * shared_out = nullptr; if (has_shared) { - ggml_tensor * shared_out = build_shared_expert_subgraph(ctx, desc, inp, cfg.swiglu_clamp); - if (shared_out) { - combined_out = ggml_add(ctx, combined_out, shared_out); - } + shared_out = build_shared_expert_subgraph(ctx, desc, inp, cfg.swiglu_clamp); + } + + ggml_tensor * combined_out = nullptr; + if (moe_hybrid_graph_policy().fused_combine) { + // The production expert-major MMID path used to materialize the + // weighted route tensor, transpose it, reduce it, and finally add + // the shared expert. Reduce the owner-local routes directly from + // down_e instead. The same operation handles the cold owner with a + // null shared tensor, so both GPU owners avoid the legacy chain. + combined_out = ggml_ds4_moe_fused_combine_shared( + ctx, down_e, owner_weights_tensor, shared_out); + } else { + ggml_tensor * weights_3d = ggml_reshape_3d( + ctx, owner_weights_tensor, 1, n_used, n_tokens); + ggml_tensor * routed_out = ggml_mul(ctx, down_e, weights_3d); + routed_out = ggml_cont( + ctx, ggml_permute(ctx, routed_out, 1, 0, 2, 3)); + routed_out = ggml_sum_rows(ctx, routed_out); + routed_out = ggml_reshape_2d(ctx, routed_out, n_embd, n_tokens); + combined_out = shared_out + ? ggml_add(ctx, routed_out, shared_out) + : routed_out; } ggml_cgraph * gf = ggml_new_graph_custom(ctx, 256, false); diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 4158a0ffd..3471a60b4 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -359,6 +359,15 @@ struct DeepSeek4CachedDecodeAttnGraph { } }; +static bool ds4_moe_fused_combine_enabled() { + static const bool enabled = []() { + const char * val = getenv("DFLASH_MOE_FUSED_COMBINE"); + if (!val) return true; // Default ON in production + return atoi(val) != 0; + }(); + return enabled; +} + struct DeepSeek4CachedLayerAlloc { const ggml_context * owner_ctx = nullptr; ggml_backend_t backend = nullptr; @@ -477,14 +486,18 @@ static bool build_cached_decode_ffn_graph( weights = ggml_scale(out.sg.ctx, weights, w.expert_weight_scale); } - ggml_tensor * weights_3d = ggml_reshape_3d(out.sg.ctx, weights, 1, n_used, n_tokens); - ggml_tensor * routed_out = ggml_mul(out.sg.ctx, down_e, weights_3d); - routed_out = ggml_cont( - out.sg.ctx, ggml_permute(out.sg.ctx, routed_out, 1, 0, 2, 3)); - routed_out = ggml_sum_rows(out.sg.ctx, routed_out); - routed_out = ggml_reshape_2d(out.sg.ctx, routed_out, w.n_embd, n_tokens); + if (ds4_moe_fused_combine_enabled()) { + ffn_out = ggml_ds4_moe_fused_combine_shared(out.sg.ctx, down_e, weights, shared_out); + } else { + ggml_tensor * weights_3d = ggml_reshape_3d(out.sg.ctx, weights, 1, n_used, n_tokens); + ggml_tensor * routed_out = ggml_mul(out.sg.ctx, down_e, weights_3d); + routed_out = ggml_cont( + out.sg.ctx, ggml_permute(out.sg.ctx, routed_out, 1, 0, 2, 3)); + routed_out = ggml_sum_rows(out.sg.ctx, routed_out); + routed_out = ggml_reshape_2d(out.sg.ctx, routed_out, w.n_embd, n_tokens); - ffn_out = ggml_add(out.sg.ctx, shared_out, routed_out); + ffn_out = ggml_add(out.sg.ctx, shared_out, routed_out); + } } else { ffn_out = build_moe_ffn(out.sg.ctx, ffn_normed, w, L, layer_idx, n_tokens); } @@ -3348,11 +3361,16 @@ static ggml_tensor * build_moe_ffn( ggml_tensor * down_e = ggml_mul_mat_id(ctx, L.ffn_down_exps, mid_e, routing.selected); down_e = ggml_reshape_3d(ctx, down_e, n_embd, n_used, n_tokens); - ggml_tensor * weights_3d = ggml_reshape_3d(ctx, routing.weights, 1, n_used, n_tokens); - routed_out = ggml_mul(ctx, down_e, weights_3d); - routed_out = ggml_cont(ctx, ggml_permute(ctx, routed_out, 1, 0, 2, 3)); - routed_out = ggml_sum_rows(ctx, routed_out); - routed_out = ggml_reshape_2d(ctx, routed_out, n_embd, n_tokens); + if (ds4_moe_fused_combine_enabled()) { + return ggml_ds4_moe_fused_combine_shared(ctx, down_e, routing.weights, shared_out); + } else { + ggml_tensor * weights_3d = ggml_reshape_3d(ctx, routing.weights, 1, n_used, n_tokens); + routed_out = ggml_mul(ctx, down_e, weights_3d); + routed_out = ggml_cont(ctx, ggml_permute(ctx, routed_out, 1, 0, 2, 3)); + routed_out = ggml_sum_rows(ctx, routed_out); + routed_out = ggml_reshape_2d(ctx, routed_out, n_embd, n_tokens); + return ggml_add(ctx, shared_out, routed_out); + } } return ggml_add(ctx, shared_out, routed_out); @@ -5083,6 +5101,10 @@ static ggml_tensor * ds4_build_hash_routed_ffn( weights = ggml_scale(ctx, weights, w.expert_weight_scale); } + if (ds4_moe_fused_combine_enabled()) { + return ggml_ds4_moe_fused_combine_shared(ctx, down_e, weights, shared_out); + } + ggml_tensor * weights_3d = ggml_reshape_3d( ctx, weights, 1, n_used, n_tokens); ggml_tensor * routed_out = ggml_mul(ctx, down_e, weights_3d); diff --git a/server/test/test_ds4_moe_combine_cuda.cpp b/server/test/test_ds4_moe_combine_cuda.cpp new file mode 100644 index 000000000..4c7f2d084 --- /dev/null +++ b/server/test/test_ds4_moe_combine_cuda.cpp @@ -0,0 +1,399 @@ +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml-cpu.h" +#include "ggml-cuda.h" +#include "ggml.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr int kEmbeddings = 260; + +struct CombineCase { + int top_k; + int tokens; + bool include_shared; +}; + +struct Inputs { + std::vector down; + std::vector weights; + std::vector shared; + std::vector expected; + int zero_weight_routes = 0; +}; + +Inputs make_inputs(const CombineCase & test, bool poison_masked = true) { + Inputs result; + result.down.resize((size_t) kEmbeddings * test.top_k * test.tokens); + result.weights.resize((size_t) test.top_k * test.tokens); + if (test.include_shared) { + result.shared.resize((size_t) kEmbeddings * test.tokens); + } + result.expected.resize((size_t) kEmbeddings * test.tokens); + + for (int token = 0; token < test.tokens; ++token) { + for (int expert = 0; expert < test.top_k; ++expert) { + const size_t weight_offset = (size_t) token * test.top_k + expert; + const bool masked = (token * 3 + expert) % 4 == 0; + result.weights[weight_offset] = masked + ? 0.0f + : 0.125f * (float) (expert + 1); + result.zero_weight_routes += masked ? 1 : 0; + + for (int embedding = 0; embedding < kEmbeddings; ++embedding) { + const size_t down_offset = + (size_t) token * test.top_k * kEmbeddings + + (size_t) expert * kEmbeddings + embedding; + result.down[down_offset] = masked + ? (poison_masked + ? std::numeric_limits::quiet_NaN() + : 0.03125f * (float) ((embedding % 7) + 1)) + : 0.03125f * (float) ((embedding % 11) - 5) * + (float) (expert + 1) + + 0.015625f * (float) token; + } + } + } + + for (int token = 0; token < test.tokens; ++token) { + for (int embedding = 0; embedding < kEmbeddings; ++embedding) { + const size_t output_offset = (size_t) token * kEmbeddings + embedding; + const float shared_value = test.include_shared + ? 0.0625f * (float) ((embedding + token) % 13 - 6) + : 0.0f; + if (test.include_shared) { + result.shared[output_offset] = shared_value; + } + float sum = 0.0f; + for (int expert = 0; expert < test.top_k; ++expert) { + const size_t route = (size_t) token * test.top_k + expert; + if (result.weights[route] == 0.0f) { + continue; + } + const size_t down_offset = + (size_t) token * test.top_k * kEmbeddings + + (size_t) expert * kEmbeddings + embedding; + const float product = result.down[down_offset] * result.weights[route]; + sum += product; + } + sum += shared_value; + result.expected[output_offset] = sum; + } + } + + return result; +} + +bool run_backend( + ggml_backend_t backend, + const CombineCase & test, + const Inputs & inputs, + bool fused, + std::vector * output) { + ggml_init_params params{}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + if (!ctx) { + std::fprintf(stderr, "[ds4-moe-combine] ggml_init failed\n"); + return false; + } + + ggml_tensor * down = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, kEmbeddings, test.top_k, test.tokens); + ggml_tensor * weights = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, test.top_k, test.tokens); + ggml_tensor * shared = test.include_shared + ? ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kEmbeddings, test.tokens) + : nullptr; + ggml_set_input(down); + ggml_set_input(weights); + if (shared) { + ggml_set_input(shared); + } + ggml_tensor * combined = nullptr; + if (fused) { + combined = ggml_ds4_moe_fused_combine_shared( + ctx, down, weights, shared); + } else { + ggml_tensor * weights_3d = ggml_reshape_3d( + ctx, weights, 1, test.top_k, test.tokens); + ggml_tensor * routed = ggml_mul(ctx, down, weights_3d); + routed = ggml_cont(ctx, ggml_permute(ctx, routed, 1, 0, 2, 3)); + routed = ggml_sum_rows(ctx, routed); + routed = ggml_reshape_2d(ctx, routed, kEmbeddings, test.tokens); + combined = shared ? ggml_add(ctx, routed, shared) : routed; + } + ggml_set_output(combined); + + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, combined); + ggml_gallocr_t alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + if (!alloc || !ggml_gallocr_alloc_graph(alloc, graph)) { + std::fprintf(stderr, "[ds4-moe-combine] graph allocation failed\n"); + if (alloc) { + ggml_gallocr_free(alloc); + } + ggml_free(ctx); + return false; + } + + ggml_backend_tensor_set(down, inputs.down.data(), 0, ggml_nbytes(down)); + ggml_backend_tensor_set(weights, inputs.weights.data(), 0, ggml_nbytes(weights)); + if (shared) { + ggml_backend_tensor_set(shared, inputs.shared.data(), 0, ggml_nbytes(shared)); + } + + const bool ok = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + if (ok) { + ggml_backend_synchronize(backend); + output->resize(inputs.expected.size()); + ggml_backend_tensor_get( + combined, output->data(), 0, output->size() * sizeof(float)); + } else { + std::fprintf(stderr, "[ds4-moe-combine] graph compute failed\n"); + } + + ggml_gallocr_free(alloc); + ggml_free(ctx); + return ok; +} + +bool equal_bytes(const std::vector & expected, + const std::vector & actual, + const char * label, + const CombineCase & test) { + if (expected.size() == actual.size() && + std::memcmp(expected.data(), actual.data(), expected.size() * sizeof(float)) == 0) { + return true; + } + + size_t first = 0; + while (first < expected.size() && first < actual.size() && + std::memcmp(&expected[first], &actual[first], sizeof(float)) == 0) { + ++first; + } + std::fprintf(stderr, + "[ds4-moe-combine] %s mismatch top_k=%d tokens=%d shared=%d index=%zu " + "expected=%g actual=%g\n", + label, test.top_k, test.tokens, test.include_shared ? 1 : 0, first, + first < expected.size() ? expected[first] : 0.0f, + first < actual.size() ? actual[first] : 0.0f); + return false; +} + +bool finite_output(const std::vector & output, const CombineCase & test) { + for (size_t i = 0; i < output.size(); ++i) { + if (!std::isfinite(output[i])) { + std::fprintf(stderr, + "[ds4-moe-combine] poisoned masked route reached output " + "top_k=%d tokens=%d shared=%d index=%zu value=%g\n", + test.top_k, test.tokens, test.include_shared ? 1 : 0, + i, output[i]); + return false; + } + } + return true; +} + +Inputs make_signed_zero_inputs(const CombineCase & test) { + Inputs result; + result.down.assign((size_t) kEmbeddings * test.top_k * test.tokens, 1.0f); + result.weights.assign((size_t) test.top_k * test.tokens, 0.0f); + result.expected.assign((size_t) kEmbeddings * test.tokens, 0.0f); + for (int token = 0; token < test.tokens; ++token) { + result.weights[(size_t) token * test.top_k] = 1.0f; + for (int embedding = 0; embedding < kEmbeddings; ++embedding) { + result.down[(size_t) token * test.top_k * kEmbeddings + embedding] = -0.0f; + } + } + return result; +} + +bool benchmark_path(ggml_backend_t backend, int tokens, bool fused, + double * median_ms, double * mad_ms) { + constexpr int n_embd = 4096; + constexpr int top_k = 6; + constexpr int warmups = 2; + constexpr int samples = 7; + + ggml_init_params params{}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + if (!ctx) { + return false; + } + + ggml_tensor * down = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, n_embd, top_k, tokens); + ggml_tensor * weights = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, top_k, tokens); + ggml_tensor * shared = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, n_embd, tokens); + ggml_set_input(down); + ggml_set_input(weights); + ggml_set_input(shared); + + ggml_tensor * output = nullptr; + if (fused) { + output = ggml_ds4_moe_fused_combine_shared( + ctx, down, weights, shared); + } else { + ggml_tensor * weights_3d = ggml_reshape_3d( + ctx, weights, 1, top_k, tokens); + ggml_tensor * routed = ggml_mul(ctx, down, weights_3d); + routed = ggml_cont(ctx, ggml_permute(ctx, routed, 1, 0, 2, 3)); + routed = ggml_sum_rows(ctx, routed); + routed = ggml_reshape_2d(ctx, routed, n_embd, tokens); + output = ggml_add(ctx, routed, shared); + } + ggml_set_output(output); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); + ggml_gallocr_t alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + if (!alloc || !ggml_gallocr_alloc_graph(alloc, graph)) { + if (alloc) { + ggml_gallocr_free(alloc); + } + ggml_free(ctx); + return false; + } + + std::vector down_h((size_t) n_embd * top_k * tokens, 0.03125f); + std::vector weights_h((size_t) top_k * tokens, 0.125f); + std::vector shared_h((size_t) n_embd * tokens, -0.0625f); + ggml_backend_tensor_set(down, down_h.data(), 0, ggml_nbytes(down)); + ggml_backend_tensor_set(weights, weights_h.data(), 0, ggml_nbytes(weights)); + ggml_backend_tensor_set(shared, shared_h.data(), 0, ggml_nbytes(shared)); + + std::vector timings; + timings.reserve(samples); + bool ok = true; + for (int i = 0; i < warmups + samples; ++i) { + const auto start = std::chrono::steady_clock::now(); + ok = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS && ok; + ggml_backend_synchronize(backend); + const auto end = std::chrono::steady_clock::now(); + if (i >= warmups) { + timings.push_back(std::chrono::duration( + end - start).count()); + } + } + + std::sort(timings.begin(), timings.end()); + *median_ms = timings[timings.size() / 2]; + std::vector deviations; + deviations.reserve(timings.size()); + for (double value : timings) { + deviations.push_back(std::fabs(value - *median_ms)); + } + std::sort(deviations.begin(), deviations.end()); + *mad_ms = deviations[deviations.size() / 2]; + + ggml_gallocr_free(alloc); + ggml_free(ctx); + return ok; +} + +} // namespace + +int main(int argc, char ** argv) { + if (ggml_backend_cuda_get_device_count() <= 0) { + std::puts("[ds4-moe-combine] SKIP: HIP device unavailable"); + return 77; + } + + ggml_backend_t cpu = ggml_backend_cpu_init(); + ggml_backend_t hip = ggml_backend_cuda_init(0); + if (!cpu || !hip) { + std::fprintf(stderr, "[ds4-moe-combine] backend initialization failed\n"); + if (hip) { + ggml_backend_free(hip); + } + if (cpu) { + ggml_backend_free(cpu); + } + return 1; + } + ggml_backend_cpu_set_n_threads(cpu, 1); + + const CombineCase cases[] = { + {4, 1, false}, {4, 3, true}, {4, 33, false}, {4, 401, true}, + {6, 1, true}, {6, 3, false}, {6, 33, true}, {6, 401, false}, + }; + + bool ok = true; + for (const CombineCase & test : cases) { + const Inputs inputs = make_inputs(test); + std::vector cpu_output; + std::vector hip_output; + ok = run_backend(cpu, test, inputs, true, &cpu_output) && ok; + ok = run_backend(hip, test, inputs, true, &hip_output) && ok; + ok = finite_output(cpu_output, test) && finite_output(hip_output, test) && ok; + ok = equal_bytes(inputs.expected, cpu_output, "CPU reference", test) && ok; + ok = equal_bytes(cpu_output, hip_output, "HIP exact parity", test) && ok; + + const Inputs finite_inputs = make_inputs(test, false); + std::vector legacy_cpu_output; + std::vector legacy_hip_output; + std::vector fused_cpu_output; + std::vector fused_hip_output; + ok = run_backend(cpu, test, finite_inputs, false, &legacy_cpu_output) && ok; + ok = run_backend(hip, test, finite_inputs, false, &legacy_hip_output) && ok; + ok = run_backend(cpu, test, finite_inputs, true, &fused_cpu_output) && ok; + ok = run_backend(hip, test, finite_inputs, true, &fused_hip_output) && ok; + ok = equal_bytes(legacy_cpu_output, fused_cpu_output, + "CPU legacy differential", test) && ok; + ok = equal_bytes(legacy_hip_output, fused_hip_output, + "HIP legacy differential", test) && ok; + std::printf("[ds4-moe-combine] top_k=%d tokens=%d shared=%d zero_routes=%d %s\n", + test.top_k, test.tokens, test.include_shared ? 1 : 0, + inputs.zero_weight_routes, ok ? "PASS" : "FAIL"); + } + + const CombineCase signed_zero_case{6, 3, false}; + const Inputs signed_zero_inputs = make_signed_zero_inputs(signed_zero_case); + std::vector signed_zero_legacy; + std::vector signed_zero_fused; + ok = run_backend(hip, signed_zero_case, signed_zero_inputs, false, + &signed_zero_legacy) && ok; + ok = run_backend(hip, signed_zero_case, signed_zero_inputs, true, + &signed_zero_fused) && ok; + ok = equal_bytes(signed_zero_legacy, signed_zero_fused, + "HIP signed-zero legacy differential", signed_zero_case) && ok; + ok = equal_bytes(signed_zero_inputs.expected, signed_zero_fused, + "HIP signed-zero +0 result", signed_zero_case) && ok; + + if (argc == 2 && std::strcmp(argv[1], "--benchmark") == 0) { + for (int tokens : {401, 2048}) { + double legacy_ms = 0.0; + double legacy_mad = 0.0; + double fused_ms = 0.0; + double fused_mad = 0.0; + ok = benchmark_path( + hip, tokens, false, &legacy_ms, &legacy_mad) && ok; + ok = benchmark_path( + hip, tokens, true, &fused_ms, &fused_mad) && ok; + std::printf( + "[ds4-moe-combine-bench] tokens=%d legacy_ms=%.6f " + "legacy_mad=%.6f fused_ms=%.6f fused_mad=%.6f speedup=%.6fx\n", + tokens, legacy_ms, legacy_mad, fused_ms, fused_mad, + fused_ms > 0.0 ? legacy_ms / fused_ms : 0.0); + } + } + + ggml_backend_free(hip); + ggml_backend_free(cpu); + return ok ? 0 : 1; +} From 3f958e6c5a8729ace3fd58f17265ae18bfb9f303 Mon Sep 17 00:00:00 2001 From: cheese-cakee Date: Thu, 27 Aug 2026 15:32:51 +0530 Subject: [PATCH 13/16] fix(ds4): harden fused MoE combine backend contracts --- .../llama.cpp/ggml/src/ggml-backend-meta.cpp | 38 +++++ .../llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu | 1 + server/test/test_ds4_moe_combine_cuda.cpp | 133 +++++++++++++++--- 3 files changed, 156 insertions(+), 16 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp b/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp index a082f7565..14bd289e2 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp @@ -873,6 +873,41 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( return {GGML_BACKEND_SPLIT_AXIS_0, {0}, 1, {1}}; }; + auto handle_ds4_moe_combine = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + GGML_ASSERT(tensor->src[2] == nullptr || + src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + return src_ss[0]; + } + + // Splitting embeddings is safe when the optional shared branch uses + // the identical embedding partition. Route weights remain mirrored. + if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) { + GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + GGML_ASSERT(tensor->src[2] == nullptr || + split_states_equal(src_ss[0], src_ss[2])); + return src_ss[0]; + } + + // Splitting experts partitions the reduced dimension. Each device + // produces a partial sum, so a following meta-backend synchronization + // must reduce those sums. A shared result cannot be added locally here + // because it would then be counted once per device. + if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1) { + GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_0); + ggml_backend_meta_split_state weights_ss = src_ss[1]; + weights_ss.axis = GGML_BACKEND_SPLIT_AXIS_1; + GGML_ASSERT(split_states_equal(src_ss[0], weights_ss)); + GGML_ASSERT(tensor->src[2] == nullptr); + return {assume_sync ? GGML_BACKEND_SPLIT_AXIS_MIRRORED : + GGML_BACKEND_SPLIT_AXIS_PARTIAL, + {0}, 1, {1}}; + } + + GGML_ABORT("unsupported DS4 MoE combine split"); + }; + auto calculate_split_state = [&]() -> ggml_backend_meta_split_state { if (ggml_nelements(tensor) == 0) { return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1, {1}}; @@ -1112,6 +1147,9 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( // neither may run on unreduced dot-product shards. split_state = handle_mirrored(src_ss); } break; + case GGML_OP_DS4_MOE_COMBINE: { + split_state = handle_ds4_moe_combine(src_ss); + } break; case GGML_OP_UNARY: { split_state = handle_generic(src_ss, /*scalar_only =*/ false); } break; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu index a7c8c0a66..f403981de 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu @@ -6076,6 +6076,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g op->src[0]->ne[0] % 4 == 0 && op->src[0]->nb[1] % sizeof(float4) == 0 && op->src[0]->nb[2] % sizeof(float4) == 0 && + op->src[1]->nb[1] % sizeof(float) == 0 && op->nb[1] % sizeof(float4) == 0 && (op->src[2] == nullptr || (op->src[2]->type == GGML_TYPE_F32 && op->src[2]->nb[1] % sizeof(float4) == 0)); case GGML_OP_MUL_MAT: diff --git a/server/test/test_ds4_moe_combine_cuda.cpp b/server/test/test_ds4_moe_combine_cuda.cpp index 4c7f2d084..5d4b1c3d2 100644 --- a/server/test/test_ds4_moe_combine_cuda.cpp +++ b/server/test/test_ds4_moe_combine_cuda.cpp @@ -30,6 +30,11 @@ struct Inputs { int zero_weight_routes = 0; }; +ggml_backend_meta_split_state mirrored_split_state( + const ggml_tensor *, void *) { + return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, 1, {1}}; +} + Inputs make_inputs(const CombineCase & test, bool poison_masked = true) { Inputs result; result.down.resize((size_t) kEmbeddings * test.top_k * test.tokens); @@ -168,20 +173,100 @@ bool run_backend( return ok; } -bool equal_bytes(const std::vector & expected, - const std::vector & actual, - const char * label, - const CombineCase & test) { - if (expected.size() == actual.size() && - std::memcmp(expected.data(), actual.data(), expected.size() * sizeof(float)) == 0) { - return true; +bool meta_allocation_test(ggml_backend_t simple_backend, bool include_shared) { + ggml_backend_dev_t simple_dev = ggml_backend_get_device(simple_backend); + ggml_backend_dev_t meta_dev = ggml_backend_meta_device( + &simple_dev, 1, mirrored_split_state, nullptr); + ggml_backend_t meta = meta_dev ? ggml_backend_dev_init(meta_dev, nullptr) : nullptr; + if (!meta) { + std::fprintf(stderr, "[ds4-moe-combine] meta backend initialization failed\n"); + return false; + } + + ggml_init_params params{}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + ggml_tensor * down = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, kEmbeddings, 6, 3); + ggml_tensor * weights = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 6, 3); + ggml_tensor * shared = include_shared + ? ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kEmbeddings, 3) + : nullptr; + ggml_set_input(down); + ggml_set_input(weights); + if (shared) { + ggml_set_input(shared); + } + ggml_tensor * combined = ggml_ds4_moe_fused_combine_shared( + ctx, down, weights, shared); + ggml_set_output(combined); + + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, combined); + ggml_gallocr_t alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(meta)); + const bool ok = alloc && ggml_gallocr_alloc_graph(alloc, graph); + if (!ok) { + std::fprintf(stderr, + "[ds4-moe-combine] meta graph allocation failed shared=%d\n", + include_shared ? 1 : 0); + } + + if (alloc) { + ggml_gallocr_free(alloc); } + ggml_free(ctx); + ggml_backend_free(meta); + return ok; +} + +bool rejects_unaligned_weight_stride(ggml_backend_t hip) { + ggml_init_params params{}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + constexpr int top_k = 4; + constexpr int tokens = 3; + ggml_tensor * down = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, kEmbeddings, top_k, tokens); + ggml_tensor * weights_storage = ggml_new_tensor_1d( + ctx, GGML_TYPE_F32, 16); + ggml_tensor * weights = ggml_view_2d( + ctx, weights_storage, top_k, tokens, + top_k * sizeof(float) + 2, 0); + ggml_tensor * combined = ggml_ds4_moe_fused_combine_shared( + ctx, down, weights, nullptr); + + const bool rejected = !ggml_backend_dev_supports_op( + ggml_backend_get_device(hip), combined); + if (!rejected) { + std::fprintf(stderr, + "[ds4-moe-combine] HIP accepted an unaligned weights row stride\n"); + } + ggml_free(ctx); + return rejected; +} + +bool equal_floats(const std::vector & expected, + const std::vector & actual, + const char * label, + const CombineCase & test) { + constexpr float abs_tol = 1.0e-6f; + constexpr float rel_tol = 1.0e-6f; size_t first = 0; - while (first < expected.size() && first < actual.size() && - std::memcmp(&expected[first], &actual[first], sizeof(float)) == 0) { + while (first < expected.size() && first < actual.size()) { + const float lhs = expected[first]; + const float rhs = actual[first]; + const float tolerance = abs_tol + rel_tol * std::max(std::fabs(lhs), std::fabs(rhs)); + if (!std::isfinite(lhs) || !std::isfinite(rhs) || std::fabs(lhs - rhs) > tolerance) { + break; + } ++first; } + if (first == expected.size() && first == actual.size()) { + return true; + } std::fprintf(stderr, "[ds4-moe-combine] %s mismatch top_k=%d tokens=%d shared=%d index=%zu " "expected=%g actual=%g\n", @@ -191,6 +276,20 @@ bool equal_bytes(const std::vector & expected, return false; } +bool equal_bytes(const std::vector & expected, + const std::vector & actual, + const char * label, + const CombineCase & test) { + if (expected.size() == actual.size() && + std::memcmp(expected.data(), actual.data(), expected.size() * sizeof(float)) == 0) { + return true; + } + std::fprintf(stderr, + "[ds4-moe-combine] %s bit mismatch top_k=%d tokens=%d shared=%d\n", + label, test.top_k, test.tokens, test.include_shared ? 1 : 0); + return false; +} + bool finite_output(const std::vector & output, const CombineCase & test) { for (size_t i = 0; i < output.size(); ++i) { if (!std::isfinite(output[i])) { @@ -333,7 +432,9 @@ int main(int argc, char ** argv) { {6, 1, true}, {6, 3, false}, {6, 33, true}, {6, 401, false}, }; - bool ok = true; + bool ok = meta_allocation_test(hip, false) && + meta_allocation_test(hip, true) && + rejects_unaligned_weight_stride(hip); for (const CombineCase & test : cases) { const Inputs inputs = make_inputs(test); std::vector cpu_output; @@ -341,8 +442,8 @@ int main(int argc, char ** argv) { ok = run_backend(cpu, test, inputs, true, &cpu_output) && ok; ok = run_backend(hip, test, inputs, true, &hip_output) && ok; ok = finite_output(cpu_output, test) && finite_output(hip_output, test) && ok; - ok = equal_bytes(inputs.expected, cpu_output, "CPU reference", test) && ok; - ok = equal_bytes(cpu_output, hip_output, "HIP exact parity", test) && ok; + ok = equal_floats(inputs.expected, cpu_output, "CPU reference", test) && ok; + ok = equal_floats(cpu_output, hip_output, "HIP parity", test) && ok; const Inputs finite_inputs = make_inputs(test, false); std::vector legacy_cpu_output; @@ -353,10 +454,10 @@ int main(int argc, char ** argv) { ok = run_backend(hip, test, finite_inputs, false, &legacy_hip_output) && ok; ok = run_backend(cpu, test, finite_inputs, true, &fused_cpu_output) && ok; ok = run_backend(hip, test, finite_inputs, true, &fused_hip_output) && ok; - ok = equal_bytes(legacy_cpu_output, fused_cpu_output, - "CPU legacy differential", test) && ok; - ok = equal_bytes(legacy_hip_output, fused_hip_output, - "HIP legacy differential", test) && ok; + ok = equal_floats(legacy_cpu_output, fused_cpu_output, + "CPU legacy differential", test) && ok; + ok = equal_floats(legacy_hip_output, fused_hip_output, + "HIP legacy differential", test) && ok; std::printf("[ds4-moe-combine] top_k=%d tokens=%d shared=%d zero_routes=%d %s\n", test.top_k, test.tokens, test.include_shared ? 1 : 0, inputs.zero_weight_routes, ok ? "PASS" : "FAIL"); From 70dcd1e06b45c9315969bed491d50ad429f63ef6 Mon Sep 17 00:00:00 2001 From: cheese-cakee Date: Thu, 27 Aug 2026 15:35:27 +0530 Subject: [PATCH 14/16] fix(ggml): require aligned meta combine partitions --- server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp b/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp index 14bd289e2..731541936 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp @@ -887,6 +887,13 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); GGML_ASSERT(tensor->src[2] == nullptr || split_states_equal(src_ss[0], src_ss[2])); + // Each local embedding slice must preserve the float4 layout + // required by the GPU combine kernel. + for (size_t s = 0; s < src_ss[0].n_segments; ++s) { + for (size_t j = 0; j < n_bufs; ++j) { + GGML_ASSERT(src_ss[0].ne[s*n_bufs + j] % 4 == 0); + } + } return src_ss[0]; } From cf43e667c73e2e08825066214500ea98f2673bbc Mon Sep 17 00:00:00 2001 From: cheese-cakee Date: Thu, 27 Aug 2026 15:48:54 +0530 Subject: [PATCH 15/16] fix(dflash): require matching MoE meta partition layouts --- server/CMakeLists.txt | 6 ++ .../ggml/src/ggml-backend-meta-impl.h | 28 +++++ .../llama.cpp/ggml/src/ggml-backend-meta.cpp | 8 +- server/test/test_ggml_meta_split_layout.cpp | 102 ++++++++++++++++++ 4 files changed, 140 insertions(+), 4 deletions(-) create mode 100644 server/deps/llama.cpp/ggml/src/ggml-backend-meta-impl.h create mode 100644 server/test/test_ggml_meta_split_layout.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 4905e4c8c..28fa0ae13 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -1635,6 +1635,12 @@ if(DFLASH27B_TESTS) endif() unset(_client_timeout_test) + add_executable(test_ggml_meta_split_layout test/test_ggml_meta_split_layout.cpp) + target_include_directories(test_ggml_meta_split_layout PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/src) + list(APPEND _raw_unit_test_targets test_ggml_meta_split_layout) + # CPU-only contract test for the fail-closed layer-split tree boundary. add_executable(test_qwen35_split_tree_guard test/test_qwen35_split_tree_guard.cpp) diff --git a/server/deps/llama.cpp/ggml/src/ggml-backend-meta-impl.h b/server/deps/llama.cpp/ggml/src/ggml-backend-meta-impl.h new file mode 100644 index 000000000..eb32fff9a --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-backend-meta-impl.h @@ -0,0 +1,28 @@ +#pragma once + +#include "ggml-backend.h" + +// Axes are checked by the caller: paired tensors can store the same expert +// partition on different axes. Require identical layout representations, not +// just equal per-device totals, so local element ordering cannot differ. +inline bool ggml_backend_meta_split_layout_equal( + const ggml_backend_meta_split_state & a, + const ggml_backend_meta_split_state & b, + size_t n_devices) { + if (n_devices == 0 || n_devices > GGML_BACKEND_META_MAX_DEVICES || + a.n_segments == 0 || a.n_segments > sizeof(a.nr) / sizeof(a.nr[0]) || + a.n_segments != b.n_segments) { + return false; + } + for (size_t s = 0; s < a.n_segments; ++s) { + if (a.nr[s] != b.nr[s]) { + return false; + } + for (size_t j = 0; j < n_devices; ++j) { + if (a.ne[s*n_devices + j] != b.ne[s*n_devices + j]) { + return false; + } + } + } + return true; +} diff --git a/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp b/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp index 731541936..a99f62996 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp @@ -2,6 +2,7 @@ #include "ggml-impl.h" #include "ggml-backend.h" #include "ggml-backend-impl.h" +#include "ggml-backend-meta-impl.h" #include "ggml-alloc.h" #include "ggml-cpp.h" @@ -886,7 +887,8 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) { GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); GGML_ASSERT(tensor->src[2] == nullptr || - split_states_equal(src_ss[0], src_ss[2])); + (src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_0 && + ggml_backend_meta_split_layout_equal(src_ss[0], src_ss[2], n_bufs))); // Each local embedding slice must preserve the float4 layout // required by the GPU combine kernel. for (size_t s = 0; s < src_ss[0].n_segments; ++s) { @@ -903,9 +905,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( // because it would then be counted once per device. if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1) { GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_0); - ggml_backend_meta_split_state weights_ss = src_ss[1]; - weights_ss.axis = GGML_BACKEND_SPLIT_AXIS_1; - GGML_ASSERT(split_states_equal(src_ss[0], weights_ss)); + GGML_ASSERT(ggml_backend_meta_split_layout_equal(src_ss[0], src_ss[1], n_bufs)); GGML_ASSERT(tensor->src[2] == nullptr); return {assume_sync ? GGML_BACKEND_SPLIT_AXIS_MIRRORED : GGML_BACKEND_SPLIT_AXIS_PARTIAL, diff --git a/server/test/test_ggml_meta_split_layout.cpp b/server/test/test_ggml_meta_split_layout.cpp new file mode 100644 index 000000000..63b74d2e2 --- /dev/null +++ b/server/test/test_ggml_meta_split_layout.cpp @@ -0,0 +1,102 @@ +#include "ggml-backend-meta-impl.h" + +#include + +int main() { + int checks = 0; + int failures = 0; + const auto check = [&](bool ok, const char * label) { + ++checks; + if (!ok) { + ++failures; + std::fprintf(stderr, "[meta-split-layout] FAIL: %s\n", label); + } + }; + + const ggml_backend_meta_split_state single = { + GGML_BACKEND_SPLIT_AXIS_1, {2, 2}, 1, {1}}; + check(ggml_backend_meta_split_layout_equal(single, single, 2), "single segment"); + + const ggml_backend_meta_split_state down = { + GGML_BACKEND_SPLIT_AXIS_1, {1, 1, 1, 1}, 2, {1, 1}}; + auto weights = down; + weights.axis = GGML_BACKEND_SPLIT_AXIS_0; + check(ggml_backend_meta_split_layout_equal(down, weights, 2), + "matching expert partition on different tensor axes"); + + // Both states assign two experts per device, but down assigns [0,2]/[1,3] + // while weights assigns [0,1]/[2,3]. Totals alone incorrectly accept this. + const ggml_backend_meta_split_state wrong_weights = { + GGML_BACKEND_SPLIT_AXIS_0, {2, 0, 0, 2}, 2, {1, 1}}; + check(!ggml_backend_meta_split_layout_equal(down, wrong_weights, 2), + "equal totals with different expert identities"); + + const ggml_backend_meta_split_state nonempty = { + GGML_BACKEND_SPLIT_AXIS_1, {8, 8, 8, 8}, 2, {1, 1}}; + const ggml_backend_meta_split_state wrong_nonempty = { + GGML_BACKEND_SPLIT_AXIS_0, {12, 4, 4, 12}, 2, {1, 1}}; + check(!ggml_backend_meta_split_layout_equal(nonempty, wrong_nonempty, 2), + "equal totals with different nonempty expert segments"); + + const ggml_backend_meta_split_state repeated = { + GGML_BACKEND_SPLIT_AXIS_1, {1, 1, 1, 1}, 2, {2, 1}}; + check(ggml_backend_meta_split_layout_equal(repeated, repeated, 2), + "matching repeated multi-segment layout"); + auto wrong_repeats = repeated; + wrong_repeats.nr[0] = 1; + wrong_repeats.nr[1] = 2; + check(!ggml_backend_meta_split_layout_equal(repeated, wrong_repeats, 2), + "conservatively reject different repeat encodings even when equivalent"); + + const ggml_backend_meta_split_state varied_repeats = { + GGML_BACKEND_SPLIT_AXIS_1, {1, 1, 2, 2, 1, 1}, 3, {2, 1, 1}}; + auto reordered_repeats = varied_repeats; + reordered_repeats.nr[0] = 1; + reordered_repeats.nr[2] = 2; + check(!ggml_backend_meta_split_layout_equal(varied_repeats, reordered_repeats, 2), + "equal totals but different repeated segment ordering"); + + check(!ggml_backend_meta_split_layout_equal(down, single, 2), + "different segment counts with equal totals"); + auto padded = down; + padded.ne[4] = 123; + padded.nr[2] = 456; + check(ggml_backend_meta_split_layout_equal(down, padded, 2), + "inactive storage is ignored"); + + const ggml_backend_meta_split_state embedding = { + GGML_BACKEND_SPLIT_AXIS_0, {4, 4, 4, 4}, 2, {1, 1}}; + const ggml_backend_meta_split_state wrong_shared = { + GGML_BACKEND_SPLIT_AXIS_0, {8, 0, 0, 8}, 2, {1, 1}}; + check(ggml_backend_meta_split_layout_equal(embedding, embedding, 2), + "matching shared embedding partition"); + check(!ggml_backend_meta_split_layout_equal(embedding, wrong_shared, 2), + "equal totals with different shared embedding identities"); + + auto invalid = down; + invalid.n_segments = 0; + check(!ggml_backend_meta_split_layout_equal(invalid, invalid, 2), "empty layout rejected"); + invalid.n_segments = 17; + check(!ggml_backend_meta_split_layout_equal(invalid, invalid, 2), "oversized layout rejected"); + check(!ggml_backend_meta_split_layout_equal(down, down, 0), "zero devices rejected"); + check(!ggml_backend_meta_split_layout_equal(down, down, GGML_BACKEND_META_MAX_DEVICES + 1), + "too many devices rejected"); + + auto maximum = down; + maximum.n_segments = 16; + for (auto & count : maximum.ne) { + count = 1; + } + for (auto & repeat : maximum.nr) { + repeat = 1; + } + check(ggml_backend_meta_split_layout_equal(maximum, maximum, GGML_BACKEND_META_MAX_DEVICES), + "maximum supported layout"); + auto wrong_last = maximum; + ++wrong_last.ne[16 * GGML_BACKEND_META_MAX_DEVICES - 1]; + check(!ggml_backend_meta_split_layout_equal(maximum, wrong_last, GGML_BACKEND_META_MAX_DEVICES), + "last segment and device participate in comparison"); + + std::printf("[meta-split-layout] %d checks, %d failures\n", checks, failures); + return failures == 0 ? 0 : 1; +} From 6edb98f4a20ff6a2504db8c7d75fca823953cfbe Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Thu, 3 Sep 2026 15:29:14 +0200 Subject: [PATCH 16/16] perf(ds4): make fused MoE combine structural --- .../ggml/src/ggml-cuda/moe-fused-combine.cu | 20 ------- server/src/deepseek4/deepseek4_graph.cpp | 57 ++----------------- 2 files changed, 6 insertions(+), 71 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused-combine.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused-combine.cu index 19762791b..bce1fdf0a 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused-combine.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused-combine.cu @@ -1,11 +1,6 @@ #include "common.cuh" #include "moe-fused-combine.cuh" -#include -#include -#include -#include - // Fused post-MMID route reduction and shared-expert add. The down projection // remains a separate operation and materializes down_e before this kernel. // Computes dst[t, d] = (shared_out ? shared_out[t, d] : 0) + sum_{e=0}^{n_used-1} (down_e[t, e, d] * weights[t, e]) @@ -86,21 +81,6 @@ void ggml_cuda_op_ds4_moe_combine(ggml_backend_cuda_context & ctx, ggml_tensor * const int n_used = (int) down_e->ne[1]; const int n_tokens = (int) down_e->ne[2]; - static const bool trace_enabled = []() { - const char * raw = std::getenv("DFLASH_MOE_FUSED_COMBINE_TRACE"); - return raw && *raw && std::strcmp(raw, "0") != 0; - }(); - if (trace_enabled) { - static std::atomic launch_count{0}; - int device = -1; - CUDA_CHECK(cudaGetDevice(&device)); - const unsigned long long launch = - launch_count.fetch_add(1, std::memory_order_relaxed) + 1; - std::fprintf(stderr, - "[moe-fused-combine] launch=%llu device=%d tokens=%d routes=%d shared=%d\n", - launch, device, n_tokens, n_used, shared_out != nullptr ? 1 : 0); - } - GGML_ASSERT(n_embd % 4 == 0 && "n_embd must be a multiple of 4 for float4 vectorization"); GGML_ASSERT(reinterpret_cast(down_e->data) % 16 == 0 && "down_e->data must be 16-byte aligned"); GGML_ASSERT(reinterpret_cast(dst->data) % 16 == 0 && "dst->data must be 16-byte aligned"); diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 3471a60b4..4d6bac797 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -359,15 +359,6 @@ struct DeepSeek4CachedDecodeAttnGraph { } }; -static bool ds4_moe_fused_combine_enabled() { - static const bool enabled = []() { - const char * val = getenv("DFLASH_MOE_FUSED_COMBINE"); - if (!val) return true; // Default ON in production - return atoi(val) != 0; - }(); - return enabled; -} - struct DeepSeek4CachedLayerAlloc { const ggml_context * owner_ctx = nullptr; ggml_backend_t backend = nullptr; @@ -486,18 +477,8 @@ static bool build_cached_decode_ffn_graph( weights = ggml_scale(out.sg.ctx, weights, w.expert_weight_scale); } - if (ds4_moe_fused_combine_enabled()) { - ffn_out = ggml_ds4_moe_fused_combine_shared(out.sg.ctx, down_e, weights, shared_out); - } else { - ggml_tensor * weights_3d = ggml_reshape_3d(out.sg.ctx, weights, 1, n_used, n_tokens); - ggml_tensor * routed_out = ggml_mul(out.sg.ctx, down_e, weights_3d); - routed_out = ggml_cont( - out.sg.ctx, ggml_permute(out.sg.ctx, routed_out, 1, 0, 2, 3)); - routed_out = ggml_sum_rows(out.sg.ctx, routed_out); - routed_out = ggml_reshape_2d(out.sg.ctx, routed_out, w.n_embd, n_tokens); - - ffn_out = ggml_add(out.sg.ctx, shared_out, routed_out); - } + ffn_out = ggml_ds4_moe_fused_combine_shared( + out.sg.ctx, down_e, weights, shared_out); } else { ffn_out = build_moe_ffn(out.sg.ctx, ffn_normed, w, L, layer_idx, n_tokens); } @@ -3361,16 +3342,8 @@ static ggml_tensor * build_moe_ffn( ggml_tensor * down_e = ggml_mul_mat_id(ctx, L.ffn_down_exps, mid_e, routing.selected); down_e = ggml_reshape_3d(ctx, down_e, n_embd, n_used, n_tokens); - if (ds4_moe_fused_combine_enabled()) { - return ggml_ds4_moe_fused_combine_shared(ctx, down_e, routing.weights, shared_out); - } else { - ggml_tensor * weights_3d = ggml_reshape_3d(ctx, routing.weights, 1, n_used, n_tokens); - routed_out = ggml_mul(ctx, down_e, weights_3d); - routed_out = ggml_cont(ctx, ggml_permute(ctx, routed_out, 1, 0, 2, 3)); - routed_out = ggml_sum_rows(ctx, routed_out); - routed_out = ggml_reshape_2d(ctx, routed_out, n_embd, n_tokens); - return ggml_add(ctx, shared_out, routed_out); - } + return ggml_ds4_moe_fused_combine_shared( + ctx, down_e, routing.weights, shared_out); } return ggml_add(ctx, shared_out, routed_out); @@ -5101,26 +5074,8 @@ static ggml_tensor * ds4_build_hash_routed_ffn( weights = ggml_scale(ctx, weights, w.expert_weight_scale); } - if (ds4_moe_fused_combine_enabled()) { - return ggml_ds4_moe_fused_combine_shared(ctx, down_e, weights, shared_out); - } - - ggml_tensor * weights_3d = ggml_reshape_3d( - ctx, weights, 1, n_used, n_tokens); - ggml_tensor * routed_out = ggml_mul(ctx, down_e, weights_3d); - if (n_tokens == 1) { - // Preserve the established q=1 graph and reduction order. - routed_out = ggml_cont(ctx, ggml_permute(ctx, routed_out, 1, 0, 2, 3)); - routed_out = ggml_sum_rows(ctx, routed_out); - routed_out = ggml_reshape_2d(ctx, routed_out, w.n_embd, 1); - } else { - ggml_tensor * sum_shape = ggml_new_tensor_3d( - ctx, GGML_TYPE_F32, w.n_embd, 1, n_tokens); - routed_out = ggml_repeat_back(ctx, routed_out, sum_shape); - routed_out = ggml_reshape_2d( - ctx, routed_out, w.n_embd, n_tokens); - } - return ggml_add(ctx, shared_out, routed_out); + return ggml_ds4_moe_fused_combine_shared( + ctx, down_e, weights, shared_out); } static bool ds4_fused_ensure_fn_mirrors(