ggml-openvino: auto-enable full-MoE GPU path + dodge GPU rms_fusion bug

Two coupled changes that make gemma4 26B MoE run correctly on the GPU device
with no manual flag.

1. Auto-detect MoE and stop fragmenting the GPU graph. The per-node "force this
   MoE op to CPU on GPU" gates fragment the graph into dozens of submodels with
   cross-boundary copies (which mis-handle e.g. the layer-5 argsort indices and
   crash). GGML_OPENVINO_GPU_FULL_MOE kept the whole MoE on one OV submodel but
   had to be set by hand. Now ggml_openvino_gpu_full_moe_enabled() auto-enables
   it when running a MoE model on GPU: a GGML_OP_MUL_MAT_ID op (the expert-routed
   matmul, the defining op of a MoE model) latches a process-global flag from
   supports_op() at op-placement time. The scheduler queries placement before the
   expert weights are streamed in and makes several placement passes, so the first
   pass that sees MUL_MAT_ID sets the flag and later passes converge on the
   full-MoE layout. The GGML_OPENVINO_GPU_FULL_MOE env var still overrides
   (non-zero forces on, "0" forces off as an escape hatch). CPU/NPU behavior is
   unchanged: the gates are GPU-guarded, and the auto path only fires on GPU.

2. Dodge the GPU rms_fusion bug on that path. OpenVINO's rms_fusion folds
   Power(x, 2) -> ... into the internal RMS op; on the GPU plugin that fused RMS
   primitive's dynamic multi-token kernel writes only token 0 (tokens 1..N read
   back as 0). For gemma4 this collapsed the per-layer router RMSNorm (~7x summed
   over the prefill tokens), flattening the router softmax and flipping the top-8
   expert selection, so the GPU output drifted ("France" -> " only"). On the GPU
   full-MoE path only, compute the square as Multiply(x, x): algebraically
   identical, but it does not match the fusion pattern, so the GPU runs the
   unfused primitives and writes every token. Every other configuration keeps the
   fused fast path (CPU, NPU, and non-MoE GPU models such as Llama-3.2-1B).

Verified: gemma4 26B MoE on GPU with NO flag now matches CPU byte-for-byte on
prefill and multi-token decode; the GGML_OPENVINO_GPU_FULL_MOE=0 escape hatch
restores the old (fragmented) path; gemma4 CPU output unchanged; dense
Llama-3.2-1B on GPU still correct with rms_fusion active.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: cavusmustafa

ggml/src/ggml-openvino/ggml-openvino-extra.cpp

@@ -45,6 +45,7 @@ void ggml_openvino_device_config::init() {
         "GGML_OPENVINO_DISABLE_CACHE",
         "GGML_OPENVINO_DISABLE_KV_SLICE",
         "GGML_OPENVINO_MANUAL_GQA_ATTN",
+        "GGML_OPENVINO_GPU_FULL_MOE",
     };
 
     for (const char * const & env_var : env_var_names) {
@@ -173,6 +174,30 @@ bool ggml_openvino_is_npu() {
     return ggml_openvino_get_device_config().is_npu;
 }
 
+// Latched true once a MUL_MAT_ID op is seen during op placement; see header. Plain
+// non-atomic bool: placement runs single-threaded before the multi-threaded compute
+// that reads it, and the flag only ever transitions false->true (idempotent).
+static bool g_has_moe_expert_weights = false;
+
+void ggml_openvino_note_moe_expert_weight() {
+    g_has_moe_expert_weights = true;
+}
+
+bool ggml_openvino_has_moe_expert_weights() {
+    return g_has_moe_expert_weights;
+}
+
+bool ggml_openvino_gpu_full_moe_enabled() {
+    // Explicit env override (allowlisted): non-zero forces ON, "0" forces OFF.
+    if (const char * v = ggml_openvino_getenv_str("GGML_OPENVINO_GPU_FULL_MOE")) {
+        return std::atoi(v) != 0;
+    }
+    // Auto: keep the whole MoE on one OV submodel when running a quant-MoE model on
+    // GPU. On CPU/NPU the per-node gates are no-ops anyway (they are GPU-guarded), so
+    // leaving this OFF there preserves the existing behavior exactly.
+    return ggml_openvino_get_device_name() == "GPU" && ggml_openvino_has_moe_expert_weights();
+}
+
 // Get the remote context for the current device (returns empty optional for CPU)
 std::optional<ov::RemoteContext> ggml_openvino_get_remote_context() {
     return ggml_openvino_get_device_config().remote_context;

ggml/src/ggml-openvino/ggml-openvino-extra.h

@@ -99,6 +99,26 @@ int ggml_openvino_getenv_int(const char * var, int default_value = 0);
 // Check if running on NPU
 bool ggml_openvino_is_npu();
 
+// MoE detection. ggml_openvino_note_moe_expert_weight() latches a process-global flag
+// that ggml_openvino_has_moe_expert_weights() reports. It is called from supports_op()
+// the first time a GGML_OP_MUL_MAT_ID (the expert-routed matmul) is seen, which is the
+// defining op of a MoE model. The latch is set at op-placement time (not weight load):
+// the scheduler queries op placement before the expert weights are streamed in, and it
+// makes multiple placement passes, so the first pass that encounters MUL_MAT_ID sets the
+// flag and subsequent passes converge on the full-MoE layout. This lets the backend
+// recognize "this is a MoE model" without any architecture name.
+void ggml_openvino_note_moe_expert_weight();
+bool ggml_openvino_has_moe_expert_weights();
+
+// Whether to keep the whole MoE on one OV submodel instead of fragmenting at every
+// MoE node (see the per-node "force to CPU on GPU" gates). Resolution order:
+//   * GGML_OPENVINO_GPU_FULL_MOE set to non-zero -> force ON (any device)
+//   * GGML_OPENVINO_GPU_FULL_MOE set to "0"      -> force OFF (escape hatch)
+//   * unset -> AUTO: ON when running on GPU and the model has 3D quantized expert
+//     weights (a quant-MoE model), OFF otherwise.
+// CPU/NPU behavior is unchanged unless the env var is explicitly set.
+bool ggml_openvino_gpu_full_moe_enabled();
+
 // Get requantization type for a tensor type (returns nullopt if no requant needed)
 std::optional<ExtraQuantType> ggml_openvino_get_requant_type(const ggml_tensor * tensor, bool no_requant = false);
 

ggml/src/ggml-openvino/ggml-openvino.cpp

@@ -906,17 +906,17 @@ static bool checked_mul_size(size_t a, size_t b, size_t & out) {
     return true;
 }
 
-// When set (env GGML_OPENVINO_GPU_FULL_MOE), keep the entire MoE — including the routing
-// gather/softmax/argsort/normalization and the expert matmuls — on the OpenVINO device so
-// the whole model compiles as ONE submodel instead of fragmenting at every MoE node. The
-// per-node "force to CPU on GPU" gates below were added to work around GPU-plugin numerical
-// issues, but they fragment the graph into dozens of submodels with cross-boundary tensor
-// copies (which mis-handles e.g. the layer-5 argsort indices). With the dynamic-shape
-// frontend fix in place the un-fragmented graph is numerically correct, so this toggle lets
-// us run the whole MoE on one OV submodel.
+// Keep the entire MoE — including the routing gather/softmax/argsort/normalization and the
+// expert matmuls — on the OpenVINO device so the whole model compiles as ONE submodel instead
+// of fragmenting at every MoE node. The per-node "force to CPU on GPU" gates below were added
+// to work around GPU-plugin numerical issues, but they fragment the graph into dozens of
+// submodels with cross-boundary tensor copies (which mis-handles e.g. the layer-5 argsort
+// indices). With the dynamic-shape frontend fix in place the un-fragmented graph is
+// numerically correct, so this keeps the whole MoE on one OV submodel. Auto-enabled for
+// quant-MoE models on GPU; see ggml_openvino_gpu_full_moe_enabled() for the resolution order
+// and the GGML_OPENVINO_GPU_FULL_MOE override.
 static bool gpu_full_moe_enabled() {
-    static const bool v = getenv("GGML_OPENVINO_GPU_FULL_MOE") != nullptr;
-    return v;
+    return ggml_openvino_gpu_full_moe_enabled();
 }
 
 static bool mul_mat_id_requires_large_tmp(const ggml_tensor * op) {
@@ -1265,6 +1265,13 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {
 static bool ggml_backend_openvino_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) {
     GGML_ASSERT(dev->reg != nullptr);
 
+    // A MUL_MAT_ID op is the expert-routed matmul: its presence means this is a MoE
+    // model. Latch it here (placement time) rather than at weight load, because the
+    // scheduler queries op placement before the expert weights are streamed in.
+    if (op->op == GGML_OP_MUL_MAT_ID) {
+        ggml_openvino_note_moe_expert_weight();
+    }
+
     static std::unordered_set<ggml_type> supported_types{
         GGML_TYPE_F32,  GGML_TYPE_F16,  GGML_TYPE_BF16, GGML_TYPE_I64,  GGML_TYPE_I32,  GGML_TYPE_Q4_0,
         GGML_TYPE_Q4_1, GGML_TYPE_Q4_K, GGML_TYPE_Q5_1, GGML_TYPE_Q5_K, GGML_TYPE_Q8_0, GGML_TYPE_Q6_K};

ggml/src/ggml-openvino/openvino/op/rms_norm.cpp

@@ -1,6 +1,7 @@
 #include "../node_context.h"
 #include "../op_table.h"
 #include "../utils.h"
+#include "ggml-openvino/ggml-openvino-extra.h"
 
 #include <memory>
 #include <openvino/op/add.hpp>
@@ -20,8 +21,32 @@ OutputVector translate_rms_norm(const NodeContext & context) {
     num_inputs_check(context, 1, 1);
 
     auto input_node = process_view_input_new(context, 0);
-    auto square = std::make_shared<ov::op::v1::Power>(
-        input_node, ov::op::v0::Constant::create(ov::element::f32, ov::Shape{1}, {2.0f}));
+
+    // Build the mean-of-squares numerator. Normally use Power(x, 2): the OpenVINO
+    // rms_fusion pass matches that Power node and folds the whole decomposition into
+    // the internal RMS op (a perf win, e.g. dense Llama on GPU), so we keep it by
+    // default for every model and device.
+    //
+    // EXCEPTION — quant-MoE model on the GPU full-MoE path: the fused GPU RMS primitive's
+    // dynamic multi-token kernel writes only token 0 (tokens 1..N read back as 0). That
+    // silently collapses the per-layer MoE router RMSNorm summed over the prefill tokens
+    // (~7x), flattening the router softmax and flipping the top-8 expert selection, so the
+    // GPU output drifts from CPU (task #16). On that exact path only, compute the square as
+    // Multiply(x, x) — algebraically identical, but it does not match the rms_fusion
+    // pattern, so the GPU runs the unfused primitives and writes every token. Keyed to the
+    // same predicate as the full-MoE GPU path (auto-enabled for quant-MoE on GPU; see
+    // ggml_openvino_gpu_full_moe_enabled), so it never affects CPU/NPU or non-MoE GPU
+    // models, which keep the fused fast path.
+    static const bool dodge_rms_fusion =
+        ggml_openvino_get_device_name() == "GPU" && ggml_openvino_gpu_full_moe_enabled();
+
+    std::shared_ptr<ov::Node> square;
+    if (dodge_rms_fusion) {
+        square = std::make_shared<ov::op::v1::Multiply>(input_node, input_node);
+    } else {
+        square = std::make_shared<ov::op::v1::Power>(
+            input_node, ov::op::v0::Constant::create(ov::element::f32, ov::Shape{1}, {2.0f}));
+    }
 
     auto mean = std::make_shared<ov::op::v1::ReduceMean>(
         square, ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {-1}), true);