Enable asynchronous adaptive FP16/BF16 MoE expert swaps

[Xavier Dupré (xadupre)]

Description

Adds adaptive FP16/BF16 CUDA MoE expert swaps on top of #33039. After each successful inference, the session-global pilot ranks the hottest CPU expert against the coldest CUDA expert in each eligible node and schedules an exchange only when cpu_max > (1 + epsilon) * cuda_min. Each kernel permits one pending exchange, and the scheduler permits at most four concurrent exchanges per CUDA device.

The swap protocol keeps the published placement valid throughout the transfer:

1. StartSwap() waits for the outgoing expert's last use on the compute stream. CUDA kernels transpose its FC1/FC2 weights into MLAS layout before a dedicated D2H stream copies them, along with biases, into reusable outgoing pinned storage.
2. Incoming expert weights are copied contiguously into separate pinned staging in MLAS layout; BF16 CPU weights are converted linearly from FP32. No host transpose callback runs. Published CPU weights remain unchanged until publication.
3. A dedicated H2D stream uploads the incoming MLAS-layout weights and CUDA kernels transpose them into reusable CUDA staging. This revision waits for outgoing D2H completion before reusing the shared layout buffer. The persistent CUDA slot and mapping remain unchanged during transfer.
4. At the start of a later invocation, the kernel queries the transfer event. If incomplete, it uses the old placement. If complete, it copies the staged expert into the persistent slot on that invocation's compute stream, updates the host and device expert maps, and publishes the same mapping to KernelPilot.
5. After the publication event completes, the kernel returns to the idle state. The pinned and CUDA staging buffers remain allocated for reuse by the next exchange and are released only when the kernel is destroyed.

This delayed-publication protocol ensures that an invocation sees either the complete old placement or the complete new placement, never partially transferred weights. Up to four kernels per CUDA device can pipeline independent D2H and H2D transfers; hardware with fewer copy engines serializes them without changing correctness. Per-node slot counts and the session-global CUDA expert budget remain unchanged; cross-node slot redistribution and dynamic adjustment of the concurrency limit remain follow-up work.

Validation:

• cmake --build build/linux_cuda/Release --target onnxruntime_test_all onnxruntime_provider_test -j 8 — passed
• CUDA_VISIBLE_DEVICES=5 build/linux_cuda/Release/onnxruntime_test_all --gtest_filter='MoeExpertCountingTest.AdaptiveCpuOffloadPublishesCompleted*SwapsAtNextInvocation:KernelPilotMoeExpertStateTest.*' — 17 passed, including real FP16 and BF16 delayed-publication swaps
• CUDA_VISIBLE_DEVICES=5 build/linux_cuda/Release/onnxruntime_provider_test --gtest_filter='MoETest.MoECuda*16StaticCpuOffloadRunsDisabledMixedAndAllCpuExpertsAndCountsUsage' — 2 passed, covering FP16 and BF16 disabled, mixed, and all-CPU placement
• cmake --build build/linux_moe_minimal/MinSizeRel --target onnxruntime_test_all -j 8 — passed
• build/linux_moe_minimal/MinSizeRel/onnxruntime_test_all --gtest_filter=OrtModelOnlyTests.RejectsMoeExpertCountingInMinimalBuild — 1 passed
• lintrunner -a — passed

Latest transfer and execution update:

• CUDA tiled kernels transpose FP16/BF16 swap weights between CUDA and MLAS layouts. No heavy host callback transposes complete outgoing/incoming experts.
• Separate pinned incoming/outgoing buffers preserve published CPU and CUDA weights until delayed publication. At this revision, the directions still share one CUDA layout buffer and are serialized; bidirectional overlap has not yet been claimed.
• GPU counters and compact CPU workers are integrated from the independently validated Enable static CPU offload for FP16/BF16 CUDA MoE #33039 update. Successful usage is consumed before swap scheduling; asynchronous worker errors are returned to the caller.
• CPU work overlaps CUDA expert computation, with a CUDA-visible upload dependency established only at output combination. Empty single-row CPU routing can complete directly on GPU.

Matched unprofiled performance, Qwen3.5-35B FP16, GPU 5, 1,000 CPU experts:

Revision	Decode mean (ms/token)
Historical baseline	26.038651
CUDA swap transposition and reduced host waits	21.139596
GPU counters and compact CPU worker	18.207928

The final row measures prompts 3–5 from the same five-prompt runner, 256 generated tokens per prompt, alpha 0.9, epsilon 0.1, swaps enabled, counter logging disabled and CUDA Graph disabled. Prefill and the first generation call without an actual decode inference are excluded: 255 actual decode runs per scored prompt, 765 total. Per-prompt means are 18.718456, 16.516421 and 19.388907 ms/token. The measured gain over the historical baseline is 7.830724 ms/token, not yet the requested 10 ms.

Latest validation:

• cmake --build build/linux_cuda/Release --target onnxruntime onnxruntime_test_all onnxruntime_provider_test onnxruntime_providers_cuda_ut -j 16 — passed.
• cd build/linux_cuda/Release && CUDA_VISIBLE_DEVICES=5 timeout 180 ./onnxruntime_test_all --gtest_filter='KernelPilotMoeExpertStateTest.*:MoeExpertCountingTest.*' — 49 passed, none skipped.
• cd build/linux_cuda/Release && CUDA_VISIBLE_DEVICES=5 timeout 180 ./onnxruntime_provider_test --gtest_filter='MoETest.MoECuda*AdaptiveSwap*:MoETest.MoECuda*StaticCpuOffload*' — 13 passed, including repeated FP16/BF16 swaps with asymmetric weights and biases.
• cd build/linux_cuda/Release && CUDA_VISIBLE_DEVICES=5 GTEST_FILTER='KernelPilotMoeExpertSelectionCudaTest.*' timeout 180 ./onnxruntime_provider_test --gtest_filter='CUDA_EP_Unittest.All' — 8 internal CUDA tests passed, including actual v2 stream-memory completion.
• cd build/linux_cuda/Release && CUDA_VISIBLE_DEVICES=5 timeout 120 ./onnxruntime_test_all --gtest_filter='MoeExpertCountingTest.StaticCpuOffloadKeepsBranchLocalWeightsSeparateFromAncestor' --gtest_repeat=30 — all 30 repetitions passed after fixing the upload dependency.

Motivation and Context

Static FP16/BF16 placement from #33039 cannot adapt when routing popularity changes between requests. Asynchronous swaps improve placement over time without making the next MoE invocation wait for expert-weight transfers or invalidating the mapping used by in-flight CUDA work.

[Copilot]

Copilot review overview

🔵 Needs a closer look

Multi-stream CUDA synchronization, callbacks, and mutable expert placement warrant final human review despite the focused tests.

Review effort: Balanced
Findings: None

What changed in this PR

Adds adaptive asynchronous expert swapping to FP16 CUDA MoE CPU offload. No actionable issues were identified.

Changes:

• Adds session-wide swap policy with epsilon thresholds and two concurrent exchanges per device.
• Implements provider-side asynchronous D2H/H2D staging and deferred publication.
• Adds configuration, documentation, policy tests, and CUDA integration coverage.

File	Description
onnxruntime/​test/​framework/​ort_model_only_test.cc	Covers minimal-build rejection.
onnxruntime/​test/​framework/​moe_expert_counting_test.cc	Tests end-to-end swap publication and results.
onnxruntime/​test/​framework/​kernel_pilot_moe_expert_state_test.cc	Tests policy, limits, and validation.
onnxruntime/​core/​session/​inference_session.cc	Validates configuration and schedules only after successful runs.
onnxruntime/​core/​framework/​session_state.cc	Initializes the swap threshold.
onnxruntime/​core/​framework/​kernel_pilot.h	Exposes cache attachment and placement publication.
onnxruntime/​core/​framework/​kernel_pilot.cc	Retains registered expert counts.
onnxruntime/​core/​framework/​kernel_pilot_moe_expert_state.h	Declares adaptive policy state.
onnxruntime/​core/​framework/​kernel_pilot_moe_expert_state.cc	Ranks experts and schedules bounded swaps.
onnxruntime/​core/​framework/​kernel_pilot_moe_expert_cache.h	Defines the provider cache interface.
onnxruntime/​contrib_ops/​cuda/​moe/​moe.h	Declares CUDA swap resources and phases.
onnxruntime/​contrib_ops/​cuda/​moe/​moe.cc	Implements asynchronous transfer and publication.
include/​onnxruntime/​core/​session/​onnxruntime_session_options_config_keys.h	Adds the epsilon configuration key.
docs/​moe-cpu-offload/​README.md	Documents adaptive swapping and remaining work.

---

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

[Xavier Dupré (xadupre)]

Profiling update for dynamic offload with 1,000 CPU-resident experts on H200 (alpha=0.9, epsilon=0.1, no counter logs, no CUDA Graph). I added NVTX ranges and compared the same 765 decoded tokens from prompts 3-5 against a matching all-GPU profile.

Measurement per decoded token	All GPU	1,000 CPU experts	Difference
Wall-clock latency	14.34 ms	37.83 ms	+23.49 ms
All GPU kernels	5.94 ms	6.04 ms	+0.10 ms
All GPU GEMM kernels	3.79 ms	3.78 ms	-0.01 ms
GPU MoE GEMM kernels	1.14 ms	1.13 ms	-0.01 ms
GPU transfer-engine activity	0.09 ms	0.60 ms	+0.51 ms
Host time issuing cudaMemcpyAsync	0.47 ms	1.18 ms	+0.71 ms
Host time in cudaEventSynchronize	0 ms	0.17 ms	+0.17 ms
Time exposed at final cudaStreamSynchronize	0.07 ms	15.42 ms	+15.35 ms

The transfer engines spend only 0.60 ms/token on copies (1.6% of profiled latency). Incremental transfer activity accounts for roughly 2.2% of the regression. GPU GEMM time is unchanged, so neither PCIe bandwidth nor CUDA GEMM execution explains the slowdown.

The offload path currently performs approximately, per decoded token:

• 78 32-byte D2H copies for routing indices and probabilities;
• 38.5 4-KiB D2H copies for input activations;
• 38.5 4-KiB H2D copies for CPU results;
• 264 cudaMemcpyAsync calls in total.

MLAS GEMM represents about 16.5 ms of cumulative CPU-core work per token across worker threads. Sampling estimates roughly 1.1 ms/token on the critical thread, but Nsight Systems cannot give an exact wall-clock MLAS duration without additional internal instrumentation.

Conclusion: the dominant issue is CPU/GPU orchestration and dependency bubbles, not transfer bandwidth or GPU GEMM. The extra 15.35 ms visible in the final stream synchronization represents unfinished work accumulated by the offload pipeline; it is not the intrinsic overhead of the synchronization call itself. The activity categories overlap and should not be added as independent latency components.

The highest-priority optimization is therefore to keep routing counters and hot/cold selection on GPU, compact only CPU-routed rows, notify CPU workers asynchronously instead of synchronizing every MoE node, and wait only when CPU and GPU outputs must be combined.

[Xavier Dupré (xadupre)]

Superseded by #33347, a self-contained main-targeted PR combining static CPU offload and adaptive expert swaps. It documents the actual shared-layout transfer design, the rejected bidirectional-layout experiment and all 33 available in-grid benchmark runs with actual repetition counts. Remaining measurements were stopped as requested. This PR is closed as superseded; its changes are included in the replacement branch.