Enable asynchronous adaptive FP16/BF16 MoE expert swaps

docs/MoEExpertCounting.md

@@ -12,10 +12,12 @@ The default is `0`. Counting alone does not emit logs, enable profiling, change
 Configure the update with:
 
 ```text
-session.moe_expert_counter_alpha=<finite value in [0, 1]> # default: 0.9
+session.moe_expert_counter_alpha=<finite value in [0, 1]> # default: 0.75
+session.moe_expert_counter_probability_weight=<finite value in [0, 1]> # default: 0
 ```
 
-The selected-expert increment is `1 - alpha`.
+The selected-expert increment is `1 - alpha`. With CPU offload, the optional probability weight `lambda` changes the
+selected contribution from `1` to `(1 - lambda) + lambda * routing_probability`.
 
 It supports the CPU and built-in CUDA `MoE` and `QMoE` kernels. Minimal builds, the CUDA plugin EP, and CUDA graph
 capture are not supported when counting, counter-update logging, or CPU offload is enabled.
@@ -25,10 +27,12 @@ sources. They reject counter configuration except for explicitly disabling count
 Each node invocation updates every expert counter belonging to that node:
 
 ```text
-c(t+1) = alpha * c(t) + (1 - alpha) * (1 if used otherwise 0)
+c(t+1) = alpha * c(t) + (1 - alpha) * contribution
 ```
 
-Selecting an expert for several token rows still contributes one `1 - alpha`. Unused experts still receive the decay term.
+By default, `contribution` is `1` for a selected expert and `0` otherwise. With probability weighting, repeated
+selections use the maximum routing probability observed for that expert during the invocation. Unused experts still
+receive the decay term.
 Counts persist across `Run()` calls. Only successful kernel invocations commit their collected usage.
 A failed run can retain updates from earlier successful kernels, including CUDA nodes whose
 expert computations were still in flight when their counters were updated. Updates are not transactional across an
@@ -56,6 +60,14 @@ Counter updates continue unchanged after the limit, and the logging budget reset
 The Python analyzer rejects any trace containing this marker, even with an otherwise valid completion footer.
 When INFO output is disabled, counter records are neither formatted nor charged against the logging budget.
 Only logging-budget reservation uses an atomic; counting-only updates remain unsynchronized across disjoint kernel ranges.
+
+After each successful CPU-offload `Run()`, INFO logging also emits one `moe_expert_swaps` JSON record. It reports
+`in_flight` swaps and `blocked_candidates` per CUDA device, plus session-wide totals. Its `scheduled_swaps` array
+identifies each exchange started by that run with the graph scope, node index, device, outgoing CUDA expert, incoming
+CPU expert, and both pre-swap scores. A blocked candidate is currently eligible for a swap but was not started because
+the per-device concurrency limit was reached. Candidates are recomputed after every run rather than retained in a
+persistent queue. This record is not part of the counter-record logging budget.
+
 A `KernelPilot` is allocated for each kernel at initialization; its `KernelPilotMoeExpertSelection` member counts repeated routing IDs only once.
 Each counter is updated in place. The coefficient constraints keep it bounded by the larger of its initial value
 and `1`, so no next-value buffer or overflow-validation pass is needed.

docs/moe-cpu-offload/README.md

@@ -9,9 +9,9 @@
 Implement adaptive expert placement for Qwen 3.6 and other Mixture-of-Experts (MoE) models whose expert weights do not
 all fit in GPU memory.
 
-CPU memory keeps the canonical copy of every expert. CUDA holds a configurable subset of expert copies. Each `MoE` and
-`QMoE` node maintains one exponentially decayed counter per expert, uses those counters to rank experts, and updates
-its CUDA placement asynchronously.
+CPU and CUDA each hold a disjoint, configurable subset of experts. Each `MoE` and `QMoE` node maintains one
+exponentially decayed counter per expert, uses those counters to rank experts, and exchanges experts between CPU and
+CUDA asynchronously.
 
 The placement policy has two levels, both evaluated by the pilot after a complete model inference:
 
@@ -28,6 +28,13 @@ Training, router changes, expert-weight quantization, and multiple CUDA devices
 - [#32738](https://github.com/microsoft/onnxruntime/pull/32738) added the session-global expert state, CPU and CUDA
   expert-selection collection, and exponentially decayed counters. It does not implement placement, CPU offload,
   swaps, or redistribution.
+- [#33039](https://github.com/microsoft/onnxruntime/pull/33039) added static FP16/BF16 CUDA `MoE` placement, hybrid
+  CPU/CUDA execution, and compact CUDA expert storage.
+- The adaptive-swap implementation exchanges hot CPU experts with cold CUDA experts asynchronously while preserving
+  each node's CUDA slot count. Global slot redistribution across nodes remains future work.
+- Expert swaps transpose matrix weights between CUDA and MLAS layouts on their dedicated CUDA streams, rather than
+  blocking the transfer pipeline with a CPU transposition callback. Incoming and outgoing experts use separate pinned
+  buffers; publication waits for both transfers to finish before changing the CPU/CUDA maps.
 
 ## Exploratory routing analysis
 
@@ -59,8 +66,10 @@ The numerical policy parameters are exposed as session configuration entries:
 | Session option | Parameter | Meaning and valid range |
 |---|---|---|
 | `session.moe_cpu_offload_experts` | Offload count | Global number of experts to execute from CPU (`>= 0`); default `0`. |
-| `session.moe_expert_counter_alpha` | `alpha` | Counter decay coefficient in `[0, 1]`; selected experts receive `1 - alpha`; default `0.9`. |
-| `session.moe_expert_swap_epsilon` | `epsilon` | Relative swap margin, finite and `>= 0`. |
+| `session.moe_expert_counter_alpha` | `alpha` | Counter decay coefficient in `[0, 1]`; selected experts receive `1 - alpha`; default `0.75`. |
+| `session.moe_expert_counter_probability_weight` | Probability weight | Blend in `[0, 1]` between binary selection and routing probability; default `0`. |
+| `session.moe_expert_swap_epsilon` | `epsilon` | Relative swap margin, finite and `>= 0`; default `0.01`. |
+| `session.moe_max_in_flight_swaps_per_device` | Swap concurrency | Maximum concurrent adaptive swaps per CUDA device (`>= 0`); `0` disables adaptive swaps; default `4`. |
 
 The optional `session.moe_expert_counter_state_file` path is configured separately.
 
@@ -92,8 +101,12 @@ count_{t+1}(e) = alpha * count_t(e) + (1 - alpha) * (1 if expert e was used, oth
 The used term is binary for one invocation. Selecting the same expert for several rows still contributes `1`, not the
 number of routed rows.
 
+For CPU-offloaded CUDA MoE nodes, `session.moe_expert_counter_probability_weight=lambda` changes the selected-expert
+contribution to `(1 - lambda) + lambda * probability`. Repeated selections within one invocation use the maximum
+routing probability observed for that expert. `0` preserves binary counting and `1` uses only routing probability.
+
 CPU-offloaded CUDA nodes update their counters on the GPU. The pilot consumes counter snapshots and selected-expert
-metadata at the end of the run rather than waiting for routing on the host at every node.
+metadata at the end of the run, before scheduling swaps, rather than waiting for routing on the host at every node.
 Only successful kernel invocations are committed, and asynchronous CPU execution errors are propagated by `EndRun`.
 Repeated invocations of the same node consume the preceding snapshot before reusing its buffers.
 
@@ -110,6 +123,46 @@ the larger of their initial value and `1`. Their defaults and constraints are co
 Expert identity is `(graph_scope, node_index, node_type, expert_id)`. Ranking is by descending counter. Ties are resolved
 by `expert_id`, then graph scope and node index, so placement is deterministic.
 
+When counter logging is enabled, each event reports only the experts selected by that invocation. Each selected-expert
+record contains the local `expert_id`, its pre-update `score`, its zero-based `selected_rank` among the selected
+experts, its zero-based `node_rank` among every expert in the node, its zero-based `global_position` among every
+expert in the model, and its current `device` (`CPU` or `CUDA`). Ranks use descending score. Local ties are broken by
+expert ID, while global ties are broken by the fixed model-wide expert index. The device reflects the currently
+published placement; an asynchronous swap is not visible until the kernel publishes it. Each event also reports
+`moe_count` and `total_expert_count`.
+
+At the end of each successful run, `moe_expert_swaps` reports the number of in-flight swaps and eligible candidates
+blocked by the concurrency limit, both per CUDA device and in total. Candidates are recalculated after every run and
+are not retained in a persistent queue. `scheduled_swaps` identifies every exchange started by that run, including
+the node, the outgoing CUDA expert, the incoming CPU expert, and their counter scores.
+
+```json
+"selected_experts": [
+  {
+    "expert_id": 7,
+    "score": 0.1,
+    "selected_rank": 0,
+    "node_rank": 3,
+    "global_position": 1024,
+    "device": "CUDA"
+  }
+]
+```
+
+```json
+"scheduled_swaps": [
+  {
+    "graph_scope": "main",
+    "node_index": 12,
+    "device_id": 0,
+    "outgoing_cuda_expert": 41,
+    "incoming_cpu_expert": 173,
+    "outgoing_score": 0.05,
+    "incoming_score": 0.2
+  }
+]
+```
+
 ## Initial counter state
 
 An optional session setting names a UTF-8 text file containing initial counter values:
@@ -185,14 +238,14 @@ inference t+1, node L
 ```
 
 CUDA devices expose their copy-engine count, but that value does not provide a portable expert-level concurrency
-guarantee. The initial implementation therefore permits at most two in-flight expert exchanges per CUDA device. Two
-independent pinned buffers, two extra CUDA staging slots, and separate device-to-host and host-to-device streams form a
-bidirectional pipeline. After exchange A finishes its device-to-host transfer, its host-to-device transfer may overlap
-the device-to-host transfer of exchange B. Events preserve the CPU-first ordering within each exchange; buffers and
+guarantee. The initial implementation therefore permits at most four in-flight expert exchanges per CUDA device.
+Independent pinned buffers, CUDA staging slots, and device-to-host and host-to-device streams form a bidirectional
+pipeline. After exchange A finishes its device-to-host transfer, its host-to-device transfer may overlap the
+device-to-host transfer of another exchange. Events preserve the CPU-first ordering within each exchange; buffers and
 CUDA slots are never shared by transfers that overlap. Hardware with one copy engine serializes the transfers without
 changing correctness. Additional exchanges remain queued for a later completion or inference boundary. For the
-measured Qwen model, where one QMoE expert occupies 1,775,616 bytes, this limit requires about 3.4 MiB of pinned staging
-memory and 3.4 MiB of temporary CUDA storage.
+measured Qwen model, where one QMoE expert occupies 1,775,616 bytes, this limit requires about 6.8 MiB of pinned staging
+memory and 6.8 MiB of temporary CUDA storage.
 
 ## End-of-inference redistribution
 
@@ -208,7 +261,7 @@ The allocation objective is lexicographic:
 Within each node, keep the experts with the highest counters. Redistribution may transfer slot ownership between
 nodes, whereas a per-node exchange changes the expert stored in a slot without changing that node's slot count.
 
-Redistribution never drains or waits for pending exchanges. It schedules up to the available two-exchange concurrency
+Redistribution never drains or waits for pending exchanges. It schedules up to the available four-exchange concurrency
 limit and leaves additional non-conflicting exchanges queued. A node with an incomplete exchange continues using its
 published pre-exchange placement. Slot metadata changes only after both transfer directions complete.
 
@@ -250,11 +303,11 @@ counter-update logic. These are internal C++ calls, not a public C API.
 Expert identity includes graph scope as well as node and expert IDs, so nodes in different
 subgraphs cannot collide. The state persists across `Run()` calls and is isolated from other sessions.
 
-The CUDA cache manager owns device-specific resources and execution state:
+Each CUDA kernel cache owns device-specific resources and execution state:
 
-- CUDA slots, two extra staging slots, and current immutable mappings;
-- two reusable pinned host staging buffers;
-- pending exchanges and redistribution transfers;
+- CUDA slots, one reusable staging slot, and the current immutable mapping;
+- one reusable pinned host staging buffer;
+- at most one pending exchange;
 - CUDA completion events.
 
 Initialization builds an immutable dictionary from `(OpKernel pointer, local expert ID)` to a global expert index.
@@ -279,8 +332,8 @@ Implemented for the built-in CUDA FP16 and BF16 `MoE` paths:
 
 - Parse and validate `session.moe_cpu_offload_experts`; `0` disables offloading.
 - Use loaded counters to select CUDA experts globally; distribute an all-zero budget round-robin across eligible nodes.
-- Keep one representation of all constant expert weights on CPU and materialize only selected expert slices in compact
-  CUDA storage. All-CUDA nodes retain the original host weights; nodes with CPU experts retain the GEMM-layout weights.
+- Keep only CPU-resident experts in compact GEMM-layout CPU storage and materialize CUDA-resident experts in compact
+  CUDA storage. A swap exchanges one expert between these fixed-capacity stores through reusable staging buffers.
 - Keep this initial placement immutable: this step has no swaps or end-of-inference redistribution.
 - Dispatch CUDA-resident routes through CUTLASS. CPU-resident FP16 routes use MLAS FP16 GEMMs; BF16 weights are
   converted once during initialization and CPU-resident BF16 routes use MLAS FP32 GEMMs.
@@ -299,20 +352,28 @@ must also be constant.
 
 ### Step 2: adaptive expert swaps
 
+Implemented in the adaptive-swap change:
+
 - Use the session-global counters and optional initial counter state to rank experts.
 - Apply the strict `cpu_max > (1 + epsilon) * cuda_min` rule and let the pilot schedule exchanges after inference.
-- Move the CUDA expert to CPU before moving its replacement to CUDA.
-- Manage CUDA slots, two staging slots, two pinned buffers, dedicated device-to-host and host-to-device streams,
-  completion events, immutable per-invocation mappings, and atomic publication of completed swaps.
-- Permit at most two in-flight exchanges per CUDA device and queue the rest.
-- Pipeline the host-to-device transfer of one exchange with the device-to-host transfer of the other when the hardware
-  exposes bidirectional copy engines.
-- Never wait for an incomplete exchange when a `MoE` starts; use the pre-exchange placement for that invocation.
-- Redistribute the global CUDA expert budget after inference without draining pending exchanges, while maximizing the
-  number of completely CUDA-resident nodes.
-- Test the epsilon boundary, transfer ordering, two-exchange concurrency, queued exchanges, nonblocking use of the old
-  mapping, asynchronous publication, global budget preservation, counter-based placement, and explicit transfer
-  failures.
+- Move the CUDA expert to pinned CPU staging before moving its replacement to the extra CUDA slot.
+- Use dedicated device-to-host and host-to-device streams, events, and separate staging allocations for each active
+  exchange.
+- Permit at most four in-flight exchanges per CUDA device; additional eligible nodes are reconsidered at the next
+  inference boundary.
+- Pipeline the host-to-device transfer of one exchange with the device-to-host transfer of another when the hardware
+  supports bidirectional copies.
+- Never wait for an incomplete exchange when a `MoE` starts. The old placement remains published until transfer
+  completion; publication copies the staged expert into the persistent slot on that invocation's compute stream.
+- Preserve each node's CUDA slot count and therefore the session-global CUDA expert budget.
+- Test the strict epsilon boundary, deterministic tie-breaking, multi-exchange concurrency, delayed
+  publication, and numerical parity before and after a swap.
+
+Still planned:
+
+- Redistribute CUDA slot ownership across nodes without draining pending exchanges, while maximizing the number of
+  completely CUDA-resident nodes.
+- Add end-to-end throughput and memory measurements for representative Qwen models.
 
 After these two implementation steps, the remaining work is end-to-end measurement. Run reproducible CPU-only,
 CUDA-only, and hybrid evaluations with identical models, prompts, and generation settings; sweep offload targets and

include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h

@@ -35,16 +35,31 @@ static const char* const kOrtSessionOptionsConfigEnableMoeExpertCounting =
 static const char* const kOrtSessionOptionsConfigMoeExpertCounterStateFile =
     "session.moe_expert_counter_state_file";
 // Exponential decay applied to every expert counter after each invocation of its MoE/QMoE node.
-// Must be finite and in [0, 1]. A selected expert receives an increment of 1 - alpha. The default is 0.9.
+// Must be finite and in [0, 1]. A selected expert receives an increment of 1 - alpha. The default is 0.75.
 static const char* const kOrtSessionOptionsConfigMoeExpertCounterAlpha =
     "session.moe_expert_counter_alpha";
 
+// Blend between binary selection and routing-probability contributions to expert counters.
+// Must be finite and in [0, 1]. Zero preserves binary counting; one uses routing probabilities. The default is 0.
+static const char* const kOrtSessionOptionsConfigMoeExpertCounterProbabilityWeight =
+    "session.moe_expert_counter_probability_weight";
+
 // Number of FP16/BF16 MoE experts to execute on CPU across the session. All expert weights remain in CPU memory, while
 // CUDA-resident copies are created for the remaining experts. The value must be a non-negative integer no greater
 // than the number of experts owned by CUDA MoE nodes. The default is 0, which disables CPU offloading.
 static const char* const kOrtSessionOptionsConfigMoeCpuOffloadExperts =
     "session.moe_cpu_offload_experts";
 
+// Relative margin required before an adaptive MoE swap replaces a CUDA-resident expert.
+// The value must be finite and non-negative. The default is 0.01.
+static const char* const kOrtSessionOptionsConfigMoeExpertSwapEpsilon =
+    "session.moe_expert_swap_epsilon";
+
+// Maximum number of adaptive MoE expert swaps that may be in flight concurrently on one CUDA device.
+// Must be a non-negative integer. Zero disables adaptive swaps. The default is 4.
+static const char* const kOrtSessionOptionsConfigMoeMaxInFlightSwapsPerDevice =
+    "session.moe_max_in_flight_swaps_per_device";
+
 // A value of "1" means allocators registered in the env will be used. "0" means the allocators created in the session
 // will be used. Use this to override the usage of env allocators on a per session level.
 static const char* const kOrtSessionOptionsConfigUseEnvAllocators = "session.use_env_allocators";