From 199743801af2d4cc4fff519a268463699cc63068 Mon Sep 17 00:00:00 2001 From: cheese-cakee Date: Thu, 27 Aug 2026 00:43:12 +0530 Subject: [PATCH 1/5] 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 925d5839d..fdd25018d 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -1019,6 +1019,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 0e1bcebf4..41937523d 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, }; @@ -2713,6 +2715,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 82052f611..2a035449d 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" @@ -737,6 +738,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) { @@ -3422,6 +3424,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; @@ -3809,14 +3814,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( @@ -5966,6 +5970,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 ef6f51532..6ae68bb28 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"); @@ -9044,3 +9048,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 8986328b1..a4e81811b 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); } @@ -3322,11 +3335,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); @@ -5051,6 +5069,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 0ad294be66a69af7c19dc9e9bd56063c9bdad427 Mon Sep 17 00:00:00 2001 From: cheese-cakee Date: Thu, 27 Aug 2026 00:58:14 +0530 Subject: [PATCH 2/5] chore: retrigger PR checks From 6282f535f21fa2b7597f22f55415ecbf0823c6ae Mon Sep 17 00:00:00 2001 From: cheese-cakee Date: Thu, 27 Aug 2026 15:32:51 +0530 Subject: [PATCH 3/5] 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 2a035449d..029aa758d 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 @@ -5976,6 +5976,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 7b1af18d0f68bc686b22c75cb435cfb37c2062f0 Mon Sep 17 00:00:00 2001 From: cheese-cakee Date: Thu, 27 Aug 2026 15:35:27 +0530 Subject: [PATCH 4/5] 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 7d0ef55732cc5c9260bcf349baec5c3533ee4094 Mon Sep 17 00:00:00 2001 From: cheese-cakee Date: Thu, 27 Aug 2026 15:48:54 +0530 Subject: [PATCH 5/5] 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 fdd25018d..dba6228d3 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -1548,6 +1548,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; +}