You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Update MoE offload implementation plan and implements a global state for MoE/QMoE - #32738
update the MoE CPU-offload implementation plan and document counter configuration and the initial-state format
add opt-in session-global expert state for CPU and built-in CUDA MoE/QMoE, shared with subgraphs and isolated between sessions
build the immutable kernel/expert index at initialization; keep global counters, reads, and updates in MoeExpertState
own one concrete KernelUsage collector per registered kernel in the session state; CPU and CUDA obtain it through OpKernelContext::GetKernelUsage and collect expert IDs directly
forward only the collector getter through the internal shared-provider bridge, without extending the public C API; keep copies, pinned buffers, and events in CudaRoutingSnapshot
commit collected usage in the executor after successful kernel execution; failed kernels do not update counters, and invocations that do not access collection do not replay stale selections
update counters in place once per node invocation; require finite non-negative alpha and beta with alpha + beta <= 1, defaulting to 0.9 and 0.1
use no RecordUsage mutex and reject overlapping runs on the same session
overlap CUDA counting with expert computation: enqueue a reusable pinned routing snapshot and completion event after top-k, then collect usage on the CPU after launching expert kernels, waiting only for the snapshot rather than the whole stream
preserve fused QMoE routing and union selected experts across row tiles; cover buffer reuse and aborted-invocation recovery
exclude counting state and collectors from minimal builds; explicitly reject unsupported configurations, including the CUDA plugin EP
This PR implements the session-global state and counting stage. It does not enable expert placement, CPU offload, swaps, or budget redistribution.
Validation
CPU and built-in CUDA targets build successfully.
35 focused framework/collector/state/counting tests passed on H200, including context collector ownership, failed-kernel commit rejection and all 4 CUDA integration cases with the default routing configuration.
All 4 CUDA counting integration cases also passed with fused routing disabled.
Generic collector tests cover deduplication, independent instances, atomic invalid-batch rejection, reset, and retained storage across invocations.
3 native minimal-build tests passed.
C++ formatting and git diff --check passed.
The earlier minimal Android size fix measured 1,594,893 bytes against the unchanged 1,595,392-byte threshold.
Rename-only follow-up: verified identifier/path-only changes and syntax/type-checked CUDA MoE, QMoE, and the renamed routing snapshot tests with the configured compiler flags.
Validation in progress: CPU and CUDA builds and focused tests are not yet confirmed.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Track the total expert population across all MoE nodes and apply configurable exponential alpha/beta updates to every expert in the invoked node.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Replace mutex-protected updates with ordinary counters and preallocated scratch. Freeze the kernel/expert registry before execution and reject overlapping runs at session entry.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Expose one internal MoeExpertUsage accessor without extending the public C API. Cache opt-in flags, skip diagnostics when disabled, and reuse CUDA collectors and buffers per kernel. Document the initial-state format and cover disabled usage access.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Capture routing IDs into reusable pinned memory before expert computation and wait only for the copy-completion event after launching the expert kernels. Preserve fused routing and union tiled QMoE selections before recording usage once per invocation. Add real-CUDA overlap and recovery coverage.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Append GetMoeExpertUsage after all existing public and protected virtual declarations without changing its public visibility.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Exclude counter state, context access, and run bookkeeping from minimal builds. Keep unsupported MoE option validation lightweight and explicit. Skip built-in CUDA counting cases for the plugin EP and cover plugin rejection and minimal-build configuration errors.
Android arm64 minimal sections: 1594893 bytes, below the unchanged 1595392-byte CI limit. Validated 24 full-runtime and 3 minimal-runtime targeted tests.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Update counters in place and retain a byte selection mask to deduplicate expert IDs. Require finite, non-negative alpha and beta with alpha + beta <= 1, using defaults 0.9 and 0.1.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Remove MoeExpertUsage and KernelUsage. Forward kernel-context recording through the shared-provider bridge to the session-global state, keeping CUDA collection independent of counter updates.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Move provider-independent collection and deduplication into KernelUsage and use it in CPU and CUDA MoE/QMoE. Keep only asynchronous transfer and event handling in the CUDA adapter; MoeExpertState retains global counters and updates. Preserve reusable selection storage across invocations.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Replace RecordMoeExpertUsage with GetKernelUsage across the runtime and shared-provider bridge. CPU and CUDA populate the same session-owned collector; successful executor invocations commit it to MoeExpertState. Remove duplicate state-side selection storage and cover failed collection commits and collector ownership.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Rename the adapter, header, test suite, and member variables to distinguish asynchronous routing snapshots from generic KernelUsage collection. No behavior changes.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Expose a flat InlinedVector<ExpertStat> (kernel, expert_id, popularity)
built directly from the kernel dictionary, for fast post-inference
access without walking the graph-scope snapshot map.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
I completed a full static review of this PR at 84c7712dabde1aa43e01c78fba40dc7816da1460. I found one major lifecycle issue and three smaller producer/consumer and maintainability issues.
Major: PartialRun() bypasses the MoE state's non-overlap guard
Normal RunImpl() brackets execution with KernelPilotMoeExpertState::BeginRun() and EndRun(), but the training-only PartialRun() path does not. PartialRun() still uses the sequential executor, so a successful MoE/QMoE kernel can still reach OpKernelContextInternal::RecordKernelUsage().
This permits overlapping PartialRun() calls, or a Run() overlapping a PartialRun(), to mutate the same session-global selection and counter state without the guard on which the lock-free design relies. Possible results include mixed selections, lost/corrupted counter updates, and logging with the wrong request/logger state.
Please either add the same scope-guarded BeginRun()/EndRun() lifecycle to PartialRun(), or explicitly reject PartialRun() while MoE counting/statistics is enabled. Preferably, centralize this lifecycle in a shared session execution scope so future entry points cannot omit it.
Relevant code: onnxruntime/core/session/inference_session.cc (PartialRun() versus the guarded normal run path) and onnxruntime/core/framework/sequential_executor.cc (post-compute usage commit).
Minor: the C++ producer can emit a record rejected by the supplied Python consumer
KernelPilotMoeExpertState::RecordUsage() accepts an empty selected-expert span and can emit:
{"selected_experts": []}
However, tools/python/qmoe_expert_distribution.py rejects an empty selected_experts list and aborts parsing the trace. Please align the contract: either omit/define empty-selection records in C++, or make the Python validator accept them with the intended semantics.
Minor: counter-update logging has no volume limit
When statistics logging is enabled, every successful MoE invocation synchronously formats and emits the selected experts plus the node's complete expert-counter vector. Large expert counts and long decode runs can create substantial CPU and log-storage overhead. Please add a per-run record or byte budget and emit one truncation summary when the budget is exhausted.
Minor: GetKernelPilot() establishes a hidden commit obligation
GetKernelPilot() appears to be an accessor, but caching its non-null result causes the executor to call RecordUsage() after a successful kernel. A future caller that merely probes availability without completing the full BeginInvocation -> Collect -> RecordUsage protocol can turn an otherwise successful kernel into a post-compute failure.
Please document this mandatory protocol adjacent to the API, or redesign the API so probing does not arm the commit path and an uninitialized invocation cannot fail during commit.
The current revisions do not run together because both MoE options reject graph-capture-enabled EPs. If that exclusion is relaxed, #32718's partition runner currently invokes kernels directly without committing MoE usage, while adaptive expert placement can change routing, GEMM work, and scratch requirements under a capture contract that assumes fixed execution and allocation sequences. A future hybrid MoE design needs an explicit internal capture boundary or a fixed worst-case schedule rather than capturing the adaptive orchestrator as one CUDA node.
Not tested in this review: concurrent PartialRun() reproduction, real-model counter semantics/performance, CUDA execution, or end-to-end analysis-tool processing.
PartialRun() is a training-only legacy path used by TrainingAgent to split forward and backward execution; that training path is deprecated and is not a supported entry point for KernelPilot MoE expert statistics.
The python consumer will evolve. The first version was used to proove it could be useful, it will be modified to finalize the study with the final modifications.
I fixed RecordUsage.
#32718 seems to be a prototype. I prefer to deal with it when the design is final.
Limit counter records to 1024 per Run without limiting counter updates. Emit one truncation marker and reject truncated analysis traces. Document collector preconditions and remove empty guards.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This CUDA collector is still compiled into an ORT_MINIMAL_BUILD whenever the full CUDA contrib provider is enabled: onnxruntime_providers_cuda.cmake:58-65,108-133 recursively adds every CUDA contrib .cc/.h and does not filter on onnxruntime_MINIMAL_BUILD. Only the framework collector files are excluded by this PR's CMake change. That contradicts the stated minimal-build exclusion and can retain the pinned-buffer/event implementation in minimal CUDA binaries; exclude these CUDA collector files or compile them out under ORT_MINIMAL_BUILD.
Preserve three-field identity formatting in error paths
This helper is also called by calculate_qmoe_expert_bytes() for the legacy three-field identities that function still supports (and that the tests still pass). On its missing-node or missing-initializer error paths, unpacking those identities as four fields raises an unrelated not enough values to unpack exception instead of the intended diagnostic. Preserve the three-field formatting path here, or remove that compatibility consistently.
Remove both CUDA expert-selection collector files from the minimal-build source list while leaving full CUDA builds unchanged.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
A failed Compute() leaves invocation_pending_ set because the executor skips RecordUsage(). If the next invocation only queries GetKernelPilot() and succeeds without calling BeginInvocation()/Collect(), this method sees the previous invocation as pending and commits its stale selection, contrary to the documented query-only no-op behavior. Explicitly discard pending usage when kernel execution fails (or bind pending state to the current context), and cover failed-then-query-only recovery.
Thanks for addressing the other review points. One producer/consumer mismatch from my previous comment remains at the current head (8168488): KernelPilotMoeExpertSelection::BeginInvocation() starts with an empty selection, and RecordUsage() can emit "selected_experts":[], but qmoe_expert_distribution.py rejects an empty selected_experts list. A successful invocation with no selected experts would therefore produce a log that the checked-in analyzer cannot read. Please align the two sides (either define and accept empty selections in the analyzer, or prevent/omit such records in the producer) and add a regression test for that case. I would hold off on approval until this contract is resolved.
Thanks for addressing the other review points. One producer/consumer mismatch from my previous comment remains at the current head (8168488): KernelPilotMoeExpertSelection::BeginInvocation() starts with an empty selection, and RecordUsage() can emit "selected_experts":[], but qmoe_expert_distribution.py rejects an empty selected_experts list. A successful invocation with no selected experts would therefore produce a log that the checked-in analyzer cannot read. Please align the two sides (either define and accept empty selections in the analyzer, or prevent/omit such records in the producer) and add a regression test for that case. I would hold off on approval until this contract is resolved.
I know, this is intended. This script will be extended in the next PR to do more statistics. The script was used to make a statistical study. Plus, no selected expert is unexpected unless the MoE fails in onnxruntime. So, I prefer to keep this case as failure.
The reason will be displayed to describe this comment to others. Learn more.
Thanks for clarifying the intended contract. I agree that the statistics analyzer may reject an empty selected_experts list as anomalous data; my earlier producer/consumer concern did not establish a successful, valid MoE invocation that produces an empty selection, so I am not blocking approval on it. Non-blocking follow-up: please document or test the invariant that a successful tracked MoE invocation selects at least one expert (including the zero-token boundary), or fail at the producer if an empty selection indicates a routing error. This keeps the intentional analyzer rejection and makes the contract explicit.
…untime/contrib_ops/cuda/moe/moe_quantization.cc. (#32932)
### Description
<!-- Describe your changes. -->
Replace old use of `routing_record` with `routing_snapshot_` in
onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc.
### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
Fix build error about undeclared identifier `routing_record`. Likely an
unlucky merge artifact from two recent PRs #32761 and #32738.
---------
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
### Description
<!-- Describe your changes. -->
Fix the CUDA QMoE packed-INT GEMV path's remaining reference to the
removed
`routing_record` API and add a focused expert-counting regression.
- Capture selected expert IDs immediately after routing and consume the
snapshot
before the packed-INT path returns, using the same lifecycle as the
neighboring
paths. Preserve the existing non-plugin/non-minimal guards.
- Add `MoeExpertCountingTest.CudaQMoEPackedIntGemv`: FP16, INT2, fused
SwiGLU,
three rows, four experts, top-k 1, 512-wide hidden/intermediate
dimensions,
and block size 64.
- Verify expert-dependent numerical outputs and per-invocation selection
counters:
duplicate routes `[0, 0, 2]`, a repeated invocation, then changed routes
`[1, 1, 1]`. Expected counters are `[0.1, 0, 0.1, 0]`,
`[0.19, 0, 0.19, 0]`, and `[0.171, 0.1, 0.171, 0]`.
- Disable CPU fallback and cap dense dequantization scratch at one byte,
so a
successful run cannot silently use the dense fallback.
Only `moe_quantization.cc` and `moe_expert_counting_test.cc` change.
Routing
selection and GEMV arithmetic are unchanged; no WebGPU/PagedAttention
changes
are included.
**Validation**
Built and executed locally on NVIDIA A100-SXM4-80GB (SM80), CUDA
12.8.93,
GCC 11.4, cuDNN 9.10.2.21, Release, warnings-as-errors, QUICK_BUILD=OFF.
The existing CUDA build tree was reconfigured from SM86 to SM80,
recompiling
the CUDA targets. This is not a claim of CUDA 13/GCC 14 CI validation.
After configuring with the same options plus `--update`, the following
build
command passed:
```bash
source /workspace/venv/bin/activate
python tools/ci_build/build.py \
--build_dir /workspace/onnxruntime/build/cuda-ci-20260929 \
--config Release --build --parallel 16 --cmake_generator Ninja \
--build_shared_lib --build_wheel --enable_onnx_tests --use_cuda \
--use_binskim_compliant_compile_flags --nvcc_threads 4 --flash_nvcc_threads 4 \
--cuda_version=12.8 --cuda_home=/usr/local/cuda-12.8 \
--cudnn_home=/workspace/venv/lib/python3.11/site-packages/nvidia/cudnn \
--enable_cuda_profiling --skip_tests --allow_running_as_root \
--targets onnxruntime_test_all onnxruntime_provider_test \
--cmake_extra_defines CMAKE_CUDA_ARCHITECTURES=80 \
onnxruntime_BUILD_UNIT_TESTS=ON \
onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON \
onnxruntime_BUILD_CUDA_QUANT_PREPROCESS=ON
```
Tests were run from the build output directory:
```bash
cd /workspace/onnxruntime/build/cuda-ci-20260929/Release
export LD_LIBRARY_PATH=/workspace/venv/lib/python3.11/site-packages/nvidia/cudnn/lib:/usr/local/cuda-12.8/lib64:${LD_LIBRARY_PATH:-}
./onnxruntime_test_all --gtest_filter='MoeExpertCountingTest.CudaQMoE*'
# PASS: 4 tests, no skips.
./onnxruntime_provider_test --gtest_filter='MoETest.QMoETest_Int2CudaPackedDecodeMultiRowExpertParity'
# PASS: 1 test, no skips.
ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO=1 ./onnxruntime_test_all \
--gtest_filter='MoeExpertCountingTest.CudaQMoEPackedIntGemv'
# PASS: 1 test; all three invocations log Route=packed_int_gemv Rows=3.
```
**Negative control:** temporarily omitted both packed-path `Capture` and
`Consume`, rebuilt with
`cmake --build build/cuda-ci-20260929/Release --target
onnxruntime_providers_cuda --parallel 16`,
and ran
`./onnxruntime_test_all
--gtest_filter='MoeExpertCountingTest.CudaQMoEPackedIntGemv'`.
It failed as intended: outputs still matched, but seven counter
assertions
failed because all counters remained zero (first failure: expert 0
expected
0.1, actual 0). Restored the repair, rebuilt the provider, and reran
both
focused suites: all five tests passed. No negative-control edits are
committed.
### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
The expert-counting migration in #32738 left a stale
`routing_record->CaptureTile`
call in the packed-INT GEMV early-return path. This prevents ordinary
built-in
CUDA compilation on current main; simply deleting that call would
compile but
silently omit expert counts for this path.
This separates the repair originally included in #32808, following
[the scope/coverage review
comment](#32808 (comment)).
The production patch is carried unchanged from commit `1f57e92abe`; the
new
regression verifies the runtime counting behavior as well as
compilation.
After this lands, merging main into #32808 will remove the unrelated MoE
file
from that PR's diff.
---------
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
### Description
<!-- Describe your changes. -->
- Move the MoE CPU offload roadmap and its images to
`docs/moe-cpu-offload/`.
- Remove the obsolete `docs/next_steps/` hierarchy.
- Restore the two exploratory routing graphs, their interpretation, and
links to the merged foundation PRs.
- Split the remaining implementation into fixed kernel-order hybrid
inference followed by adaptive swaps.
- Define nonblocking end-of-inference swaps, CPU-first transfer
ordering, and a two-exchange pinned-buffer/staging-slot limit.
- Limit the deferred measurements to token throughput and peak CPU/CUDA
memory for now.
### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
The session-global expert counters landed in #32738. The roadmap now
reflects the remaining work and separates hybrid inference correctness
from dynamic placement changes.
Documentation-only change; no tests were run.
---------
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…untime/contrib_ops/cuda/moe/moe_quantization.cc. (#32932)
### Description
<!-- Describe your changes. -->
Replace old use of `routing_record` with `routing_snapshot_` in
onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc.
### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
Fix build error about undeclared identifier `routing_record`. Likely an
unlucky merge artifact from two recent PRs #32761 and #32738.
---------
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Xavier Dupré (xadupre)
changed the title
Ti-Tai Wang (titaiwangms)
left a comment
Description
This PR implements the session-global state and counting stage. It does not enable expert placement, CPU offload, swaps, or budget redistribution.
Validation