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
Enable static CPU offload for FP16/BF16 CUDA MoE - #33039
Add the first static CPU-offload slice for built-in CUDA FP16 and BF16 MoE:
introduce session.moe_cpu_offload_experts as the session-wide number of experts offloaded to and executed on CPU;
let KernelPilot finalize deterministic expert placement before session initialization returns;
retain one host representation of every expert's constant weights and materialize only selected expert slices in compact CUDA storage;
execute CPU-resident FP16 routes with MLAS FP16 GEMMs and BF16 routes with MLAS FP32 GEMMs after one-time weight conversion, execute CUDA-resident routes with CUTLASS, and combine both contributions on CUDA;
preserve full routing usage collection and exponential counter updates for every successful invocation;
reject unsupported configurations such as dynamic offloaded weights, FC3, sparse mixer, CUDA graph capture, and the CUDA plugin EP.
The provider bridge exposes the MLAS FP16 and FP32 GEMM calls required by the shared CUDA provider. Documentation describes the current static-placement behavior and limitations.
All-CUDA nodes retain the original host weights without a GEMM-layout copy. Nodes with CPU experts retain the
transposed GEMM-layout weights and release the temporary original copy. Offload requires enabled prepacking, a
static input hidden dimension, and constant connected FC1/FC2 biases.
Mixed execution packs CPU-routed activations on the compute stream before launching an in-place CUDA runner,
protecting input/output aliasing while allowing the CPU worker and CUDA expert GEMMs to overlap. The synchronous
fallback retains its explicit input-copy dependency.
Initializer collection and use validation follow actual outer-scope captures and exclude shadowing subgraph-local
inputs and initializers, keeping unrelated same-named ancestor values out of CPU placement.
This PR also restores free-threaded dependency filtering in the AArch64 Python packaging image. The Docker build
failed while installing ONNX and its transitive ml_dtypes dependency under CPython 3.13t. The restored filtering
excludes ONNX-dependent tooling only from free-threaded environments; standard CPython requirements are unchanged.
This is a packaging CI unblock, not part of the MoE execution change. Shell syntax and requirements filtering were
checked locally; the private base-image Docker build has not been reproduced locally.
cmake --build build/linux_cuda/Release --target onnxruntime_test_all -j 8 — passed after adding the prepacking configuration regression.
cmake --build build/linux_cuda/Release --target onnxruntime_provider_test -j 8 — passed after adding the sparse-mixer rejection regression.
cmake --build build/linux_cuda/Release --target onnxruntime_provider_test -j 8 — passed after strengthening the non-square layout regression.
cd build/linux_cuda/Release && CUDA_VISIBLE_DEVICES=5 ./onnxruntime_provider_test --gtest_filter='MoETest.MoECuda*StaticCpuOffload*' — 10/10 passed. This includes BF16 disabled, mixed, and all-CPU placement plus the existing FP16 activation, layout, aliasing, and rejection coverage.
cd build/linux_cuda/Release && ./onnxruntime_test_all --gtest_filter='MoeExpertCountingTest.StaticCpuOffload*:InferenceSessionTests.RunBeforeInitializeReturnsError' — 10/10 passed, including branch-local expert weights colliding with unrelated ancestor initializers, Loop-local inputs/initializers shadowing root weights, actual outer-scope captures, planner input/output reuse, and all four combinations of disabled/enabled prepacking and zero/positive CPU offload. The conflicting prepacking combination rejects initialization with the expected diagnostic; the other three initialize successfully.
cd build/linux_cuda/Release && ./onnxruntime_test_all --gtest_filter='MoeExpertCountingTest.StaticCpuOffloadKeepsBranchLocalWeightsSeparateFromAncestor' — failed as expected with the original name-only collection restored temporarily, reproducing the incorrect initialization rejection. With the fix restored, the same regression passes in the 10/10 run above and verifies both branch results plus the untouched ancestor value and its GPU placement.
cd build/linux_cuda/Release && ./onnxruntime_provider_test --gtest_filter='MoETest.MoECudaFp16StaticCpuOffloadSupportsNonSquareLegacyWeightLayout' — passed for standard and legacy shapes with identical asymmetric physical weight bytes and nonuniform inputs, on CUDA-only, mixed, and all-CPU execution. Temporarily bypassing transposition for legacy shapes as suggested by review failed the mixed/all-CPU cases; that change was removed. The final 9/9 operator run passes with the existing conversion restored.
cd build/linux_cuda/Release && CUDA_VISIBLE_DEVICES=5 ./onnxruntime_test_all --gtest_filter='MoeExpertCountingTest.StaticCpuOffload*' — 9/9 passed with BF16 included in CUDA-offload eligibility.
CPU-offloaded CUDA MoE nodes update expert counters on GPU and defer successful usage collection to EndRun; worker failures propagate to the caller and failed invocations are drained without committing usage.
Routing packs only rows selecting CPU-resident experts. Reusable per-node workers transfer and compute those compact batches while CUDA expert kernels execute.
Result combination expresses an explicit CUDA-visible upload dependency. This avoids intermittent deadlocks with subgraphs and tactic profiling; empty single-row CPU routing can release its completion flag directly on GPU.
The synchronous fallback remains available for older CUDA versions without stream memory operations. Modern CUDA uses the v2 APIs rather than the deprecated v1 capability flag.
Latest validation (isolated build of this PR, not the adaptive branch):
cd build/linux_cuda_static_validation && CUDA_VISIBLE_DEVICES=5 GTEST_FILTER='KernelPilotMoeExpertSelectionCudaTest.*' timeout 180 ./onnxruntime_provider_test --gtest_filter='CUDA_EP_Unittest.All' — 8 internal CUDA tests passed, including GPU counter mathematics and worker-free empty decode completion.
The matched active-swap performance measurement belongs to the stacked adaptive PR #33155: 18.2079 ms/token with 1,000 CPU experts, compared with the historical unprofiled 26.0387 ms/token baseline. It is not a standalone latency claim for this static-placement PR.
Motivation and Context
Large FP16 and BF16 MoE models may not fit all expert weights in VRAM. This change establishes session-initialization-time CPU/CUDA expert placement and mixed inference while deliberately leaving runtime movement, swapping, and redistribution for later roadmap stages.
Resolve the MoE roadmap move by applying the static FP16 offload documentation to docs/moe-cpu-offload/README.md.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Treat session.moe_cpu_offload_experts=0 as the disabled default and derive the CUDA expert budget from the eligible FP16 MoE expert count.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Preserve the default CUDA input placement and add multi-node placement policy coverage for zero counters, loaded counters, and over-budget rejection.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Shall we add options like session.cpu_offload_moe_nodes and session.cpu_offload_qmoe_nodes, so that we can offload given number of MoE/QMoE nodes to CPU? That only need node placement changes, and kernel need not change.
Fixed in d4103d8. The CPU-only test helper is now CUDA-guarded, and KernelPilot-specific MoE hooks are excluded from CUDA plugin/minimal builds. I couldn’t run a local build because this environment has no configured build tree or nvcc.
Shall we add options like session.cpu_offload_moe_nodes and session.cpu_offload_qmoe_nodes, so that we can offload given number of MoE/QMoE nodes to CPU? That only need node placement changes, and kernel need not change.
I believe this mechanism is already implemented with static offloading per layer, this is not exactyl the same but it should produce similar performance.
Skip profiling and workspace allocation when no CUDA experts exist
onnxruntime/contrib_ops/cuda/moe/moe.cc:297
This forces a nonzero runner expert count even when cuda_experts_ is empty, so the tactic profiler and runner workspace are still created for the all-CPU configuration before CUTLASS execution is skipped below. Because profiling uses need_weights=true, GemmProfilerBackend allocates synthetic FC1/FC2 weights proportional to this count (moe_kernels.cu:3280-3287) and launches profiling kernels despite no CUDA expert work, adding substantial latency and potentially causing an avoidable GPU OOM. Skip tactic profiling and use zero runner workspace when there are no CUDA-resident experts.
Reject unsupported standalone FC3 bias inputs and skip CUTLASS tactic profiling and workspace allocation when every expert is offloaded to CPU.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Override the initializer allocation plan before materialization so CUDA FP16 MoE expert constants remain on CPU. PrePack then uploads only the experts selected for CUDA residency.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Validate every shared initializer occurrence, use pageable memory for persistent CPU expert weights, and pass the operator thread pool to MLAS half GEMMs.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The reason will be displayed to describe this comment to others. Learn more.
Warning
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
I see #33039/#33155 as complementary to the policy-based workspace estimation proposal. Currently, session.moe_cpu_offload_experts specifies a fixed offload count—even when VRAM is sufficient—and counters determine which experts stay on GPU.
Perhaps a follow-up PR could add an optional budget-aware auto policy. Given a VRAM budget and workload envelope, ORT would account for resident weights, KV cache, execution workspace, swap staging, and a safety reserve, then choose GPU expert-slot counts. All experts could remain on GPU when the complete memory bound fits; otherwise, ORT would offload enough to meet the budget.
The existing KernelPilot could continue selecting and swapping expert identities within that capacity. This would give FL/GenAI a budget-driven deployment policy without manually guessing an expert count, while preserving the current fixed-count option as an explicit override.
Merged main (b131c79) in da0162a. This incorporates #33168's fix for the Linux/QNN ONNX node-test failures and #33164's Android size-baseline correction. Android minimal compilation and 13 build-argument tests passed. Local size was 1,627,925 bytes, above the threshold, using NDK 27.3 rather than CI's 28.0; that check remains unverified in the CI environment. Stopped the remaining native test build as requested. Secret scans were clean; automated review was unavailable and CodeQL timed out with an upstream workflow-permissions alert. New CI is pending.
I see #33039/#33155 as complementary to the policy-based workspace estimation proposal. Currently, session.moe_cpu_offload_experts specifies a fixed offload count—even when VRAM is sufficient—and counters determine which experts stay on GPU.
Perhaps a follow-up PR could add an optional budget-aware auto policy. Given a VRAM budget and workload envelope, ORT would account for resident weights, KV cache, execution workspace, swap staging, and a safety reserve, then choose GPU expert-slot counts. All experts could remain on GPU when the complete memory bound fits; otherwise, ORT would offload enough to meet the budget.
The existing KernelPilot could continue selecting and swapping expert identities within that capacity. This would give FL/GenAI a budget-driven deployment policy without manually guessing an expert count, while preserving the current fixed-count option as an explicit override.
What's your thoughts?
Having a fixed number of expert offloaded on CPU is not really a good setting because the value depends on VRAM and the model size. I totally agree with it. I chose that parameter to force onnxruntime to offload and get figures knowing we would need another PR to change that setting and the decision to offload. The first goal of this current implementation is to draw a plot x=number of offloaded expert, y=(latency, memory consumption).
[P2] Preserve FP16/FP32 autotuning when offload is disabled — moe.cc:385.
The previous unconditional profiling branch is now restricted to BFloat16 . With the default non-deterministic setting, ordinary FP16 and FP32 sessions—including those without CPU offload—skip profiling and use the constructor’s first tactic instead of shape-specific tuned configurations. Restore the unconditional else ; the surrounding run_cuda_experts guard already excludes all-CPU execution.
[P2] Reject unsupported offload settings in minimal builds — inference_session.cc:2773–2808.
The minimal-build rejection list omits session.moe_cpu_offload_experts , while its parsing and implementation are compiled out. A positive setting therefore succeeds but silently leaves offloading disabled, contrary to the documented unsupported configuration. Reject positive values in minimal builds and add corresponding coverage, while allowing numeric zero.
[P2] Require SM80 for the new BF16 integration test — moe_test.cc:4334–4335.
This helper is instantiated for BFloat16 but accepts SM70/SM75 devices. Its disabled and mixed cases invoke CUDA MoE GEMMs, whose dispatch rejects architectures below SM80, so the test fails rather than skips on those devices. Use the dtype-dependent 800 / 700 gate already present in the new non-square-layout helper.
Restore FP16 and FP32 tactic profiling, reject positive CPU offload counts in minimal builds, and gate BF16 CUDA tests on SM80.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Reject nonzero CPU offload settings without pulling integer parsing into minimal builds, while accepting zero-only numeric representations.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The three findings in #33039 (comment) are addressed on the current branch: the normal FP16/FP32 profiling path was restored in ef78678; positive session.moe_cpu_offload_experts values are rejected in minimal builds with coverage in 2834ef3 and the validation was compacted in 07bb1cd; the BF16 integration helper now requires SM80 in 2834ef3.
Skip CPU offload path when no selected experts are on CPU
onnxruntime/contrib_ops/cuda/moe/moe.cc:787
run_cpu_experts reflects static placement, not whether this invocation actually selected an offloaded expert. Consequently, a mixed node whose current routes all hit CUDA still allocates the full host output and CPU scratch, runs the CPU helper, transfers an all-zero output tensor back to CUDA, and launches the add kernel. This is especially costly for hot-expert workloads because it adds an output-sized H2D transfer on every invocation. After the routing copy completes, detect whether any selected expert maps to CPU and skip the CPU compute/output-copy/add path when none do (the eager input copy can remain for alias safety).
Cover BF16 activation branches in offload tests
onnxruntime/test/contrib_ops/moe_test.cc:4429
BF16 offload has a separate CPU implementation, but this BF16 test only exercises ReLU. The FP16-only activation test below does not cover the BF16 GELU, SiLU, identity, or fused SwiGLU branches, so errors in BF16 activation/indexing logic can pass CI. Parameterize the all-activations test over MLFloat16 and BFloat16 (with the existing SM80 gate for BF16).
The reason will be displayed to describe this comment to others. Learn more.
We seldom use fp16/bf16 MoE for CUDA EP, since most models need quantize to int4/nvfp4 QMoE operator for performance and VRAM limit.
This PR adds a feature that might not be used.
We shall focus on make a real use case like qwen 3.8 flash next model with int4 or nvfp4 QMoE that run end-to-end with onnxruntime-genai integration.
There are a few things need consideration:
(1) CPU and CUDA kernel shall use same prepacked format, so that we need not transform the data layout after we copy expert weights from CPU to GPU.
(2) ORT and GenAI integration: we might let GenAI handle some expert weights copy or take part in planning.
The reason will be displayed to describe this comment to others. Learn more.
I started with this one to simplify the transfer between CUDA/CPU. I know it has to be extended to QMoE but the kernels has different prepacking so the transfer is not straightforward. My first goal is agree on a design valid for both first.
Remove the independent beta option and derive the selected-expert increment as 1 - alpha.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Expect the floating-point complement of alpha after removing the independent beta coefficient.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Defer successful usage snapshots to EndRun and propagate worker errors. Run compact CPU-routed batches on reusable workers, and express the result upload dependency before combining outputs.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Superseded by #33347, which combines static CPU offload and adaptive expert swaps in one PR targeting main. It includes the current design and the 33 available in-grid benchmark runs across 21 configurations. This PR is closed as superseded; its changes are included in the replacement branch.
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.
Description
Add the first static CPU-offload slice for built-in CUDA FP16 and BF16
MoE:session.moe_cpu_offload_expertsas the session-wide number of experts offloaded to and executed on CPU;KernelPilotfinalize deterministic expert placement before session initialization returns;The provider bridge exposes the MLAS FP16 and FP32 GEMM calls required by the shared CUDA provider. Documentation describes the current static-placement behavior and limitations.
All-CUDA nodes retain the original host weights without a GEMM-layout copy. Nodes with CPU experts retain the
transposed GEMM-layout weights and release the temporary original copy. Offload requires enabled prepacking, a
static input hidden dimension, and constant connected FC1/FC2 biases.
Mixed execution packs CPU-routed activations on the compute stream before launching an in-place CUDA runner,
protecting input/output aliasing while allowing the CPU worker and CUDA expert GEMMs to overlap. The synchronous
fallback retains its explicit input-copy dependency.
Initializer collection and use validation follow actual outer-scope captures and exclude shadowing subgraph-local
inputs and initializers, keeping unrelated same-named ancestor values out of CPU placement.
This PR also restores free-threaded dependency filtering in the AArch64 Python packaging image. The Docker build
failed while installing ONNX and its transitive
ml_dtypesdependency under CPython 3.13t. The restored filteringexcludes ONNX-dependent tooling only from free-threaded environments; standard CPython requirements are unchanged.
This is a packaging CI unblock, not part of the MoE execution change. Shell syntax and requirements filtering were
checked locally; the private base-image Docker build has not been reproduced locally.
Validation:
cmake --build build/linux_cuda/Release --target onnxruntime_test_all onnxruntime_provider_test -j 8— passed.cmake --build build/linux_cuda/Release --target onnxruntime_test_all -j 8— passed after adding the prepacking configuration regression.cmake --build build/linux_cuda/Release --target onnxruntime_provider_test -j 8— passed after adding the sparse-mixer rejection regression.cmake --build build/linux_cuda/Release --target onnxruntime_provider_test -j 8— passed after strengthening the non-square layout regression.cd build/linux_cuda/Release && CUDA_VISIBLE_DEVICES=5 ./onnxruntime_provider_test --gtest_filter='MoETest.MoECuda*StaticCpuOffload*'— 10/10 passed. This includes BF16 disabled, mixed, and all-CPU placement plus the existing FP16 activation, layout, aliasing, and rejection coverage.cd build/linux_cuda/Release && ./onnxruntime_test_all --gtest_filter='MoeExpertCountingTest.StaticCpuOffload*:InferenceSessionTests.RunBeforeInitializeReturnsError'— 10/10 passed, including branch-local expert weights colliding with unrelated ancestor initializers, Loop-local inputs/initializers shadowing root weights, actual outer-scope captures, planner input/output reuse, and all four combinations of disabled/enabled prepacking and zero/positive CPU offload. The conflicting prepacking combination rejects initialization with the expected diagnostic; the other three initialize successfully.cd build/linux_cuda/Release && ./onnxruntime_test_all --gtest_filter='MoeExpertCountingTest.StaticCpuOffloadKeepsBranchLocalWeightsSeparateFromAncestor'— failed as expected with the original name-only collection restored temporarily, reproducing the incorrect initialization rejection. With the fix restored, the same regression passes in the 10/10 run above and verifies both branch results plus the untouched ancestor value and its GPU placement.cd build/linux_cuda/Release && ./onnxruntime_provider_test --gtest_filter='MoETest.MoECudaFp16StaticCpuOffloadSupportsNonSquareLegacyWeightLayout'— passed for standard and legacy shapes with identical asymmetric physical weight bytes and nonuniform inputs, on CUDA-only, mixed, and all-CPU execution. Temporarily bypassing transposition for legacy shapes as suggested by review failed the mixed/all-CPU cases; that change was removed. The final 9/9 operator run passes with the existing conversion restored.cd build/linux_cuda/Release && CUDA_VISIBLE_DEVICES=5 ./onnxruntime_test_all --gtest_filter='MoeExpertCountingTest.StaticCpuOffload*'— 9/9 passed with BF16 included in CUDA-offload eligibility.git diff --check— passed before commit.bash -n tools/ci_build/github/linux/docker/inference/aarch64/python/cpu/scripts/install_deps.sh— passed.Latest execution update:
EndRun; worker failures propagate to the caller and failed invocations are drained without committing usage.Latest validation (isolated build of this PR, not the adaptive branch):
cmake --build build/linux_cuda_static_validation --target onnxruntime onnxruntime_test_all onnxruntime_provider_test onnxruntime_providers_cuda_ut -j 16— passed.cd build/linux_cuda_static_validation && CUDA_VISIBLE_DEVICES=5 timeout 180 ./onnxruntime_test_all --gtest_filter='KernelPilotMoeExpertStateTest.*:MoeExpertCountingTest.*'— 41 passed, none skipped.cd build/linux_cuda_static_validation && CUDA_VISIBLE_DEVICES=5 timeout 180 ./onnxruntime_provider_test --gtest_filter='MoETest.MoECuda*StaticCpuOffload*'— 11 passed, none skipped.cd build/linux_cuda_static_validation && CUDA_VISIBLE_DEVICES=5 GTEST_FILTER='KernelPilotMoeExpertSelectionCudaTest.*' timeout 180 ./onnxruntime_provider_test --gtest_filter='CUDA_EP_Unittest.All'— 8 internal CUDA tests passed, including GPU counter mathematics and worker-free empty decode completion.The matched active-swap performance measurement belongs to the stacked adaptive PR #33155: 18.2079 ms/token with 1,000 CPU experts, compared with the historical unprofiled 26.0387 ms/token baseline. It is not a standalone latency claim for this static-placement PR.
Motivation and Context
Large FP16 and BF16 MoE models may not fit all expert weights in VRAM. This change establishes session-initialization-time CPU/CUDA expert placement and mixed inference while deliberately leaving runtime movement, swapping, and redistribution for later roadmap stages.