feat(moe): add a generic heterogeneous stage planner

server/CMakeLists.txt

@@ -429,6 +429,7 @@ add_library(dflash_common STATIC
     src/qwen35moe/qwen35moe_daemon.cpp
     src/qwen35moe/qwen35moe_pipelined_decode.cpp
     # ── Common MoE hybrid infrastructure ──
+    src/common/heterogeneous_stage_planner.cpp
     src/common/moe_hybrid_placement.cpp
     src/common/moe_hybrid_routing_stats.cpp
     src/common/moe_hybrid_storage.cpp

server/deps/llama.cpp/ggml/include/ggml.h

@@ -2485,6 +2485,7 @@ extern "C" {
         GGML_MOE_FUSED_DEFERRED_PEER_COPY = -3,
         GGML_MOE_FUSED_OWNER_SPLIT        = -4,
         GGML_MOE_FUSED_ALIGN_IDS          = -5,
+        GGML_MOE_FUSED_BALANCED_OWNER_IDS = -6,
     };
 
     // Word offsets in ggml_tensor::op_params for the deferred peer-copy op.
@@ -2518,21 +2519,24 @@ extern "C" {
 
     // Coarse DeepSeek-V4 routed-owner op. gate_up contains concatenated gate
     // and up output rows; the backend performs fused gate/up MMVQ + clamped
-    // SwiGLU, down MMVQ, route weighting, and local reduction.
+    // SwiGLU, down MMVQ, route weighting, and local reduction. An optional
+    // owner_residual is added in that final reduction without another launch.
     GGML_API struct ggml_tensor * ggml_ds4_moe_owner(
             struct ggml_context * ctx,
             struct ggml_tensor  * input,
             struct ggml_tensor  * gate_up_w,
             struct ggml_tensor  * down_w,
             struct ggml_tensor  * expert_ids,
             struct ggml_tensor  * expert_weights,
+            struct ggml_tensor  * owner_residual,
             int64_t               ff_dim,
             float                 swiglu_clamp,
             float                 down_scale);
 
     // Coarse owner variant for checkpoints that store gate and up tensors
     // separately. The backend fuses their MMVQ traversal and SwiGLU while
-    // preserving the checkpoint's three external value scales.
+    // preserving the checkpoint's three external value scales. The optional
+    // owner_residual has the same semantics as in ggml_ds4_moe_owner.
     GGML_API struct ggml_tensor * ggml_ds4_moe_owner_split(
             struct ggml_context * ctx,
             struct ggml_tensor  * input,
@@ -2541,6 +2545,7 @@ extern "C" {
             struct ggml_tensor  * down_w,
             struct ggml_tensor  * expert_ids,
             struct ggml_tensor  * expert_weights,
+            struct ggml_tensor  * owner_residual,
             int64_t               ff_dim,
             float                 swiglu_clamp,
             float                 gate_scale,
@@ -2554,6 +2559,20 @@ extern "C" {
             struct ggml_context * ctx,
             struct ggml_tensor  * expert_ids);
 
+    // Map global route IDs directly to one owner's compact expert stack while
+    // assigning a batch-wide main-resident quota encoded as four times the
+    // desired routes per token. The peer owner receives the exact complement.
+    // This fuses the former repeated-LUT lookup and route masking chain into
+    // one small device-local operation.
+    GGML_API struct ggml_tensor * ggml_ds4_moe_balanced_owner_ids(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * global_ids,
+            struct ggml_tensor  * router_weights,
+            struct ggml_tensor  * local_id_lut,
+            struct ggml_tensor  * main_candidate_lut,
+            int                   main_slots_x4,
+            bool                  main_owner);
+
     // Copy an F32 peer-GPU tensor only after a scheduler-provided device event
     // has completed. The scheduler writes the native event handle into the op
     // parameters and deliberately leaves src on its owner backend.

server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu

@@ -719,7 +719,7 @@ static bool rocmfp4_x4_enabled() {
     return enabled;
 }
 
-static bool rocmfp4_q5_x4_plus1_enabled() {
+static bool rocmfp4_qwide_x4_tail_enabled() {
     static const bool enabled = []() {
         const char * value = std::getenv("DFLASH_CUDA_MMVQ_FP4_Q5_X4_PLUS1");
         return value && value[0] == '1' && value[1] == '\0';
@@ -767,8 +767,8 @@ static __global__ void mul_mat_vec_q(
                   "packed FP3 MMVQ specialization requires ROCmFP3 weights");
     static_assert(!c_fp4_x4 ||
                       (type == GGML_TYPE_Q4_0_ROCMFP4_FAST &&
-                       (ncols_dst == 4 || ncols_dst == 5)),
-                  "FP4 x4 MMVQ specialization requires ROCmFP4-fast q4/q5");
+                       (ncols_dst >= 4 && ncols_dst <= 6)),
+                  "FP4 x4 MMVQ specialization requires ROCmFP4-fast q4..q6");
 
     const uint32_t channel_dst = blockIdx.y;
 
@@ -982,10 +982,13 @@ static __global__ void mul_mat_vec_q(
                     tmp[1][i] += dots.y;
                     tmp[2][i] += dots.z;
                     tmp[3][i] += dots.w;
-                    if constexpr (ncols_dst == 5) {
-                        tmp[4][i] += vec_dot_q_mmvq<type, false>(
-                            vx, &y[4*stride_col_y + kby],
-                            kbx_offset + i*stride_row_x + kbx, kqs);
+                    if constexpr (ncols_dst > 4) {
+#pragma unroll
+                        for (int j = 4; j < ncols_dst; ++j) {
+                            tmp[j][i] += vec_dot_q_mmvq<type, false>(
+                                vx, &y[j*stride_col_y + kby],
+                                kbx_offset + i*stride_row_x + kbx, kqs);
+                        }
                     }
                     if constexpr (has_fusion) {
                         if (use_gate) {
@@ -1002,10 +1005,13 @@ static __global__ void mul_mat_vec_q(
                             tmp_gate[1][i] += gate_dots.y;
                             tmp_gate[2][i] += gate_dots.z;
                             tmp_gate[3][i] += gate_dots.w;
-                            if constexpr (ncols_dst == 5) {
-                                tmp_gate[4][i] += vec_dot_q_mmvq<type, false>(
-                                    vgate, &y[4*stride_col_y + kby],
-                                    kbx_offset + i*stride_row_x + kbx, kqs);
+                            if constexpr (ncols_dst > 4) {
+#pragma unroll
+                                for (int j = 4; j < ncols_dst; ++j) {
+                                    tmp_gate[j][i] += vec_dot_q_mmvq<type, false>(
+                                        vgate, &y[j*stride_col_y + kby],
+                                        kbx_offset + i*stride_row_x + kbx, kqs);
+                                }
                             }
                         }
                     }
@@ -2228,9 +2234,24 @@ static void mul_mat_vec_q_switch_ncols_dst(
                 0, ids_stride, stream);
             return;
         }
-        if (!has_ids && ncols_dst == 5 && rocmfp4_x4_enabled() &&
-            rocmfp4_q5_x4_plus1_enabled()) {
-            constexpr int c_ncols_dst = 5;
+        if (!has_ids && (ncols_dst == 5 || ncols_dst == 6) &&
+            rocmfp4_x4_enabled() && rocmfp4_qwide_x4_tail_enabled()) {
+            if (ncols_dst == 5) {
+                constexpr int c_ncols_dst = 5;
+                std::pair<dim3, dim3> dims = calc_launch_params<type>(
+                    c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst,
+                    warp_size, table_id);
+                mul_mat_vec_q_switch_fusion<type, c_ncols_dst, false, 0,
+                                            false, false, false, true>(
+                    vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd,
+                    stride_row_x, stride_col_y, stride_col_dst,
+                    channel_ratio_fd, stride_channel_x, stride_channel_y,
+                    stride_channel_dst, sample_ratio_fd, stride_sample_x,
+                    stride_sample_y, stride_sample_dst, dims.first, dims.second,
+                    0, ids_stride, stream);
+                return;
+            }
+            constexpr int c_ncols_dst = 6;
             std::pair<dim3, dim3> dims = calc_launch_params<type>(
                 c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst,
                 warp_size, table_id);

server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused.cu

@@ -265,6 +265,7 @@ static __global__ void moe_fused_kernel(
 static __global__ void laguna_moe_combine_kernel(
     const char * __restrict__ experts,
     const char * __restrict__ weights,
+    const char * __restrict__ owner_residual,
     char * __restrict__ output,
     const int n_embd,
     const int n_used,
@@ -274,6 +275,8 @@ static __global__ void laguna_moe_combine_kernel(
     const size_t experts_nb2,
     const size_t weights_nb0,
     const size_t weights_nb1,
+    const size_t residual_nb0,
+    const size_t residual_nb1,
     const size_t output_nb0,
     const size_t output_nb1,
     const float value_scale) {
@@ -304,6 +307,12 @@ static __global__ void laguna_moe_combine_kernel(
         const float prod = __fmul_rn(scaled, w);
         sum = (e == 0) ? prod : __fadd_rn(sum, prod);
     }
+    if (owner_residual != nullptr) {
+        const float residual = *(const float *)(owner_residual +
+            (size_t)h * residual_nb0 +
+            (size_t)t * residual_nb1);
+        sum = __fadd_rn(sum, residual);
+    }
     *(float *)(output +
         (size_t)h * output_nb0 +
         (size_t)t * output_nb1) = sum;
@@ -323,8 +332,10 @@ static __global__ void ds4_peer_copy_f32_kernel(
 // of every valid output encode the original route slot; invalid entries use
 // the sign bit plus the original slot.  The dedicated MoE MMVQ kernel decodes
 // this metadata and scatters its result back to the original route layout.
-// A single thread is deliberate: this is at most a 4 x 6 assignment problem
-// and runs once per owner/layer, outside the weight-streaming kernels.
+// A single thread is deliberate: practical speculative MoE batches are tiny,
+// and this runs once per owner/layer outside the weight-streaming kernels. The
+// bounded fast path covers common top-k and speculative widths; larger shapes
+// retain the original route order and exact scatter metadata.
 static __global__ void ds4_align_moe_ids_kernel(
         const int32_t * __restrict__ ids,
         int32_t * __restrict__ aligned,
@@ -336,8 +347,8 @@ static __global__ void ds4_align_moe_ids_kernel(
         return;
     }
 
-    constexpr int max_routes = 6;
-    constexpr int max_tokens = 4;
+    constexpr int max_routes = 16;
+    constexpr int max_tokens = 16;
     if (n_routes > max_routes || n_tokens > max_tokens) {
         for (int t = 0; t < n_tokens; ++t) {
             for (int r = 0; r < n_routes; ++r) {
@@ -435,6 +446,49 @@ static __global__ void ds4_align_moe_ids_kernel(
     }
 }
 
+static __global__ void ds4_balanced_owner_ids_kernel(
+        const int32_t * __restrict__ global_ids,
+        const float * __restrict__ router_weights,
+        const int32_t * __restrict__ local_id_lut,
+        const float * __restrict__ main_candidate_lut,
+        int32_t * __restrict__ owner_ids,
+        int n_routes,
+        int n_tokens,
+        int n_expert,
+        int main_slots_x4,
+        bool main_owner) {
+    if (blockIdx.x != 0 || threadIdx.x != 0) {
+        return;
+    }
+
+    // Assign one shared quota over the complete verification batch. Quarter-
+    // route units expose the otherwise unreachable 16/30 split at q=5 while
+    // keeping the decision device-local and identical on both owners.
+    const int main_quota = (main_slots_x4 * n_tokens) / 4;
+    int assigned_main = 0;
+    for (int token = 0; token < n_tokens; ++token) {
+        const int row = token * n_routes;
+        for (int route = 0; route < n_routes; ++route) {
+            const int index = row + route;
+            const int32_t global_id = global_ids[index];
+            const bool active = router_weights[index] != 0.0f;
+            const bool valid_global = global_id >= 0 && global_id < n_expert;
+            const bool main_candidate = active && valid_global &&
+                main_candidate_lut[global_id] != 0.0f;
+            const bool route_on_main =
+                main_candidate && assigned_main < main_quota;
+            if (route_on_main) {
+                ++assigned_main;
+            }
+
+            const bool keep = active && valid_global &&
+                (main_owner ? route_on_main : !route_on_main);
+            const int32_t local_id = keep ? local_id_lut[global_id] : -1;
+            owner_ids[index] = local_id >= 0 ? local_id : -1;
+        }
+    }
+}
+
 static ggml_tensor make_contiguous_f32_tensor(
         float * data,
         int64_t ne0,
@@ -462,6 +516,7 @@ static void ggml_cuda_op_ds4_moe_owner(
     const ggml_tensor * down_w     = dst->src[2];  // [n_ff, n_embd, n_expert]
     const ggml_tensor * expert_ids = dst->src[3];  // [n_used, n_tokens]
     const ggml_tensor * weights    = dst->src[4];  // [n_used, n_tokens]
+    const ggml_tensor * residual   = dst->src[5];  // optional [n_embd, n_tokens]
 
     const int n_embd   = (int) input->ne[0];
     const int n_tokens = (int) input->ne[1];
@@ -477,6 +532,9 @@ static void ggml_cuda_op_ds4_moe_owner(
     GGML_ASSERT(expert_ids->ne[1] == n_tokens);
     GGML_ASSERT(weights->ne[0] == n_used && weights->ne[1] == n_tokens);
     GGML_ASSERT(dst->ne[0] == n_embd && dst->ne[1] == n_tokens);
+    GGML_ASSERT(!residual ||
+        (residual->type == GGML_TYPE_F32 &&
+         residual->ne[0] == n_embd && residual->ne[1] == n_tokens));
 
     // The checkpoint concatenates gate rows followed by up rows inside each
     // expert. Keep the original expert stride while viewing each half.
@@ -524,10 +582,13 @@ static void ggml_cuda_op_ds4_moe_owner(
     laguna_moe_combine_kernel<<<grid, block, 0, ctx.stream()>>>(
         (const char *) experts.data,
         (const char *) weights->data,
+        residual ? (const char *) residual->data : nullptr,
         (char *) dst->data,
         n_embd, n_used, n_tokens,
         experts.nb[0], experts.nb[1], experts.nb[2],
         weights->nb[0], weights->nb[1],
+        residual ? residual->nb[0] : 0,
+        residual ? residual->nb[1] : 0,
         dst->nb[0], dst->nb[1],
         down_scale);
 }
@@ -541,6 +602,7 @@ static void ggml_cuda_op_ds4_moe_owner_split(
     const ggml_tensor * down_w     = dst->src[3];  // [n_ff, n_embd, n_expert]
     const ggml_tensor * expert_ids = dst->src[4];  // [n_used, n_tokens]
     const ggml_tensor * weights    = dst->src[5];  // [n_used, n_tokens]
+    const ggml_tensor * residual   = dst->src[6];  // optional [n_embd, n_tokens]
 
     const int n_embd   = (int) input->ne[0];
     const int n_tokens = (int) input->ne[1];
@@ -557,6 +619,9 @@ static void ggml_cuda_op_ds4_moe_owner_split(
     GGML_ASSERT(expert_ids->ne[1] == n_tokens);
     GGML_ASSERT(weights->ne[0] == n_used && weights->ne[1] == n_tokens);
     GGML_ASSERT(dst->ne[0] == n_embd && dst->ne[1] == n_tokens);
+    GGML_ASSERT(!residual ||
+        (residual->type == GGML_TYPE_F32 &&
+         residual->ne[0] == n_embd && residual->ne[1] == n_tokens));
 
     // MUL_MAT_ID consumes token columns in dimension 2. This descriptor is
     // stack-local, so it deliberately does not participate in q8 memoization.
@@ -597,16 +662,49 @@ static void ggml_cuda_op_ds4_moe_owner_split(
     laguna_moe_combine_kernel<<<grid, block, 0, ctx.stream()>>>(
         (const char *) experts.data,
         (const char *) weights->data,
+        residual ? (const char *) residual->data : nullptr,
         (char *) dst->data,
         n_embd, n_used, n_tokens,
         experts.nb[0], experts.nb[1], experts.nb[2],
         weights->nb[0], weights->nb[1],
+        residual ? residual->nb[0] : 0,
+        residual ? residual->nb[1] : 0,
         dst->nb[0], dst->nb[1],
         down_scale);
 }
 
 void ggml_cuda_op_moe_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     const int mode = ggml_get_op_params_i32(dst, 0);
+    if (mode == GGML_MOE_FUSED_BALANCED_OWNER_IDS) {
+        const ggml_tensor * global_ids = dst->src[0];
+        const ggml_tensor * weights = dst->src[1];
+        const ggml_tensor * local_lut = dst->src[2];
+        const ggml_tensor * candidate_lut = dst->src[3];
+        GGML_ASSERT(global_ids && global_ids->type == GGML_TYPE_I32);
+        GGML_ASSERT(weights && weights->type == GGML_TYPE_F32);
+        GGML_ASSERT(local_lut && local_lut->type == GGML_TYPE_I32);
+        GGML_ASSERT(candidate_lut && candidate_lut->type == GGML_TYPE_F32);
+        GGML_ASSERT(dst->type == GGML_TYPE_I32);
+        GGML_ASSERT(ggml_is_contiguous(global_ids));
+        GGML_ASSERT(ggml_is_contiguous(weights));
+        GGML_ASSERT(ggml_is_contiguous(local_lut));
+        GGML_ASSERT(ggml_is_contiguous(candidate_lut));
+        GGML_ASSERT(ggml_is_contiguous(dst));
+
+        const int n_routes = (int) global_ids->ne[0];
+        const int n_tokens = (int) global_ids->ne[1];
+        const int n_expert = (int) ggml_nelements(local_lut);
+        const int main_slots_x4 = ggml_get_op_params_i32(dst, 1);
+        const bool main_owner = ggml_get_op_params_i32(dst, 2) != 0;
+        ds4_balanced_owner_ids_kernel<<<1, 1, 0, ctx.stream()>>>(
+            (const int32_t *) global_ids->data,
+            (const float *) weights->data,
+            (const int32_t *) local_lut->data,
+            (const float *) candidate_lut->data,
+            (int32_t *) dst->data,
+            n_routes, n_tokens, n_expert, main_slots_x4, main_owner);
+        return;
+    }
     if (mode == GGML_MOE_FUSED_ALIGN_IDS) {
         const ggml_tensor * ids = dst->src[0];
         GGML_ASSERT(ids && ids->type == GGML_TYPE_I32);
@@ -678,10 +776,12 @@ void ggml_cuda_op_moe_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst)
         laguna_moe_combine_kernel<<<grid, block, 0, ctx.stream()>>>(
             (const char *) experts->data,
             (const char *) weights->data,
+            nullptr,
             (char *) dst->data,
             n_embd, n_used, n_tokens,
             experts->nb[0], experts->nb[1], experts->nb[2],
             weights->nb[0], weights->nb[1],
+            0, 0,
             dst->nb[0], dst->nb[1],
             1.0f);
         return;