Skip to content

Xiong/moe perf #36498

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
intelgaoxiong wants to merge 7 commits into
base: master
Choose a base branch
from
Draft

Xiong/moe perf #36498

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
fa3bb85
[NPUW]Support GQ MoE
intelgaoxiong Jun 17, 2026
482027f
Remove dead code.
intelgaoxiong Jun 23, 2026
164fdf1
Critial ac fix by ensuring batched experts are unrolled.
intelgaoxiong Jun 23, 2026
8d4008a
Critial ac issue fix with correct hidden dim size.
intelgaoxiong Jun 23, 2026
766d17c
Fixed graph isolation.
intelgaoxiong Jun 24, 2026
568235a
Cut attention along with MoE.
intelgaoxiong Jun 16, 2026
2b76239
Prevent GQ MoE + Device_routed.
intelgaoxiong Jul 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 33 additions & 2 deletions src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<std::string>().empty()) {
isol_it->second = isol_it->second.as<std::string>() + ",MOE";
LOG_INFO("MoE config: appended MOE to NPUW_ONLINE_ISOLATE -> " << isol_it->second.as<std::string>());
} 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. "
Expand Down Expand Up @@ -1125,6 +1136,11 @@ ov::npuw::LLMCompiledModel::LLMCompiledModel(const std::shared_ptr<ov::Model>& m
}

if (generate_moe_hint == ::intel_npu::npuw::llm::MoEHint::DEVICE_ROUTED) {
NPUW_ASSERT(is_cw_compressed(generate_model_variants.front()) &&
"MoE DEVICE_ROUTED mode is not supported for group quantized models. "
"NPU compiler does not yet support device-routed expert selection with group "
"quantization. Please use channel-wise (CW) quantized model or switch to "
"HOST_ROUTED MoE hint.");
// Apply model transformations only to GENERATE stage (PREFILL doesn't support DEVICE_ROUTED
// transformations)
for (auto&& model_variant : generate_model_variants) {
Expand Down Expand Up @@ -1194,6 +1210,21 @@ ov::npuw::LLMCompiledModel::LLMCompiledModel(const std::shared_ptr<ov::Model>& 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<std::string>() : "(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);
Expand Down
109 changes: 43 additions & 66 deletions src/plugins/intel_npu/src/plugin/npuw/moe/moe_executor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

#include "moe_executor.hpp"

#include <cstring>
#include <optional>

#include "../compiled_model.hpp" // For CompiledModel::CompiledModelDesc
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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) {
Expand All @@ -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<ov::ITensor> 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
Expand Down
17 changes: 12 additions & 5 deletions src/plugins/intel_npu/src/plugin/npuw/moe/moe_infer_utils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand All @@ -95,7 +100,7 @@ std::vector<size_t> parse_selected_experts_from_router(const ov::SoPtr<ov::ITens
expert_to_tokens.clear();

auto shape = router_output->get_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!");
}

Expand Down Expand Up @@ -168,8 +173,10 @@ void gather_router_scores(const ov::SoPtr<ov::ITensor>& 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
Expand Down
22 changes: 18 additions & 4 deletions src/plugins/intel_npu/src/plugin/npuw/moe_transformations/device_routed_moe_transform.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

Expand All @@ -298,14 +297,21 @@ LayerNodes collect_from_expert_output(const RouterInfo& router) {
std::dynamic_pointer_cast<ov::op::v0::Constant>(reshape->input_value(1).get_node_shared_ptr());
auto shape_data = shape_const->cast_vector<int64_t>();
if (!shape_data.empty() && shape_data[0] == static_cast<int64_t>(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<ov::op::v1::Multiply>(
reshape->input_value(0).get_node_shared_ptr())) {
continue;
}
nodes.constant_reshapes.push_back(reshape);
}
}

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<bool(const ov::Output<ov::Node>&, size_t, bool&)> weight_has_expert_dim;
weight_has_expert_dim = [&](const ov::Output<ov::Node>& input, size_t expert_dim, bool& chain_recognized) -> bool {
Expand All @@ -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<ov::op::v1::Reshape>(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;
};
Expand Down Expand Up @@ -645,8 +656,11 @@ bool DeviceRoutedMoETransform::run_on_model(const std::shared_ptr<ov::Model>& 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());
Expand Down
Loading
Loading