Enable static CPU offload for FP16/BF16 CUDA MoE

docs/MoEExpertCounting.md

@@ -19,7 +19,7 @@ session.moe_expert_counter_beta=<finite non-negative value>  # default: 0.1
 The coefficients must satisfy `alpha + beta <= 1`; zero is allowed for either coefficient.
 
 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 either counting or counter-update logging is enabled.
+capture are not supported when counting, counter-update logging, or CPU offload is enabled.
 Minimal builds omit the counter state, its initialization/run bookkeeping, and the CUDA routing-snapshot collector
 sources. They reject counter configuration except for explicitly disabling counting.
 
@@ -38,7 +38,7 @@ entire model invocation.
 The root `SessionState` owns a `KernelPilotMoeExpertState` shared with its subgraph states. Sessions never share counters.
 While counter-update logging is active, the state also stores the current request ID and logger. Because expert
 selections and counters are session-owned mutable state, a concurrent Run is rejected whenever counting or
-counter-update logging is enabled.
+counter-update logging is enabled, or when static CPU offload uses the counters.
 Registration uses graph scope and resolved node index, with one counter per expert. The router's expert dimension must
 be statically known when the session is initialized. The global expert count is the sum of those per-node dimensions;
 for `N` equally sized MoE nodes with `E` experts, it is `N * E`.
@@ -124,6 +124,36 @@ kernel skips collector construction, collection calls, and statistics-only size
 There are no statistics allocations, host transfers, or stream synchronizations on that path; only cached flag checks
 remain. Enabling these diagnostics adds overhead.
 
+## Static FP16/BF16 CUDA expert placement
+
+Set the session-wide CPU offload count with:
+
+```text
+session.moe_cpu_offload_experts=<non-negative integer>
+```
+
+This first offload implementation applies to built-in CUDA `MoE` nodes with FP16 or BF16 activations and constant
+expert weights. Before session initialization returns, `KernelPilot` chooses the experts that remain resident in CUDA
+memory. All-zero counters distribute slots round-robin across eligible nodes; initialized counters rank experts by
+descending value with deterministic ties. The offload count is global and must not exceed the number of experts in
+eligible CUDA nodes. The default is `0`, which disables CPU offloading.
+
+Every expert retains a CPU weight representation. Only selected experts are copied to CUDA. During inference, CUDA
+experts run through the existing CUTLASS MoE path. CPU-resident FP16 experts use MLAS FP16 GEMMs; BF16 expert weights
+are converted once during initialization and CPU-resident BF16 experts use MLAS FP32 GEMMs. Their weighted outputs are
+combined on CUDA. Routing still records every selected expert, so counter updates continue unchanged. Each kernel
+retains only one host representation of its FC1/FC2 weights: original weights for all-CUDA nodes, or GEMM-layout weights
+for nodes with CPU experts. The temporary original copy is released after transposition.
+
+Offload requires prepacking to remain enabled and the input hidden dimension to be statically known. Connected
+FC1/FC2 biases must be constant, just like the expert weights.
+Expert-weight collection and validation follow actual outer-scope captures into subgraphs; same-named local inputs
+and initializers are separate values and do not cause ancestor values to be validated or moved to CPU.
+
+Placement is static in this stage: counters continue to evolve, but no weight transfer, expert swap, or end-of-run
+redistribution occurs. Separate FC3 weights, sparse mixer routing, dynamic expert weights, QMoE, FP32, the CUDA plugin
+EP, CUDA graphs, and minimal builds are not supported by this option.
+
 ## Initial counter file
 
 Optionally set:
@@ -132,7 +162,7 @@ Optionally set:
 session.moe_expert_counter_state_file=/path/to/counters.txt
 ```
 
-This requires counting or counter-update logging to be enabled. The UTF-8 text file begins with
+This requires counting, counter-update logging, or static CPU offload to be enabled. The UTF-8 text file begins with
 `moe_expert_state 1`, followed by whitespace-separated records:
 
 ```text

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.
 
-Each expert is resident either on CPU or CUDA. The configured global offload target determines how many experts remain
-on CPU, while the complementary set resides on CUDA. Each `MoE` and `QMoE` node maintains one exponentially decayed
-counter per expert, uses those counters to rank experts, and updates its placement asynchronously.
+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.
 
 The placement policy has two levels, both evaluated by the pilot after a complete model inference:
 
@@ -58,30 +58,28 @@ The four numerical policy parameters are exposed as session configuration entrie
 
 | Session option | Parameter | Meaning and valid range |
 |---|---|---|
-| `session.moe_cpu_offload_experts` | Offload target | Global integer expert count (`>= 1`) or proportion (`0 < value < 1`). |
+| `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, finite and `>= 0`; default `0.9`. |
 | `session.moe_expert_counter_beta` | `beta` | Increment for a used expert, finite and `>= 0`; default `0.1`. |
 | `session.moe_expert_swap_epsilon` | `epsilon` | Relative swap margin, finite and `>= 0`. |
 
 The optional `session.moe_expert_counter_state_file` path is configured separately from these four numerical parameters.
 
-The offload target has the following meaning:
+The CPU offload count has the following meaning:
 
-- An integer greater than or equal to `1` is the total number of experts to offload to CPU.
-- A value strictly between `0` and `1` is the proportion of all experts to offload to CPU. The concrete expert count is
-  `ceil(value * total_expert_count)`.
-- Zero, negative values, non-integral values greater than `1`, non-finite values, and counts larger than the total
-  number of experts are invalid.
-- When the option is absent, expert offloading is disabled.
+- A non-negative integer is the total number of experts to offload to CPU.
+- Negative values, non-integer values, and counts larger than the total number of eligible CUDA FP16/BF16 `MoE` experts
+  are invalid.
+- Zero disables expert offloading.
 
 The global CUDA budget is:
 
 ```text
 cuda_expert_count = total_expert_count - cpu_offload_expert_count
 ```
 
-One CUDA slot contains all weights required to execute one expert. Moving an expert changes its residency: the expert
-leaving CUDA is transferred to CPU before the replacement expert is transferred from CPU to CUDA.
+One CUDA slot contains all weights required to execute one expert. CUDA slots contain copies only; moving an expert
+into or out of CUDA never removes or modifies its canonical CPU weights.
 
 ## Expert counters
 
@@ -121,12 +119,10 @@ The file starts with a format-version line and contains one line per
 
 Experts omitted from the file start at zero. When the setting is absent, all counters start at zero.
 
-The first offload implementation does not use counters to choose its initial placement. It visits participating
-`MoE` and `QMoE` kernels in graph order and fills the global CUDA expert budget from the first kernels. If the budget
-ends within one kernel, the lowest expert IDs of that boundary kernel remain on CUDA. All remaining experts, including
-those of the last kernels, execute on CPU. This placement remains fixed for the lifetime of the session and separates
-hybrid CPU/CUDA inference correctness from the adaptive policy. Counter-based placement changes are added only in the
-second implementation step.
+The first offload implementation uses loaded counters to select the highest-ranked experts globally. If all counters
+are zero, it distributes CUDA slots round-robin across eligible nodes and selects lower local expert IDs first. This
+placement remains fixed for the lifetime of the session and separates hybrid CPU/CUDA inference correctness from
+adaptive placement changes.
 
 ## Per-node placement update
 
@@ -224,7 +220,7 @@ regressing the regular CUDA path, an internal graph transformer may insert an ex
 operator. That operator must reuse `MoE`/`QMoE` schema semantics and kernels and must not become part of the exported
 model contract.
 
-When `session.moe_cpu_offload_experts` is absent, CPU and CUDA `MoE`/`QMoE` behavior remains unchanged.
+When `session.moe_cpu_offload_experts` is `0` or absent, CPU and CUDA `MoE`/`QMoE` behavior remains unchanged.
 
 ## State ownership and concurrency
 
@@ -267,19 +263,29 @@ The session-global expert state and counters are already implemented. The remain
 steps so that hybrid inference is validated before placement starts changing at runtime. Each pull request includes
 the tests and documentation for its own scope.
 
-### Step 1: fixed kernel-order placement and hybrid inference
+### Step 1: static FP16/BF16 placement and hybrid inference
 
-- Parse and validate the count-or-proportion offload target.
-- Fill the global CUDA expert budget from the first `MoE` and `QMoE` kernels in graph order, splitting only the
-  boundary kernel when the budget does not contain a whole number of kernels.
-- Keep the remaining experts, including all experts of the last kernels, on CPU.
+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 this initial placement immutable: this step has no swaps or end-of-inference redistribution.
-- Materialize each expert only on its assigned device.
-- Dispatch resident experts on CUDA and offloaded experts through the shared CPU expert-compute path.
-- Combine CPU and CUDA expert results without changing the exported `MoE` or `QMoE` model contract.
-- Preserve the existing CUDA implementation when offloading is disabled.
-- Test CPU-only, CUDA-only, and mixed expert execution, the global offload count, numerical agreement, bounded memory,
-  repeated inference with a fixed placement, and unchanged disabled-path behavior.
+- 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.
+- Complete the host input copy before CUDA expert kernels can overwrite an aliased input/output buffer, then overlap
+  CPU and CUDA expert GEMMs.
+- Combine CPU and CUDA expert results on CUDA without changing the exported `MoE` model contract.
+- Continue collecting complete routing selections and updating counters.
+- Preserve dynamic FP16/BF16 weights and the existing CUDA implementation when offloading is disabled.
+- Test CPU-only, CUDA-only, and mixed expert execution, numerical agreement, counter updates, invalid configuration,
+  and unchanged disabled-path behavior.
+
+Current limitations are deliberate: constant FC1/FC2 weights are required; FC3, sparse mixer, QMoE, FP32,
+minimal builds, CUDA graph capture, and the CUDA plugin EP are not supported by this first slice.
+Prepacking must remain enabled, the input hidden dimension must be statically known, and connected FC1/FC2 biases
+must also be constant.
 
 ### Step 2: adaptive expert swaps
 

include/onnxruntime/core/framework/op_kernel.h

@@ -36,6 +36,10 @@ class OpKernelContext;
 
 namespace onnxruntime {
 
+#if !defined(ORT_MINIMAL_BUILD)
+class KernelPilot;
+#endif
+
 std::unique_ptr<OpKernelInfo> CopyOpKernelInfo(const OpKernelInfo& info);
 
 class OpKernel {
@@ -170,6 +174,14 @@ class OpKernel {
     return Status::OK();
   }
 
+#if !defined(ORT_MINIMAL_BUILD)
+  // Called after session-global piloting state is finalized and all constant inputs have been offered to PrePack().
+  // Kernels may use this hook to materialize resources selected by the pilot before session initialization returns.
+  virtual Status InitializeKernelPilot(KernelPilot* /*pilot*/) {
+    return Status::OK();
+  }
+#endif
+
   const OrtDevice GetDevice(OrtMemType mem_type) const;
   const OpKernelInfo& Info() const {
     return *op_kernel_info_;

include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h

@@ -43,6 +43,12 @@ static const char* const kOrtSessionOptionsConfigMoeExpertCounterAlpha =
 static const char* const kOrtSessionOptionsConfigMoeExpertCounterBeta =
     "session.moe_expert_counter_beta";
 
+// 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";
+
 // 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";