perf(ds4): fuse expert-major MoE route combine

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
@@ -1538,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)

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,

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;
+}

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"
 
@@ -873,6 +874,47 @@ 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<ggml_backend_meta_split_state> & 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 ||
+                        (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) {
+                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];
+        }
+
+        // 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_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,
+                    {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 +1154,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;

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:

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,15 @@ 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->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:
         case GGML_OP_MUL_MAT_GROUPED_SRC:
         case GGML_OP_MUL_MAT_ID: