Stack CUDA shared MoE gate-up

Author: farkasmark

include/backend_model.h

@@ -36,6 +36,7 @@ typedef enum {
     BN_BACKEND_HANDLE_FFN_DOWN_PREFILL = 23,
     BN_BACKEND_HANDLE_SHARED_EXPERT_GATE = 24,
     BN_BACKEND_HANDLE_MOE_ROUTER_DIFF = 25,
+    BN_BACKEND_HANDLE_SHARED_GATEUP_STACKED = 26,
 } BnBackendHandleRole;
 
 BnBackendModel *bn_backend_model_create(void);

src/model_gpu.c

@@ -238,6 +238,18 @@ int bn_model_upload_weights(BnModel *model, BnGPUBackend *gpu) {
             return -1;
         }
 
+        void *shared_gateup_stacked_gpu =
+            bn_backend_layout_upload_stacked2(
+                gpu, &lw->shared.shared_gate, &lw->shared.shared_up);
+        if (register_gpu_handle(model, l,
+                                BN_BACKEND_HANDLE_SHARED_GATEUP_STACKED,
+                                shared_gateup_stacked_gpu) != 0) {
+            if (shared_gateup_stacked_gpu)
+                gpu->buffer_destroy(gpu->ctx, shared_gateup_stacked_gpu);
+            bn_model_release_gpu(model);
+            return -1;
+        }
+
         void *ssm_qkvz_stacked_gpu =
             bn_backend_layout_upload_stacked2(gpu, &lw->ssm.wqkv, &lw->ssm.wz);
         if (register_gpu_handle(model, l, BN_BACKEND_HANDLE_SSM_QKVZ_STACKED,

src/transformer/gpu_emit.c

@@ -1560,20 +1560,32 @@ void bn_transformer_gpu_emit_context_moe(BnTransformerGPUEmitContext *ctx,
     }
 
     if (lw->shared.shared_gate.data && shared && shared->shared_gate) {
-        uint32_t shared_gate_flags =
-            lw->shared.shared_gate.type == BN_GGUF_TENSOR_Q4_K ? 1u : 0u;
-        uint32_t shared_up_flags =
-            lw->shared.shared_up.type == BN_GGUF_TENSOR_Q4_K ? 1u : 0u;
-        emit_context_matvec_flags(
-            ctx, lw->shared.shared_gate.type,
-            shared->shared_gate,
-            BN_GPU_VALUE_XB, BN_GPU_VALUE_HB, lw->shared.shared_gate.rows,
-            lw->shared.shared_gate.cols, 0, shared_gate_flags);
-        emit_context_matvec_flags(
-            ctx, lw->shared.shared_up.type,
-            shared->shared_up,
-            BN_GPU_VALUE_XB, BN_GPU_VALUE_HB2, lw->shared.shared_up.rows,
-            lw->shared.shared_up.cols, 0, shared_up_flags);
+        if (shared->shared_gateup_stacked) {
+            emit_context_matvec_split(
+                ctx, lw->shared.shared_gate.type,
+                shared->shared_gateup_stacked,
+                BN_GPU_VALUE_XB, BN_GPU_VALUE_HB, BN_GPU_VALUE_HB2, -1,
+                lw->shared.shared_gate.rows + lw->shared.shared_up.rows,
+                lw->shared.shared_gate.cols, lw->shared.shared_gate.rows,
+                0, 0, 0, 0);
+        } else {
+            uint32_t shared_gate_flags =
+                lw->shared.shared_gate.type == BN_GGUF_TENSOR_Q4_K ? 1u : 0u;
+            uint32_t shared_up_flags =
+                lw->shared.shared_up.type == BN_GGUF_TENSOR_Q4_K ? 1u : 0u;
+            emit_context_matvec_flags(
+                ctx, lw->shared.shared_gate.type,
+                shared->shared_gate,
+                BN_GPU_VALUE_XB, BN_GPU_VALUE_HB,
+                lw->shared.shared_gate.rows, lw->shared.shared_gate.cols, 0,
+                shared_gate_flags);
+            emit_context_matvec_flags(
+                ctx, lw->shared.shared_up.type,
+                shared->shared_up,
+                BN_GPU_VALUE_XB, BN_GPU_VALUE_HB2,
+                lw->shared.shared_up.rows, lw->shared.shared_up.cols, 0,
+                shared_up_flags);
+        }
         bn_transformer_gpu_emit_context_activation(
             ctx, BN_GPU_VALUE_HB, BN_GPU_VALUE_HB2,
             lw->shared.shared_gate.rows, 0, BN_GPU_IR_ACTIVATION_SILU);

src/transformer/gpu_internal.h

@@ -77,6 +77,7 @@ typedef struct {
     void *shared_up;
     void *shared_down;
     void *shared_expert_gate;
+    void *shared_gateup_stacked;
 } BnTransformerGPUMoESharedResources;
 
 typedef struct {

src/transformer/gpu_resources.c

@@ -177,5 +177,7 @@ bn_transformer_gpu_resolve_moe_shared_resources(
         .shared_down = qweight_backend_buf(backend, &lw->shared.shared_down),
         .shared_expert_gate = backend_handle_or(
             backend, layer, BN_BACKEND_HANDLE_SHARED_EXPERT_GATE),
+        .shared_gateup_stacked = backend_handle_or(
+            backend, layer, BN_BACKEND_HANDLE_SHARED_GATEUP_STACKED),
     };
 }