Add opt-in CUDA routed MoE FFN

Author: farkasmark

include/backend_model.h

@@ -38,6 +38,9 @@ typedef enum {
     BN_BACKEND_HANDLE_MOE_ROUTER_DIFF = 25,
     BN_BACKEND_HANDLE_SHARED_GATEUP_STACKED = 26,
     BN_BACKEND_HANDLE_MOE_ROUTER = 27,
+    BN_BACKEND_HANDLE_MOE_GATE_ALL = 28,
+    BN_BACKEND_HANDLE_MOE_UP_ALL = 29,
+    BN_BACKEND_HANDLE_MOE_DOWN_ALL = 30,
 } BnBackendHandleRole;
 
 BnBackendModel *bn_backend_model_create(void);

src/gpu_cuda.cu

@@ -16,6 +16,7 @@
 #include <stdlib.h>
 #include <string.h>
 #include <time.h>
+#include <limits.h>
 
 typedef struct {
     void *data;
@@ -3444,6 +3445,87 @@ static __global__ void moe_route_topk_kernel(float *route,
         route[k + i] = (float)selected[i];
 }
 
+static __global__ void moe_q4k_gateup_routed_mid_kernel(
+    float *mid,
+    const BnBlockQ4K *gate,
+    const BnBlockQ4K *up,
+    const BnCudaBlockQ8_1 *xq,
+    const float *route,
+    int hidden,
+    int cols,
+    int n_experts,
+    int k) {
+    int lane = threadIdx.x & 31;
+    int warp = threadIdx.x >> 5;
+    int warps_per_block = blockDim.x >> 5;
+    int task = blockIdx.x * warps_per_block + warp;
+    if (task >= hidden * k) return;
+
+    int slot = task / hidden;
+    int row = task - slot * hidden;
+    int expert = (int)(route[k + slot] + 0.5f);
+    if (expert < 0) expert = 0;
+    if (expert >= n_experts) expert = n_experts - 1;
+    float route_weight = route[slot];
+
+    int n_bpr = cols / BN_QK_K;
+    int kbx = lane / 16;
+    int iqs = 2 * (lane & 15);
+    size_t expert_row = ((size_t)expert * (size_t)hidden + (size_t)row);
+    const BnBlockQ4K *gate_blocks = gate + expert_row * (size_t)n_bpr;
+    const BnBlockQ4K *up_blocks = up + expert_row * (size_t)n_bpr;
+    float gate_sum = 0.0f;
+    float up_sum = 0.0f;
+    for (int b = kbx; b < n_bpr; b += 2) {
+        const BnCudaBlockQ8_1 *xqb = xq + (size_t)b * 8;
+        gate_sum += cuda_vec_dot_q4k_q8_1(&gate_blocks[b], xqb, iqs);
+        up_sum += cuda_vec_dot_q4k_q8_1(&up_blocks[b], xqb, iqs);
+    }
+    for (int offset = 16; offset > 0; offset >>= 1) {
+        gate_sum += __shfl_down_sync(0xffffffffu, gate_sum, offset);
+        up_sum += __shfl_down_sync(0xffffffffu, up_sum, offset);
+    }
+    if (lane == 0) {
+        float silu = gate_sum / (1.0f + __expf(-gate_sum));
+        mid[(size_t)slot * (size_t)hidden + (size_t)row] =
+            route_weight * silu * up_sum;
+    }
+}
+
+static __global__ void moe_q6k_down_routed_q8k_accum_kernel(
+    float *out,
+    const BnBlockQ6K *down,
+    const BnBlockQ8K *mid_q,
+    const float *route,
+    int dim,
+    int hidden,
+    int n_experts,
+    int k) {
+    int lane = threadIdx.x & 31;
+    int warp = threadIdx.x >> 5;
+    int warps_per_block = blockDim.x >> 5;
+    int row = blockIdx.x * warps_per_block + warp;
+    if (row >= dim) return;
+
+    int n_bpr = hidden / BN_QK_K;
+    float sum = 0.0f;
+    for (int slot = 0; slot < k; slot++) {
+        int expert = (int)(route[k + slot] + 0.5f);
+        if (expert < 0) expert = 0;
+        if (expert >= n_experts) expert = n_experts - 1;
+        const BnBlockQ6K *row_blocks =
+            down + (((size_t)expert * (size_t)dim + (size_t)row) *
+                    (size_t)n_bpr);
+        const BnBlockQ8K *slot_mid_q = mid_q + (size_t)slot * (size_t)n_bpr;
+        for (int b = lane; b < n_bpr; b += 32)
+            sum += cuda_vec_dot_q6k_q8k(&row_blocks[b], slot_mid_q + b);
+    }
+    for (int offset = 16; offset > 0; offset >>= 1)
+        sum += __shfl_down_sync(0xffffffffu, sum, offset);
+    if (lane == 0)
+        out[row] = sum;
+}
+
 static __device__ __forceinline__ float cuda_fast_exp(float x) {
     x = fminf(88.7f, fmaxf(-87.3f, x));
     float n_f = floorf(x * 1.4426950409f + 0.5f);
@@ -5173,6 +5255,10 @@ static int cuda_init_activations(void *vctx, const void *config_ptr) {
             moe_scratch = c->n_experts;
         if (2 * c->n_experts_active > moe_scratch)
             moe_scratch = 2 * c->n_experts_active;
+        if (c->n_experts_active > 0 &&
+            c->moe_intermediate_size <= INT_MAX / c->n_experts_active &&
+            c->moe_intermediate_size * c->n_experts_active > moe_scratch)
+            moe_scratch = c->moe_intermediate_size * c->n_experts_active;
         sizes[BN_GPU_VALUE_MOE_HB] =
             (size_t)moe_scratch * sizeof(float);
         sizes[BN_GPU_VALUE_MOE_HB2] =
@@ -8369,6 +8455,7 @@ static const char *cuda_op_name(int code) {
     case BN_GPU_CODE_SILU_ACT: return "silu_act";
     case BN_GPU_CODE_RELU2_ACT: return "relu2_act";
     case BN_GPU_CODE_MOE_ROUTE_TOPK: return "moe_route_topk";
+    case BN_GPU_CODE_MOE_ROUTED_FFN: return "moe_routed_ffn";
     case BN_GPU_CODE_ROPE: return "rope";
     case BN_GPU_CODE_ROPE_QK: return "rope_qk";
     case BN_GPU_CODE_GQA_SCORES: return "gqa_scores";
@@ -8434,6 +8521,7 @@ static int cuda_op_reads_buf(const BnGPUOp *op, int buf) {
     case BN_GPU_CODE_Q5K_MATVEC_SPLIT:
     case BN_GPU_CODE_FUSED_GATEUP_SILU:
     case BN_GPU_CODE_MOE_ROUTE_TOPK:
+    case BN_GPU_CODE_MOE_ROUTED_FFN:
     case BN_GPU_CODE_RMSNORM:
     case BN_GPU_CODE_PER_HEAD_RMSNORM:
     case BN_GPU_CODE_COPY:
@@ -9433,6 +9521,65 @@ static int cuda_execute(void *vctx, const void *ops_raw, int n_ops,
                 route, logits, n_experts, k);
             break;
         }
+        case BN_GPU_CODE_MOE_ROUTED_FFN: {
+            BnCudaBuffer *gate = (BnCudaBuffer *)op->W_buf;
+            BnCudaBuffer *up = (BnCudaBuffer *)op->W_buf2;
+            BnCudaBuffer *down = (BnCudaBuffer *)op->W_buf3;
+            float *in = cuda_act(ctx, op->buf_in);
+            float *route = cuda_act(ctx, op->buf_aux);
+            int mid_buf = (int)op->p[4];
+            float *mid = cuda_act(ctx, mid_buf);
+            float *out = cuda_act(ctx, op->buf_out);
+            int hidden = (int)op->p[0];
+            int n_experts = (int)op->p[1];
+            int k = (int)op->p[2];
+            int down_type = (int)op->p[3];
+            int dim = op->cols;
+            if (!gate || !gate->data || !up || !up->data ||
+                !down || !down->data || !in || !route || !mid || !out ||
+                op->type != BN_GGUF_TENSOR_Q4_K ||
+                down_type != BN_GGUF_TENSOR_Q6_K ||
+                dim <= 0 || hidden <= 0 || n_experts <= 0 || k <= 0 ||
+                (dim % BN_QK_K) != 0 || (hidden % BN_QK_K) != 0 ||
+                gate->type != BN_GGUF_TENSOR_Q4_K ||
+                up->type != BN_GGUF_TENSOR_Q4_K ||
+                down->type != BN_GGUF_TENSOR_Q6_K ||
+                gate->rows < hidden * n_experts || gate->cols < dim ||
+                up->rows < hidden * n_experts || up->cols < dim ||
+                down->rows < dim * n_experts || down->cols < hidden ||
+                ctx->act_sizes[op->buf_aux] <
+                    (size_t)(2 * k) * sizeof(float) ||
+                ctx->act_sizes[mid_buf] <
+                    (size_t)k * (size_t)hidden * sizeof(float) ||
+                ctx->act_sizes[op->buf_out] < (size_t)dim * sizeof(float))
+                return -1;
+            if (cuda_ensure_q8_1(ctx, dim) != 0) return -1;
+            BnCudaBlockQ8_1 *xq = (BnCudaBlockQ8_1 *)ctx->d_q8_1;
+            BN_CUDA_LAUNCH(ctx, quantize_q8_1_kernel,
+                (dim + 31) / 32, 32, 0, xq, in, dim);
+            {
+                int route_threads = 256;
+                int warps = route_threads / 32;
+                int gateup_tasks = hidden * k;
+                int gateup_blocks = (gateup_tasks + warps - 1) / warps;
+                BN_CUDA_LAUNCH(ctx, moe_q4k_gateup_routed_mid_kernel,
+                    gateup_blocks, route_threads, 0,
+                    mid, (const BnBlockQ4K *)gate->data,
+                    (const BnBlockQ4K *)up->data, xq, route, hidden, dim,
+                    n_experts, k);
+                if (cuda_ensure_q8_k(ctx, hidden, k) != 0) return -1;
+                BnBlockQ8K *mid_q = (BnBlockQ8K *)ctx->d_q8_k;
+                BN_CUDA_LAUNCH(ctx, quantize_q8k_batch_kernel,
+                    dim3(hidden / BN_QK_K, k, 1), BN_QK_K, 0,
+                    mid_q, mid, hidden, k);
+                int down_blocks = (dim + warps - 1) / warps;
+                BN_CUDA_LAUNCH(ctx, moe_q6k_down_routed_q8k_accum_kernel,
+                    down_blocks, route_threads, 0,
+                    out, (const BnBlockQ6K *)down->data, mid_q, route,
+                    dim, hidden, n_experts, k);
+            }
+            break;
+        }
         case BN_GPU_CODE_RMSNORM: {
             BnCudaBuffer *w = (BnCudaBuffer *)op->W_buf;
             float *in = cuda_act(ctx, op->buf_in);

src/gpu_shader_ir_internal.h

@@ -82,14 +82,17 @@ typedef enum {
     BN_GPU_CODE_RELU2_ACT,
     BN_GPU_CODE_WEIGHTED_ADD_SIGMOID,
     BN_GPU_CODE_MOE_ROUTE_TOPK,
+    BN_GPU_CODE_MOE_ROUTED_FFN,
 } BnGPUOpCode;
 
 // A single backend shader command in the lowered forward pass.
 typedef struct BnGPUOp {
     int op_kind;         // BnGPUOpKind semantic op; 0 = infer from op_code
     int op_code;         // BnGPUOpCode concrete shader operation
     int type;            // BN_GGUF_TENSOR_* (matvec only, -1 otherwise)
-    void *W_buf;         // weight buffer handle (matvec only, NULL otherwise)
+    void *W_buf;         // primary weight buffer handle
+    void *W_buf2;        // optional secondary weight buffer handle
+    void *W_buf3;        // optional tertiary weight buffer handle
     int buf_in;          // BN_GPU_VALUE_* primary input
     int buf_out;         // BN_GPU_VALUE_* output
     int buf_aux;         // secondary BN_GPU_VALUE_* (-1 if unused)
@@ -134,6 +137,7 @@ static inline BnGPUOpKind bn_gpu_op_kind_from_code(int code) {
             return BN_GPU_OP_COPY;
         case BN_GPU_CODE_FUSED_GATEUP_SILU:
         case BN_GPU_CODE_MOE_ROUTE_TOPK:
+        case BN_GPU_CODE_MOE_ROUTED_FFN:
             return BN_GPU_OP_FFN;
         case BN_GPU_CODE_SSM_CONV_SILU:
         case BN_GPU_CODE_SSM_L2NORM:

src/model_gpu.c

@@ -2,8 +2,10 @@
 #include "backend_layout.h"
 #include "backend_model.h"
 #include "gpu_backend.h"
+#include "moe_internal.h"
 #include <stdlib.h>
 #include <stdint.h>
+#include <limits.h>
 
 static int checked_mul_size(size_t a, size_t b, size_t *out) {
     if (a != 0 && b > SIZE_MAX / a) return -1;
@@ -69,6 +71,59 @@ static void *upload_moe_router_diff2(BnGPUBackend *gpu,
     return handle;
 }
 
+static void *upload_moe_all_proj(BnModel *model,
+                                 BnGPUBackend *gpu,
+                                 const BnMoEExpertMap *em,
+                                 int proj,
+                                 int n_experts) {
+    if (!model || !gpu || !em || n_experts <= 0)
+        return NULL;
+    size_t offset = 0;
+    size_t expert_bytes = 0;
+    if (bn_moe_proj_info(em, 0, proj, &offset, &expert_bytes) != 0 ||
+        expert_bytes == 0)
+        return NULL;
+    size_t stride = 0;
+    int type = 0;
+    int rows = 0;
+    int cols = 0;
+    switch (proj) {
+        case 0:
+            stride = em->gate_stride ? em->gate_stride : em->expert_gate_bytes;
+            type = em->gate_type;
+            rows = em->gate_rows;
+            cols = em->gate_cols;
+            break;
+        case 1:
+            stride = em->up_stride ? em->up_stride : em->expert_up_bytes;
+            type = em->up_type;
+            rows = em->up_rows;
+            cols = em->up_cols;
+            break;
+        case 2:
+            stride = em->down_stride ? em->down_stride : em->expert_down_bytes;
+            type = em->down_type;
+            rows = em->down_rows;
+            cols = em->down_cols;
+            break;
+        default:
+            return NULL;
+    }
+    if (stride != expert_bytes)
+        return NULL;
+    const uint8_t *base = bn_moe_mmap_base_for_proj(
+        bn_model_moe_io(model), em, proj);
+    if (!base)
+        return NULL;
+    size_t total_bytes = 0;
+    if (checked_mul_size(expert_bytes, (size_t)n_experts, &total_bytes) != 0)
+        return NULL;
+    if ((size_t)n_experts > (size_t)INT_MAX / (size_t)rows)
+        return NULL;
+    return gpu->buffer_create(gpu->ctx, base + offset, total_bytes,
+                              type, rows * n_experts, cols);
+}
+
 int bn_model_upload_weights(BnModel *model, BnGPUBackend *gpu) {
     if (!model || !gpu || !gpu->buffer_create) return -1;
     if (bn_model_ensure_backend(model) != 0) return -1;
@@ -144,6 +199,20 @@ int bn_model_upload_weights(BnModel *model, BnGPUBackend *gpu) {
                 (size_t)c->n_experts * (size_t)c->dim * sizeof(float),
                 BN_GGUF_TENSOR_F32, c->n_experts, c->dim)
             : NULL;
+        int upload_moe_all = lw->moe.router_weight &&
+                             getenv("BN_CUDA_ENABLE_MOE_ROUTED_FFN");
+        void *moe_gate_all_gpu = upload_moe_all
+            ? upload_moe_all_proj(model, gpu, &lw->moe.expert_map, 0,
+                                  c->n_experts)
+            : NULL;
+        void *moe_up_all_gpu = upload_moe_all
+            ? upload_moe_all_proj(model, gpu, &lw->moe.expert_map, 1,
+                                  c->n_experts)
+            : NULL;
+        void *moe_down_all_gpu = upload_moe_all
+            ? upload_moe_all_proj(model, gpu, &lw->moe.expert_map, 2,
+                                  c->n_experts)
+            : NULL;
         void *shared_expert_gate_gpu = lw->shared.shared_expert_gate
             ? gpu->buffer_create(
                 gpu->ctx, lw->shared.shared_expert_gate,
@@ -158,6 +227,12 @@ int bn_model_upload_weights(BnModel *model, BnGPUBackend *gpu) {
                                 moe_router_diff_gpu) != 0 ||
             register_gpu_handle(model, l, BN_BACKEND_HANDLE_MOE_ROUTER,
                                 moe_router_gpu) != 0 ||
+            register_gpu_handle(model, l, BN_BACKEND_HANDLE_MOE_GATE_ALL,
+                                moe_gate_all_gpu) != 0 ||
+            register_gpu_handle(model, l, BN_BACKEND_HANDLE_MOE_UP_ALL,
+                                moe_up_all_gpu) != 0 ||
+            register_gpu_handle(model, l, BN_BACKEND_HANDLE_MOE_DOWN_ALL,
+                                moe_down_all_gpu) != 0 ||
             register_gpu_handle(model, l, BN_BACKEND_HANDLE_SHARED_EXPERT_GATE,
                                 shared_expert_gate_gpu) != 0) {
             if (attn_norm_gpu) gpu->buffer_destroy(gpu->ctx, attn_norm_gpu);
@@ -166,6 +241,12 @@ int bn_model_upload_weights(BnModel *model, BnGPUBackend *gpu) {
                 gpu->buffer_destroy(gpu->ctx, moe_router_diff_gpu);
             if (moe_router_gpu)
                 gpu->buffer_destroy(gpu->ctx, moe_router_gpu);
+            if (moe_gate_all_gpu)
+                gpu->buffer_destroy(gpu->ctx, moe_gate_all_gpu);
+            if (moe_up_all_gpu)
+                gpu->buffer_destroy(gpu->ctx, moe_up_all_gpu);
+            if (moe_down_all_gpu)
+                gpu->buffer_destroy(gpu->ctx, moe_down_all_gpu);
             if (shared_expert_gate_gpu)
                 gpu->buffer_destroy(gpu->ctx, shared_expert_gate_gpu);
             bn_model_release_gpu(model);

src/transformer/gpu.c

@@ -939,6 +939,50 @@ static float *bn_transformer_gpu_forward_impl(BnModel *m, BnSession *sess,
                 backend, l, BN_BACKEND_HANDLE_MOE_ROUTER);
             int gpu_route_topk =
                 moe_router && !getenv("BN_CUDA_DISABLE_MOE_ROUTER_TOPK");
+            void *moe_gate_all = bn_backend_model_handle(
+                backend, l, BN_BACKEND_HANDLE_MOE_GATE_ALL);
+            void *moe_up_all = bn_backend_model_handle(
+                backend, l, BN_BACKEND_HANDLE_MOE_UP_ALL);
+            void *moe_down_all = bn_backend_model_handle(
+                backend, l, BN_BACKEND_HANDLE_MOE_DOWN_ALL);
+            int gpu_routed_ffn =
+                gpu_route_topk && moe_gate_all && moe_up_all && moe_down_all &&
+                getenv("BN_CUDA_ENABLE_MOE_ROUTED_FFN") &&
+                !c->has_shared_expert &&
+                lw->moe.expert_map.gate_type == BN_GGUF_TENSOR_Q4_K &&
+                lw->moe.expert_map.up_type == BN_GGUF_TENSOR_Q4_K &&
+                lw->moe.expert_map.down_type == BN_GGUF_TENSOR_Q6_K &&
+                lw->moe.expert_map.gate_rows == c->moe_intermediate_size &&
+                lw->moe.expert_map.up_rows == c->moe_intermediate_size &&
+                lw->moe.expert_map.gate_cols == dim &&
+                lw->moe.expert_map.up_cols == dim &&
+                lw->moe.expert_map.down_rows == dim &&
+                lw->moe.expert_map.down_cols == c->moe_intermediate_size &&
+                !getenv("BN_CUDA_DISABLE_MOE_ROUTED_FFN");
+            if (gpu_routed_ffn) {
+                if (bn_transformer_gpu_emit_context_moe_route_topk(
+                        &emit, moe_router, BN_GPU_VALUE_XB,
+                        BN_GPU_VALUE_MOE_HB, BN_GPU_VALUE_MOE_HB2,
+                        dim, c->n_experts, c->n_experts_active) != 0)
+                    return bn_transformer_gpu_reject_forward(
+                        &emit, "gpu moe route emit failed");
+                if (bn_transformer_gpu_emit_context_moe_routed_ffn(
+                        &emit, moe_gate_all, moe_up_all, moe_down_all,
+                        BN_GPU_VALUE_XB, BN_GPU_VALUE_MOE_HB2,
+                        BN_GPU_VALUE_MOE_HB, BN_GPU_VALUE_MOE_OUT,
+                        lw->moe.expert_map.gate_type,
+                        lw->moe.expert_map.down_type, dim,
+                        c->moe_intermediate_size, c->n_experts,
+                        c->n_experts_active) != 0)
+                    return bn_transformer_gpu_reject_forward(
+                        &emit, "gpu moe routed ffn emit failed");
+                bn_transformer_gpu_emit_context_residual_add(
+                    &emit, BN_GPU_VALUE_X, BN_GPU_VALUE_MOE_OUT, dim);
+                bn_transformer_gpu_emit_context_rmsnorm(
+                    &emit, next_norm, BN_GPU_VALUE_X, BN_GPU_VALUE_XB, dim,
+                    u_eps);
+                continue;
+            }
             int did_gpu_route_topk = 0;
             if (gpu_route_topk) {
                 if (bn_transformer_gpu_emit_context_moe_route_topk(