diff --git a/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.cpp b/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.cpp index 5af19d54c9fc..8c8059f24c10 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.cpp @@ -488,14 +488,25 @@ void apply_moe_config(ov::AnyMap& stage_config, const std::string& stage_name) { if (moe_hint == ::intel_npu::npuw::llm::MoEHint::HOST_ROUTED) { LOG_INFO("MoE config for " << stage_name << " stage: HOST_ROUTED (host-side expert routing)"); - // MoE expert and router pattern isolation options + // MoE expert and router pattern isolation options. + // Note: NPUW_ONLINE_ISOLATE is handled separately below to support coexistence + // with attention isolation (e.g. "ATTN" already set by dyn_attn_opts). A plain + // merge_config_with would overwrite the existing value and silently drop ATTN. const ov::AnyMap expert_opts = { {"NPUW_ONLINE_PIPELINE", "REP"}, - {"NPUW_ONLINE_ISOLATE", "MOE"}, {"NPUW_ONLINE_KEEP_BLOCK_SIZE", "4"}, {"NPUW_UNFOLD_IREQS", "NO"}, }; merge_config_with(stage_config, expert_opts); + // Append "MOE" to any pre-existing isolation preset (e.g. "ATTN") rather than + // overwriting it. getIsolates() accepts a comma-separated list of presets. + auto isol_it = stage_config.find("NPUW_ONLINE_ISOLATE"); + if (isol_it != stage_config.end() && !isol_it->second.as().empty()) { + isol_it->second = isol_it->second.as() + ",MOE"; + LOG_INFO("MoE config: appended MOE to NPUW_ONLINE_ISOLATE -> " << isol_it->second.as()); + } else { + stage_config["NPUW_ONLINE_ISOLATE"] = "MOE"; + } } else if (moe_hint == ::intel_npu::npuw::llm::MoEHint::DEVICE_ROUTED) { if (stage_name == "PREFILL") { NPUW_ASSERT(false && "MoE DEVICE_ROUTED is not supported for PREFILL stage. " @@ -1194,6 +1205,21 @@ ov::npuw::LLMCompiledModel::LLMCompiledModel(const std::shared_ptr& m } // Compile multiple generate model variants with different sizes + // DBG: dump final effective config keys relevant to online partitioning / attention isolation + { + auto dump_key = [](const ov::AnyMap& cfg, const std::string& key, const std::string& label) { + auto it = cfg.find(key); + std::cout << "[LLM_CFG] " << label << " " << key << " = " + << (it != cfg.end() ? it->second.as() : "(not set)") << std::endl; + }; + dump_key(prefill_config, "NPUW_ONLINE_PIPELINE", "PREFILL"); + dump_key(prefill_config, "NPUW_ONLINE_ISOLATE", "PREFILL"); + dump_key(prefill_config, "NPUW_ATTN", "PREFILL"); + dump_key(generate_config, "NPUW_ONLINE_PIPELINE", "GENERATE"); + dump_key(generate_config, "NPUW_ONLINE_ISOLATE", "GENERATE"); + dump_key(generate_config, "NPUW_ATTN", "GENERATE"); + dump_key(generate_config, "NPUW_UNFOLD_IREQS", "GENERATE"); + } compile_generate_model_variants(generate_model_variants, plugin, generate_config); m_prefill_compiled = m_compiled_model_factory(prefill_model, plugin, prefill_config); diff --git a/src/plugins/intel_npu/src/plugin/npuw/moe/moe_executor.cpp b/src/plugins/intel_npu/src/plugin/npuw/moe/moe_executor.cpp index c307f6bda647..81ecc8acbe11 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/moe/moe_executor.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/moe/moe_executor.cpp @@ -4,6 +4,7 @@ #include "moe_executor.hpp" +#include #include #include "../compiled_model.hpp" // For CompiledModel::CompiledModelDesc @@ -347,9 +348,17 @@ void MoEExecutor::run_expert_iterative(size_t idx) { auto expert_input_source = io.expert_input; - // Output embedding dimension - const auto output_shape = m_config.compiled_models.begin()->second->outputs()[0].get_shape(); - const size_t embed_dim = (output_shape.size() == 4) ? output_shape[3] : output_shape[1]; + // Use the embed_dim already validated against the accumulator buffer during prepare(). + // Defensively assert that the compiled model's output last-dim agrees, so any future + // layout change (2D/3D/4D) is caught immediately at runtime instead of producing + // silently wrong scatter results. + const size_t embed_dim = m_config.expert_hidden_dim; + { + const auto output_shape = m_config.compiled_models.begin()->second->outputs()[0].get_shape(); + NPUW_ASSERT(!output_shape.empty() && output_shape.back() == embed_dim && + "Chunk model output last-dim does not match expert_hidden_dim — " + "layout may have changed, update embed_dim derivation"); + } // num_tokens and num_experts come from config, validated once during prepare(). // The router tensor's token dimension is guaranteed to match input_token_count because @@ -676,32 +685,20 @@ void MoEExecutor::unpack_single_expert_closure(std::size_t idx, RqPtr request, s // Slice expert weight using view (no copy) - returns ov::Tensor object auto sliced_weight_tensor = ov::npuw::moe::slice_expert_weight(closure, expert_id, num_experts); - // Get impl pointer for use in unpacking/setting + // Get impl pointer for use in setting auto sliced_weight = ov::get_tensor_impl(sliced_weight_tensor); - // Handle unpacking if needed - if (m_accessor.unpack_required(idx, cidx)) { - auto clparam = request->get_tensor(iport); + // Unpack is not expected for MoE sliced weights — assert to catch if this changes. + NPUW_ASSERT(!m_accessor.unpack_required(idx, cidx) && + "EXPERT_ITERATIVE: unpack of sliced MoE weight is not supported"); - if (!comp_model_desc.scales.empty() && comp_model_desc.scales[cidx] && comp_model_desc.zerops[cidx]) { - ov::npuw::util::unpack(sliced_weight, - ov::get_tensor_impl(comp_model_desc.zerops[cidx]), - ov::get_tensor_impl(comp_model_desc.scales[cidx]), - clparam); - } else if (!comp_model_desc.scales.empty() && comp_model_desc.scales[cidx]) { - ov::npuw::util::unpack(sliced_weight, ov::get_tensor_impl(comp_model_desc.scales[cidx]), clparam); - } else { - ov::npuw::util::unpack(sliced_weight, clparam); - } + // Direct set (no unpacking needed) + // When cache is enabled: Use direct set_tensor to avoid polluting shared input tensors + // When cache is disabled: Use set_tensor_optimized for better performance (copies small tensors) + if (m_resources.request_cache) { + request->set_tensor(iport, sliced_weight); } else { - // Direct set (no unpacking needed) - // When cache is enabled: Use direct set_tensor to avoid polluting shared input tensors - // When cache is disabled: Use set_tensor_optimized for better performance (copies small tensors) - if (m_resources.request_cache) { - request->set_tensor(iport, sliced_weight); - } else { - ov::npuw::moe::set_tensor_optimized(request, iport, sliced_weight); - } + ov::npuw::moe::set_tensor_optimized(request, iport, sliced_weight); } } else { // This closure parameter doesn't need slicing, use original logic @@ -767,18 +764,24 @@ void MoEExecutor::unpack_multiple_experts_closure(std::size_t idx, // Calculate original parameter index in the model const size_t original_param_idx = comp_model_desc.param_base + closure_idx; + auto& batched_closure = desc_closure[closure_idx]; + const auto& closure_shape = batched_closure.get_shape(); + // Check if this parameter has unrolled variants (is in param_mapping) auto mapping_it = param_mapping.find(original_param_idx); if (mapping_it == param_mapping.end()) { - continue; // Not unrolled + // If the closure is batched [num_experts,...] but has no param_mapping entry, the + // expert-dimension unrolling (UnrollMoEMatMul) failed to fire for this weight. + // This is a transformation bug — assert to surface it early. + const bool is_batched_nounroll = !closure_shape.empty() && closure_shape[0] == num_experts; + NPUW_ASSERT(!is_batched_nounroll && + "Batched closure has no param_mapping entry — UnrollMoEMatMul did not fire"); + continue; // Non-batched shared param not in mapping is expected — skip silently } const auto& unrolled_indices = mapping_it->second; NPUW_ASSERT(unrolled_indices.size() == K); - auto& batched_closure = desc_closure[closure_idx]; - const auto& closure_shape = batched_closure.get_shape(); - // Verify this is a batched parameter [num_experts, ...] const bool is_batched = !closure_shape.empty() && closure_shape[0] == num_experts; if (!is_batched) { @@ -799,53 +802,27 @@ void MoEExecutor::unpack_multiple_experts_closure(std::size_t idx, // ========== Step 3: Process batched parameters (K experts) ========== - // Pre-determine unpack configuration (same for all K experts) + // Unpack is not expected for MoE unrolled weights — assert to catch if this changes. const auto batched_elem_type = batched_closure.get_element_type(); const auto target_elem_type = request->get_tensor(compiled_inputs[unrolled_indices[0]])->get_element_type(); - const bool needs_unpack = (batched_elem_type != target_elem_type); - - ov::SoPtr scales_impl, zerops_impl; - if (needs_unpack) { - if (!comp_model_desc.scales.empty() && comp_model_desc.scales[closure_idx]) { - scales_impl = ov::get_tensor_impl(comp_model_desc.scales[closure_idx]); - } - if (!comp_model_desc.zerops.empty() && comp_model_desc.zerops[closure_idx]) { - zerops_impl = ov::get_tensor_impl(comp_model_desc.zerops[closure_idx]); - } - } + NPUW_ASSERT(batched_elem_type == target_elem_type && + "EXPERT_BATCH: dtype mismatch in MoE closure, unpack path is not supported"); - // Process K experts + // Process K experts (direct set, no unpacking) for (size_t position = 0; position < K; ++position) { const size_t expert_id = expert_ids[position]; const auto& iport = compiled_inputs[unrolled_indices[position]]; - // Slice expert weight (zero-copy view) + // Slice expert weight (zero-copy view) and bind directly to the request port. ov::Tensor sliced_expert = ov::npuw::moe::slice_expert_weight(batched_closure, expert_id, num_experts); + auto sliced_impl = ov::get_tensor_impl(sliced_expert); - if (needs_unpack) { - // Unpack path (dtype mismatch) - auto sliced_impl = ov::get_tensor_impl(sliced_expert); - auto clparam = request->get_tensor(iport); - - if (scales_impl && zerops_impl) { - ov::npuw::util::unpack(sliced_impl, zerops_impl, scales_impl, clparam); - } else if (scales_impl) { - ov::npuw::util::unpack(sliced_impl, scales_impl, clparam); - } else if (zerops_impl) { - ov::npuw::util::unpack(sliced_impl, zerops_impl, clparam); - } else { - ov::npuw::util::unpack(sliced_impl, clparam); - } + // When cache is enabled: use set_tensor to avoid polluting shared input tensors. + // When cache is disabled: use set_tensor_optimized (copies small tensors, faster). + if (m_resources.request_cache) { + request->set_tensor(iport, sliced_impl); } else { - auto sliced_impl = ov::get_tensor_impl(sliced_expert); - // Direct set (no unpacking needed) - // When cache is enabled: Use direct set_tensor to avoid polluting shared input tensors - // When cache is disabled: Use set_tensor_optimized for better performance (copies small tensors) - if (m_resources.request_cache) { - request->set_tensor(iport, sliced_impl); - } else { - ov::npuw::moe::set_tensor_optimized(request, iport, sliced_impl); - } + ov::npuw::moe::set_tensor_optimized(request, iport, sliced_impl); } } // for each expert } // for each closure parameter diff --git a/src/plugins/intel_npu/src/plugin/npuw/moe/moe_infer_utils.cpp b/src/plugins/intel_npu/src/plugin/npuw/moe/moe_infer_utils.cpp index 789cdc253a67..f6f745d52ef7 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/moe/moe_infer_utils.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/moe/moe_infer_utils.cpp @@ -72,11 +72,16 @@ namespace { // Layout B: [num_experts, token_num, 1, 1] → returns shape[1] // ASSERTs if neither layout is recognised. static size_t get_router_token_count(const ov::Shape& router_shape) { + if (router_shape.size() == 3) { + // 3-D [num_experts, token_count, 1] (Qwen3 prefill) + NPUW_ASSERT(router_shape[2] == 1 && "Unexpected 3-D router shape: trailing dim must be 1"); + return router_shape[1]; + } NPUW_ASSERT(router_shape.size() == 4); if (router_shape[1] == 1 && router_shape[3] == 1) - return router_shape[2]; // Layout A + return router_shape[2]; // Layout A: [num_experts, 1, token_count, 1] if (router_shape[2] == 1 && router_shape[3] == 1) - return router_shape[1]; // Layout B + return router_shape[1]; // Layout B: [num_experts, token_count, 1, 1] NPUW_ASSERT(false && "Unexpected router output shape - cannot determine token dimension!"); return 0; // unreachable, suppress warning } @@ -95,7 +100,7 @@ std::vector parse_selected_experts_from_router(const ov::SoPtrget_shape(); - if (shape.size() != 4 || shape[0] != num_experts) { + if ((shape.size() != 4 && shape.size() != 3) || shape[0] != num_experts) { NPUW_ASSERT(false && "Unexpected router output shape!"); } @@ -168,8 +173,10 @@ void gather_router_scores(const ov::SoPtr& router_source, size_t expert_offset; if (router_source_shape.size() == 4) { expert_offset = expert_id * get_router_token_count(router_source_shape); - } else if (router_source_shape.size() == 2) { - expert_offset = expert_id * router_source_shape[1]; // [num_experts, token_num] + } else if (router_source_shape.size() == 3 || router_source_shape.size() == 2) { + // 3-D [num_experts, token_num, 1] or 2-D [num_experts, token_num] + // trailing singleton dim (if any) does not affect linear stride between experts + expert_offset = expert_id * router_source_shape[1]; } else { NPUW_ASSERT(false && "Unexpected router source shape"); expert_offset = 0; // unreachable, suppress uninitialized warning diff --git a/src/plugins/intel_npu/src/plugin/npuw/moe_transformations/device_routed_moe_transform.cpp b/src/plugins/intel_npu/src/plugin/npuw/moe_transformations/device_routed_moe_transform.cpp index 1f91cc23cf72..8f46a918c531 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/moe_transformations/device_routed_moe_transform.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/moe_transformations/device_routed_moe_transform.cpp @@ -288,7 +288,6 @@ LayerNodes collect_from_expert_output(const RouterInfo& router) { } if (nodes.num_experts == 0) { - LOG_WARN("collect_from_expert_output: no Tile with repeats[0] > k_value found"); return nodes; } @@ -298,6 +297,13 @@ LayerNodes collect_from_expert_output(const RouterInfo& router) { std::dynamic_pointer_cast(reshape->input_value(1).get_node_shared_ptr()); auto shape_data = shape_const->cast_vector(); if (!shape_data.empty() && shape_data[0] == static_cast(nodes.num_experts)) { + // Skip group-quant weight-chain Reshapes: Multiply(weight, scale) -> Reshape -> Convert -> MatMul. + // These must NOT be modified by transform_constant_reshapes (the Gather inserted by + // transform_matmuls on the final Convert output handles the expert slicing instead). + if (std::dynamic_pointer_cast( + reshape->input_value(0).get_node_shared_ptr())) { + continue; + } nodes.constant_reshapes.push_back(reshape); } } @@ -305,7 +311,7 @@ LayerNodes collect_from_expert_output(const RouterInfo& router) { nodes.dynamic_reshapes = std::move(all_dyn_reshapes); // Returns true if a constant somewhere in the weight chain has shape[0] == expert_dim. - // Recursively peels Convert and both inputs of Multiply (the two recognized chain patterns). + // Recursively peels Convert, Multiply, and Reshape (group-quant: Multiply→Reshape→Convert). // Returns false for any unrecognized node type, signalling an unknown weight chain. std::function&, size_t, bool&)> weight_has_expert_dim; weight_has_expert_dim = [&](const ov::Output& input, size_t expert_dim, bool& chain_recognized) -> bool { @@ -320,7 +326,12 @@ LayerNodes collect_from_expert_output(const RouterInfo& router) { return weight_has_expert_dim(mul->input_value(0), expert_dim, chain_recognized) || weight_has_expert_dim(mul->input_value(1), expert_dim, chain_recognized); } - // Unrecognized node in weight chain (e.g. Reshape, Gather, custom op). + if (auto reshape = std::dynamic_pointer_cast(node)) { + // Group-quant: Multiply(weight, scale) -> Reshape([N,O,G*Gs]) -> Convert -> MatMul + // Peel through the reshape to find the expert dimension in the upstream Multiply. + return weight_has_expert_dim(reshape->input_value(0), expert_dim, chain_recognized); + } + // Unrecognized node in weight chain (e.g. Gather, custom op). chain_recognized = false; return false; }; @@ -645,8 +656,11 @@ bool DeviceRoutedMoETransform::run_on_model(const std::shared_ptr& mo // Step 1: Detect router by topology (name-independent, works for GPT-OSS and Qwen3) auto router = detect_router_by_topology(node); - if (!router.has_value()) + if (!router.has_value()) { + LOG_DEBUG("DeviceRoutedMoETransform: Scatter node '" << node->get_friendly_name() + << "' does not match expected MoE router topology; skipping"); continue; + } // Step 2: Collect expert nodes via backward BFS from output_multiply auto layer_nodes = collect_from_expert_output(router.value()); diff --git a/src/plugins/intel_npu/src/plugin/npuw/moe_transformations/moe_transformation.cpp b/src/plugins/intel_npu/src/plugin/npuw/moe_transformations/moe_transformation.cpp index d6102c19da43..e2f8d0127b7c 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/moe_transformations/moe_transformation.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/moe_transformations/moe_transformation.cpp @@ -389,27 +389,27 @@ std::optional detect_and_transform_moe_downstream(const std::shar // Looking for a parameter with shape [N, ..., W] where: // dim 0 = N (total experts, > 1) // dim n-1 = W (hidden_dim) - // at least one of dims 1..n-2 equals 1 (singleton "batch" dimension) - // Both [N, 1, H, W] (GPT-OSS) and [N, H, 1, W] (Qwen) are accepted. + // Accepted layouts: + // 4-D: [N, 1, H, W] (GPT-OSS) or [N, H, 1, W] (Qwen, decoding) + // 3-D: [N, token, W] (Qwen, prefill — token may be > 1) + // The primary discriminator is check_downstream_pattern() (Parameter -> [Convert] -> ReduceSum). const auto& params = model->get_parameters(); for (size_t param_idx = 0; param_idx < params.size(); ++param_idx) { const auto& param = params[param_idx]; - // Validate shape: must be 4-D with dim 0 > 1 + // Validate shape: must be 3-D or 4-D with dim 0 > 1 auto param_shape = param->get_partial_shape(); - if (!param_shape.rank().is_static() || param_shape.rank().get_length() != 4) { + if (!param_shape.rank().is_static()) { continue; } - - auto shape = param_shape.to_shape(); - if (shape[0] <= 1) { + const auto rank = param_shape.rank().get_length(); + if (rank < 3 || rank > 4) { continue; } - // At least one of dims 1..2 must be 1 (singleton expansion dim) - bool has_singleton_dim = (shape[1] == 1 || shape[2] == 1); - if (!has_singleton_dim) { + auto shape = param_shape.to_shape(); + if (shape[0] <= 1) { continue; } @@ -468,6 +468,7 @@ std::optional create_moe_downstream(const std::shared_ptris_valid()) { @@ -481,6 +482,7 @@ std::optional create_moe_downstream(const std::shared_ptr MoEExperts::from(const std::shared_ptr& mod // For EXPERT_BATCH mode (K experts), unrolling creates the same mapping structure auto param_mapping = build_parameter_mapping_from_rtinfo(model, transformed_models.begin()->second); + // Step 5.5: Compile-time guard — verify UnrollMoEMatMul fired for every batched parameter. + // + // In EXPERT_BATCH mode every original parameter with shape[0] == num_experts MUST appear in + // param_mapping with exactly K unrolled entries. A missing entry means the transformation did + // not recognise the weight chain (e.g., a new quantisation pattern such as GPTQ introducing an + // extra Reshape node) and inference would produce silently wrong results at runtime. + // + // Failing loudly here — at model-load time — surfaces the bug immediately instead of letting + // it silently corrupt outputs. The symmetric runtime guard in unpack_multiple_experts_closure + // serves as a second line of defence (e.g., for models loaded from a serialised cache that + // bypasses this code path). + if (structure_info->is_expert_batch_mode()) { + const auto& orig_params = model->get_parameters(); + for (size_t pi = 0; pi < orig_params.size(); ++pi) { + const auto& pshape = orig_params[pi]->get_partial_shape(); + if (!pshape.rank().is_static() || !pshape[0].is_static()) + continue; + if (static_cast(pshape[0].get_length()) != structure_info->num_experts) + continue; + // This parameter is batched over the expert dimension and must be unrolled. + auto it = param_mapping.find(pi); + if (it == param_mapping.end()) { + OPENVINO_THROW("MoE EXPERT_BATCH: parameter '", + orig_params[pi]->get_friendly_name(), + "' (index=", + pi, + ", shape=", + pshape, + ") has shape[0]==num_experts but was not unrolled. " + "UnrollMoEMatMul did not recognise its weight chain. " + "Inspect the weight graph topology and add the missing chain variant."); + } + if (it->second.size() != static_cast(k_value)) { + OPENVINO_THROW("MoE EXPERT_BATCH: parameter '", + orig_params[pi]->get_friendly_name(), + "' was unrolled into ", + it->second.size(), + " entries but expected K=", + k_value); + } + } + LOG_DEBUG("Compile-time unroll validation passed: all " + << param_mapping.size() << " batched parameters are correctly mapped (K=" << k_value << ")"); + } + // Step 6: Populate and validate MoEExperts structure MoEExperts moe_experts; moe_experts._num_experts = structure_info->num_experts; diff --git a/src/plugins/intel_npu/src/plugin/npuw/moe_transformations/moe_unroll_patterns.cpp b/src/plugins/intel_npu/src/plugin/npuw/moe_transformations/moe_unroll_patterns.cpp index 3bece991c8d5..34742c4f0a4d 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/moe_transformations/moe_unroll_patterns.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/moe_transformations/moe_unroll_patterns.cpp @@ -73,6 +73,7 @@ struct ExpertBranchContext { ov::Output scale_param_source; ov::Output weights_param_source; std::shared_ptr multiply_node; + std::shared_ptr reshape_after_multiply; // group-quant: Multiply→Reshape→Convert std::shared_ptr convert_after_multiply; std::shared_ptr matmul; ov::Shape scale_new_shape; @@ -140,15 +141,34 @@ inline ov::Output create_expert_branch_weights(const ExpertBranchConte new_multiply->set_friendly_name(ctx.multiply_node->get_friendly_name() + "/expert_" + std::to_string(ctx.expert_idx)); - // 6. Convert after Multiply if needed + ov::Output after_multiply = new_multiply->output(0); + + // 5.5. Reshape after Multiply if present (group-quant: Multiply→Reshape→Convert→MatMul). + // The original Reshape changes [N, out, num_groups, group_size] → [N, out, num_groups*group_size]. + // For the per-expert branch, change dim 0 from N to 1. + if (ctx.reshape_after_multiply) { + auto orig_shape_const = std::dynamic_pointer_cast( + ctx.reshape_after_multiply->input_value(1).get_node_shared_ptr()); + OPENVINO_ASSERT(orig_shape_const, "Reshape shape input must be Constant"); + auto shape_data = orig_shape_const->cast_vector(); + shape_data[0] = 1; // per-expert slice + auto new_shape_const = + ov::op::v0::Constant::create(ov::element::i64, orig_shape_const->get_shape(), shape_data); + auto new_reshape = std::make_shared(after_multiply, new_shape_const, false); + new_reshape->set_friendly_name(ctx.reshape_after_multiply->get_friendly_name() + "/expert_" + + std::to_string(ctx.expert_idx)); + after_multiply = new_reshape->output(0); + } + + // 6. Convert after Multiply (or Reshape) if needed if (ctx.convert_after_multiply) { auto new_convert_after_multiply = - std::make_shared(new_multiply, ctx.convert_after_multiply->get_destination_type()); + std::make_shared(after_multiply, ctx.convert_after_multiply->get_destination_type()); new_convert_after_multiply->set_friendly_name(ctx.convert_after_multiply->get_friendly_name() + "/expert_" + std::to_string(ctx.expert_idx)); return new_convert_after_multiply->output(0); } - return new_multiply->output(0); + return after_multiply; } /** @@ -267,9 +287,19 @@ UnrollMoEMatMul::UnrollMoEMatMul(std::shared_ptr model) : model_(mode std::shared_ptr convert_after_multiply; std::shared_ptr multiply_node; + std::shared_ptr reshape_after_multiply; if (auto conv = std::dynamic_pointer_cast(input1_node)) { convert_after_multiply = conv; - multiply_node = std::dynamic_pointer_cast(conv->input_value(0).get_node_shared_ptr()); + auto inner = conv->input_value(0).get_node_shared_ptr(); + // Group-quant chain: Multiply → Reshape → Convert → MatMul + // Peel through the optional Reshape between Multiply and the outer Convert. + if (auto reshape = std::dynamic_pointer_cast(inner)) { + reshape_after_multiply = reshape; + multiply_node = std::dynamic_pointer_cast( + reshape->input_value(0).get_node_shared_ptr()); + } else { + multiply_node = std::dynamic_pointer_cast(inner); + } } else { multiply_node = std::dynamic_pointer_cast(input1_node); } @@ -374,6 +404,7 @@ UnrollMoEMatMul::UnrollMoEMatMul(std::shared_ptr model) : model_(mode scale_param_source, weights_param_source, multiply_node, + reshape_after_multiply, convert_after_multiply, matmul, scale_new_shape, diff --git a/src/plugins/intel_npu/src/plugin/npuw/partitioning/patterns/moe.cpp b/src/plugins/intel_npu/src/plugin/npuw/partitioning/patterns/moe.cpp index 11dac1004c4e..b6014446ee13 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/partitioning/patterns/moe.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/partitioning/patterns/moe.cpp @@ -75,6 +75,26 @@ void isolate_reduce_sum_after(const std::shared_ptr& output_multiply, } } +// Weight dequantization chain used by Qwen3Expert projections and Qwen3Router. +// Supports both regular and group-quantized weight layouts via opp::optional: +// Regular quant: [Convert_in?] -> Multiply(weight, scale) -> [Reshape?] -> Convert_out -> MatMul +// Group quant: Convert_in -> Multiply(weight, scale) -> Reshape -> Convert_out -> MatMul +struct WeightChainPattern { + std::shared_ptr convert_in; // optional: Convert(I4->FP) for INT4/group quant + std::shared_ptr multiply; // required: Multiply(weight, scale) + std::shared_ptr reshape; // optional: Reshape(3D->2D) for group quant + std::shared_ptr convert_out; // required: final Convert before MatMul +}; + +static WeightChainPattern make_weight_chain() { + WeightChainPattern wc; + wc.convert_in = opp::optional({opp::any_input()}); + wc.multiply = opp::wrap_type({wc.convert_in, opp::any_input()}); + wc.reshape = opp::optional({wc.multiply, opp::any_input()}); + wc.convert_out = opp::wrap_type({wc.reshape}); + return wc; +} + // Extract K from a TopK node's constant second input and write it to rt_info // under the RT_INFO_MOE_K key so that PartitioningCallbacks::find_node_with_rt_info // can retrieve it during the partition stage. Returns true on success. @@ -327,25 +347,32 @@ GPTOSSRouter::GPTOSSRouter([[maybe_unused]] const std::shared_ptr Reshape1 + Each projection (gate/up/down) uses a WeightChainPattern (gw/uw/dw) that handles + both regular and group-quantized weight dequantization transparently: + Regular quant: [Convert_in?] -> Multiply(weight, scale) -> [Reshape?] -> Convert_out -> MatMul + Group quant: Convert_in -> Multiply(weight, scale) -> Reshape -> Convert_out -> MatMul + Gate projection (SwiGLU gate branch): - Reshape1 -> MatMul_gate (with weights: Multiply -> Convert) -> Swish + Reshape1 -> [gw] -> MatMul_gate -> Swish Up projection (SwiGLU up branch): - Reshape1 -> MatMul_up (with weights: Multiply -> Convert) + Reshape1 -> [uw] -> MatMul_up SwiGLU merge: Swish + MatMul_up -> Multiply_swiglu Down projection: - Multiply_swiglu -> MatMul_down (with weights: Multiply -> Convert) -> Reshape2 + Multiply_swiglu -> [dw] -> MatMul_down -> Reshape2 Output (scaled by router scores): Reshape2 * router_score -> Multiply_output <-- pattern root (router_score = opp::any_input(), produced entirely by Qwen3Router) - Expert isolates: Tile, Reshape1, weight-dequant nodes, MatMuls, SwiGLU Multiply, Reshape2, output Multiply. - Router isolates: Softmax, TopK, ReduceSum, Divide, Scatter, Transpose, Reshape, Unsqueeze. - Shape-compute chains (ShapeOf->Gather->...) stay outside both as subgraph parameters. + Isolation boundary: + - Expert claims: Tile, Reshape1, all weight-dequant nodes, MatMuls, SwiGLU Multiply, Reshape2, output Multiply + - Router claims: Softmax, TopK, ReduceSum, Divide, ScatterElementsUpdate, Transpose, Reshape_score, Unsqueeze_score + - Shared shape-compute nodes (ShapeOf->Gather->Unsqueeze->Concat chains) stay outside both, + becoming parameter inputs at subgraph boundaries. */ Qwen3Expert::Qwen3Expert(const std::shared_ptr& snapshot, const std::string& isol_tag) { LOG_DEBUG("Qwen3Expert pattern matcher registered with tag: " << isol_tag); @@ -354,32 +381,27 @@ Qwen3Expert::Qwen3Expert(const std::shared_ptr& snap auto tile = opp::wrap_type({opp::any_input(), opp::any_input()}); auto reshape1 = opp::wrap_type({tile, opp::any_input()}); - // Gate projection weights: Multiply(quantized weight, scale) -> Convert - auto gate_weights_multiply = opp::wrap_type({opp::any_input(), opp::any_input()}); - auto gate_weights_convert = opp::wrap_type({gate_weights_multiply}); - // Gate MatMul + Swish activation - auto matmul_gate = opp::wrap_type({reshape1, gate_weights_convert}); + // Per-projection weight chains (regular quant and group quant both handled via opp::optional) + auto gw = make_weight_chain(); // gate weights + auto uw = make_weight_chain(); // up weights + auto dw = make_weight_chain(); // down weights + + // Gate projection: reshape1 -> [gw] -> MatMul_gate -> Swish + auto matmul_gate = opp::wrap_type({reshape1, gw.convert_out}); auto swish = opp::wrap_type({matmul_gate}); - // Up projection weights: Multiply(quantized weight, scale) -> Convert - auto up_weights_multiply = opp::wrap_type({opp::any_input(), opp::any_input()}); - auto up_weights_convert = opp::wrap_type({up_weights_multiply}); - // Up MatMul - auto matmul_up = opp::wrap_type({reshape1, up_weights_convert}); + // Up projection: reshape1 -> [uw] -> MatMul_up + auto matmul_up = opp::wrap_type({reshape1, uw.convert_out}); - // SwiGLU: gate * up + // SwiGLU merge: gate * up auto multiply_swiglu = opp::wrap_type({swish, matmul_up}); - // Down projection weights: Multiply(quantized weight, scale) -> Convert - auto down_weights_multiply = opp::wrap_type({opp::any_input(), opp::any_input()}); - auto down_weights_convert = opp::wrap_type({down_weights_multiply}); - // Down MatMul -> Reshape - auto matmul_down = opp::wrap_type({multiply_swiglu, down_weights_convert}); + // Down projection: multiply_swiglu -> [dw] -> MatMul_down -> Reshape2 + auto matmul_down = opp::wrap_type({multiply_swiglu, dw.convert_out}); auto reshape2 = opp::wrap_type({matmul_down, opp::any_input()}); // Pattern root: expert_output * router_score - // The router score (Unsqueeze output) is produced entirely by Qwen3Router and flows - // in as opp::any_input() here to avoid double-claiming shared nodes. + // The router score (Unsqueeze output) is produced entirely by Qwen3Router. auto output_multiply = opp::wrap_type({reshape2, opp::any_input()}); auto node_to_gptr = snapshot->getNodeToGroupMap(); @@ -391,27 +413,40 @@ Qwen3Expert::Qwen3Expert(const std::shared_ptr& snap auto matched_output_multiply = node_to_output.at(output_multiply).get_node_shared_ptr(); LOG_DEBUG("Qwen3Expert pattern matched: " << matched_tile->get_friendly_name()); - - LOG_DEBUG("Qwen3 Expert Multiply output_shape: " << matched_output_multiply->get_output_partial_shape(0)); const bool is_decoding = is_decoding_stage(matched_output_multiply); LOG_DEBUG("Qwen3 Expert pattern matched (" << (is_decoding ? "Decoding" : "Prefill") << " stage)"); - auto isolate = [&](const std::shared_ptr& pattern_node) { - isolate_node(node_to_output.at(pattern_node).get_node_shared_ptr(), isol_tag, node_to_gptr); + // Isolate a required pattern node. + auto isolate = [&](const std::shared_ptr& pat) { + isolate_node(node_to_output.at(pat).get_node_shared_ptr(), isol_tag, node_to_gptr); + }; + // Isolate all nodes in a weight chain. + // convert_in and reshape are optional: when absent, opp::optional does NOT insert an + // entry into node_to_output, so we must guard with .count() before calling .at(). + auto isolate_weight_chain = [&](const WeightChainPattern& wc) { + if (node_to_output.count(wc.convert_in)) { + auto n_cin = node_to_output.at(wc.convert_in).get_node_shared_ptr(); + if (std::dynamic_pointer_cast(n_cin)) + isolate_node(n_cin, isol_tag, node_to_gptr); + } + isolate_node(node_to_output.at(wc.multiply).get_node_shared_ptr(), isol_tag, node_to_gptr); + if (node_to_output.count(wc.reshape)) { + auto n_rs = node_to_output.at(wc.reshape).get_node_shared_ptr(); + if (std::dynamic_pointer_cast(n_rs)) + isolate_node(n_rs, isol_tag, node_to_gptr); + } + isolate_node(node_to_output.at(wc.convert_out).get_node_shared_ptr(), isol_tag, node_to_gptr); }; isolate(tile); isolate(reshape1); - isolate(gate_weights_multiply); - isolate(gate_weights_convert); + isolate_weight_chain(gw); isolate(matmul_gate); isolate(swish); - isolate(up_weights_multiply); - isolate(up_weights_convert); + isolate_weight_chain(uw); isolate(matmul_up); isolate(multiply_swiglu); - isolate(down_weights_multiply); - isolate(down_weights_convert); + isolate_weight_chain(dw); isolate(matmul_down); isolate(reshape2); isolate_node(matched_output_multiply, isol_tag, node_to_gptr); @@ -431,16 +466,19 @@ Qwen3Expert::Qwen3Expert(const std::shared_ptr& snap Qwen3 Router Pattern: Router weights (quantized, dequantized via weight chain): - Convert(weight) -> Multiply(weight, scale) -> Convert -> MatMul(input, weight) + Convert(weight) -> Multiply(weight, scale) -> [Reshape?] -> Convert -> MatMul(input, weight) Score computation: - MatMul -> Softmax -> TopK(values, indices) + MatMul -> Softmax -> TopK(values[out0], indices[out1]) Score normalization: - TopK(values) -> ReduceSum -> Divide(values, sum) [renormalize over K selected] + TopK[out0] -> ReduceSum + TopK[out0] / ReduceSum = Divide -> [Slice?] (optional, present in group-quantized models) Scatter to full expert dimension: - TopK(indices) + Divide(scores) -> ScatterElementsUpdate(zero_broadcast, indices, scores) + TopK[out1] -> Convert(i64->i32) ──┐ + Divide -> [Slice?] ────────────────┤ + zero_broadcast ────────────────────→ ScatterElementsUpdate Shape to [num_experts, token_count, 1, 1] for expert broadcast: ScatterElementsUpdate -> Transpose -> Reshape -> Unsqueeze <-- pattern root @@ -451,11 +489,10 @@ Qwen3Router::Qwen3Router([[maybe_unused]] const std::shared_ptr Multiply(weight, scale) -> Convert -> MatMul - auto weights_convert_in = opp::wrap_type({opp::any_input()}); - auto weights_multiply = opp::wrap_type({weights_convert_in, opp::any_input()}); - auto weights_convert_out = opp::wrap_type({weights_multiply}); - auto matmul = opp::wrap_type({opp::any_input(), weights_convert_out}); + // Router weight dequantization chain (same layout as Qwen3Expert projections). + // Supports both regular and group-quantized weight tensors via opp::optional. + auto wc = make_weight_chain(); + auto matmul = opp::wrap_type({opp::any_input(), wc.convert_out}); // Score: Softmax -> TopK auto softmax = opp::wrap_type({matmul}); @@ -465,9 +502,20 @@ Qwen3Router::Qwen3Router([[maybe_unused]] const std::shared_ptr({topk, opp::any_input()}); auto divide = opp::wrap_type({topk, reduce_sum}); - // Scatter to full expert shape (pattern root = Unsqueeze) - auto scatter = - opp::wrap_type({opp::any_input(), topk, divide, opp::any_input()}); + // Optional Slice of normalized scores before Scatter. + // Present in group-quantized models (Divide -> Slice -> Scatter port 2). + // Absent in regular models (Divide -> Scatter port 2 directly). + auto slice = opp::optional( + {divide, opp::any_input(), opp::any_input(), opp::any_input(), opp::any_input()}); + + // Scatter to full expert shape. + // port 0: zero broadcast (any_input) + // port 1: Convert(TopK indices i64->i32) — TopK output(1) cannot be expressed in + // wrap_type (always binds output(0)), so use any_input() here + // port 2: slice (normalized scores, or Divide directly if Slice absent) + // port 3: axis constant (any_input) + auto scatter = opp::wrap_type( + {opp::any_input(), opp::any_input(), slice, opp::any_input()}); auto transpose = opp::wrap_type({scatter, opp::any_input()}); auto reshape = opp::wrap_type({transpose, opp::any_input()}); auto unsqueeze = opp::wrap_type({reshape, opp::any_input()});