Execution plan to offload export on CPU for a MoE

[Xavier Dupré (xadupre)]

Big models with expert may be too big for CUDA, the PR proposes an implementation plan to move experts on CPU. The difficulty is to choose which one we should offload.

[Copilot]

Pull request overview

Documents a staged plan for evaluating and implementing CPU expert offloading for CUDA-hosted MoE models.

Changes:

• Defines profiling, simulation, cache, fallback, and evaluation phases.
• Adds discussion and status indexes.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 6 comments.

File	Description
docs/next_steps/discuss/2026/08-moe-cpu-offload.md	Adds the implementation proposal.
docs/next_steps/discuss/index.md	Lists the proposal.
docs/next_steps/started/index.md	Adds the started-status index.
docs/next_steps/completed/index.md	Adds the completed-status index.

Suppressed comments (4)

docs/next_steps/discuss/2026/08-moe-cpu-offload.md:166

• A sequence-length decrease does not identify a new request in onnxruntime-genai. A normal decoder request changes from a multi-token prefill to sequence_length == 1 decoding, so this rule falsely marks the first generated token as a new sequence; subsequent decode shapes can remain constant. Use the explicit request ID for boundaries and record position/cache-length metadata separately if generation progress is needed.

Input dimensions are required even when `request_id` is present. During autoregressive generation the sequence-length dimension normally increases; a decrease relative to the previous iteration marks the transition to the next sequence. The analysis script validates this boundary against `request_id` and reports inconsistent traces.

docs/next_steps/discuss/2026/08-moe-cpu-offload.md:328

• This repeats the invalid boundary heuristic: onnxruntime-genai normally decreases input_ids from the prompt length to one token within the same request and then keeps that shape during decoding. Detect boundaries from request IDs instead of tensor-shape decreases.

- Detect sequence boundaries from decreases in the sequence-length input dimension and validate them against request identifiers.

docs/next_steps/discuss/2026/08-moe-cpu-offload.md:389

• Executing every non-resident Qwen 3.6 expert on CPU is not currently possible for the documented NVFP4 model, because CPU QMoE does not implement NVFP4 block/global scale semantics. Make CPU NVFP4 support an explicit prerequisite for this integration, or scope PR 6 to a Qwen export whose weight format the CPU kernel supports.

- Execute all non-resident experts on CPU for the current MoE invocation, then schedule cache updates after that invocation completes.

docs/next_steps/discuss/2026/08-moe-cpu-offload.md:237

• The hybrid estimate omits costs incurred on every CPU miss. Because the node is CUDA-owned, a miss must move routing/activation data to CPU, synchronize before CPU execution, copy the expert result back to CUDA, and merge it with CUDA-hit outputs. Applying standalone CPU execution cost alone systematically understates miss latency and can select the wrong policy.

A checked-in script combines the CPU and CUDA measurements to estimate hybrid execution. For each expert decision, it applies the measured CPU cost to a simulated miss, the measured CUDA cost to a hit, and the measured host-to-device cost to an admission. It then compares the resulting estimated iteration time with the CPU-only and CUDA-only baselines.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟡 Changes recommended

The analysis calculations, profiler serialization, and CUDA plugin behavior contain correctness gaps.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

tools/python/qmoe_expert_distribution.py:334

• The complement is off by one expert. At threshold 0 no experts are shortlisted, so the CPU-side complement is num_experts * bytes_per_rank; at threshold num_experts - 1, one expert still remains, not zero. The current value undercounts every complement row and distorts the normalized comparison.

• Files reviewed: 30/32 changed files
• Comments generated: 7
• Review effort level: Balanced

[Xavier Dupré (xadupre)]

Addressed the expert-ID validation finding in 1ff6e3d191. The analyzer now rejects every ID outside 0 <= expert_id < num_experts before updating any counters. Regression coverage includes both -1 and the upper bound (num_experts); all 5 analyzer unit tests pass.

[Xavier Dupré (xadupre)]

Restricted MoE routing logging to batch size 1 in 0fe9fe8b60:

• 3D MoE/QMoE inputs now fail when their explicit batch dimension is greater than 1, on CPU and CUDA.
• 2D inputs remain supported as unbatched token matrices, so single-prompt prefill with multiple tokens is not rejected.
• The restriction applies only when expert-statistics logging is enabled; normal inference is unchanged.

[Xavier Dupré (xadupre)]

Addressed the remaining trace-integrity and analysis-coverage findings in 8433b4e766:

• registered test_qmoe_expert_distribution.py explicitly in Python CI using its source path;
• added ordered prompt_start/prompt_end markers and a final moe_routing_complete footer with prompt, prompt-run, and routing-record counts;
• made the analyzer fail closed on missing, overlapping, out-of-order, truncated, or count-mismatched traces and cross-check the benchmark JSON;
• validate both expert-ID and router-weight cardinality/types;
• validate declared external-data lengths against tensor size and available file ranges;
• documented the sequential, single-prompt (batch=1) scope and the unsupported concurrent-Run scope.

The expanded analyzer suite now covers completeness/order, schema failures, invalid IDs, external-data validation, rank ties, final-row extraction, ranked CSV semantics, and threshold totals.

[Ti-Tai Wang (titaiwangms)]

Thanks for the update. The previous review's main issues are substantially addressed: the analyzer test is registered in CI, the trace has ordered prompt/completion markers and a fail-closed footer, invalid expert IDs and router-weight cardinality are rejected, and routing collection is explicitly limited to the documented sequential batch-1 protocol.

One blocking analyzer validation gap remains in calculate_qmoe_expert_bytes: when an initializer uses external data, a missing location entry or a referenced file that does not exist is currently accepted because the file-range check only runs when location is truthy and external_path.is_file() succeeds. The analyzer can therefore publish an expert-size estimate for a model that ORT cannot load. Please fail with a clear error when external data has no location or when the referenced file is unavailable, and add tests for both cases.

Two smaller follow-ups would also improve confidence:

• Add direct CUDA MoE and QMoE tests for the new 3D batch > 1 rejection path; the current CUDA routing tests use only 2D inputs.
• Clarify the documented aggregation semantics. inference_expert_ids() intentionally selects only the final row of each routing record, while the document currently says the analyzer maps every selected top-k expert. If final-row/decode-token analysis is intended, please state that explicitly; otherwise aggregate all rows consistently.

The external-data fail-open behavior is the only item above that I consider blocking.

[Copilot]

🟡 Changes recommended

The analyzer has correctness and scalability gaps, and some supported configurations can silently produce no routing data.

Get a fresh assessment by requesting another Copilot review.

Review details

• Files reviewed: 39/41 changed files
• Comments generated: 7
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

The prompt loader rejects valid single-record JSONL files, and the analyzer’s eager repeated parsing is not viable for the documented workload size.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (2)

tools/python/qmoe_expert_distribution.py:113

• This parser materializes every routing event in memory, and main() retains the first full list in _ at line 595 while later passes parse the entire log again. The documented 1,000-prompt workload extrapolates the reported 59,880 records/10 prompts to about 6 million JSON dictionaries, so keeping overlapping copies can consume multiple gigabytes and make the analyzer fail before producing results. Validate/aggregate in streaming passes (at most the ranking pass plus one output pass) instead of returning a full event list.
  tools/python/qmoe_prompt_runner.py:64
• A one-record JSON Lines file is rejected here. A single JSONL string or { "prompt": ... } object is also valid standalone JSON, so json.loads(text) succeeds with a scalar/dict and the JSONL fallback never runs; line 66 then rejects it despite the documented JSONL format. Treat a non-list successful parse as one JSONL record (and add a single-record regression case).

• Files reviewed: 39/41 changed files
• Comments generated: 1
• Review effort level: Balanced

[Xavier Dupré (xadupre)]

🟡 Changes recommended

The prompt loader rejects valid single-record JSONL files, and the analyzer’s eager repeated parsing is not viable for the documented workload size.

Get a fresh assessment by requesting another Copilot review.

Review details

The python script is meant to be used from time to time to analyze. It is meant to evolve over time. This will be addressed if the script is used on a regular basis.

[Copilot]

🔵 Needs a closer look

The analyzer’s repeated full-trace materialization is unsuitable for the planned scale, and identical runner output paths corrupt both artifacts.

Review details

Suppressed comments (2)

tools/python/qmoe_expert_distribution.py:648

• This initial validation retains every decoded event in _ for the rest of main, and each later helper reparses and materializes the complete trace again. The planned 1,000-prompt run scales the 59,880-record/10-prompt example to millions of JSON dictionaries, so peak memory can become several GB and parsing is repeated four times. Use a validation-only streaming pass, or parse once and pass/reuse the events without retaining duplicate lists.
  tools/python/qmoe_prompt_runner.py:166
• Reject identical output paths before opening the routing log. If --output and --routing-log resolve to the same file, the routing capture is first written and then overwritten by write_text; the appended completion footer subsequently makes the results file invalid JSON while reporting zero routing records.

• Files reviewed: 40/42 changed files
• Comments generated: 0 new
• Review effort level: Balanced

[Ti-Tai Wang (titaiwangms)]

Thanks for the update. The new CUDA MoE and QMoE batch > 1 rejection tests address the requested coverage gap, and the expanded (node_index, node_type, node_name) identity avoids collisions between unnamed routing nodes.

The blocking external-data validation issue remains unchanged in calculate_qmoe_expert_bytes(). The current logic still performs validation only when location is truthy and only checks the range when external_path.is_file() succeeds. As a result, an EXTERNAL initializer with no location, or one that references a nonexistent file, is silently accepted and the analyzer can publish sizing output for a model that ORT cannot load.

Please fail explicitly when an external initializer has no non-empty location, and when the resolved external-data file does not exist or is not a regular readable file. Please also add unit tests for both cases. This remains the only blocking issue from the latest lightweight re-review.

Non-blocking question: the analyzer now matches runtime records to raw ONNX nodes using (node_index, node_type, node_name). Please confirm that the runtime NodeIndex is guaranteed to match the raw model.graph.node enumeration used by onnx.load, including the intended optimization configuration; otherwise optimized graph transformations could make a valid QMoE record fail lookup.

[Xavier Dupré (xadupre)]

Addressed the remaining blocking and follow-up items in 5daecc8a4f after merging current main:

• external QMoE initializers now fail closed when location is absent/empty or the referenced file cannot be opened as a regular file;
• regression coverage includes both a missing location and a nonexistent external file;
• routing validation and aggregation now happen in one streaming pass instead of retaining and reparsing every decoded event;
• the prompt runner rejects identical resolved --output and --routing-log paths before opening either output;
• the document now states the final-row/decode-token rank semantics.

For the NodeIndex question: Graph creates nodes by iterating GraphProto.node in order (graph.cc:1466-1468), AllocateNode() assigns nodes_.size() (graph.cc:5957-5963), and removing a node clears its slot without renumbering survivors (graph.cc:5969-5981). The analyzer additionally matches operator type and name, so replacement of a QMoE node fails closed rather than being attributed to another raw node. This invariant and validation behavior are now documented.

The merge also includes the current WebGPU ep::adapter::Logger declaration fix that was missing from the prior head. All 12 targeted analyzer unit tests pass, and lintrunner reports no issues on the changed Python files.