How to split operations in each layer so some run on CPU, others run on GPU

[Marco Kurzynski (mjkpolo)]

Hi,

I was wondering if it's possible to select certain operations to run on the CPU and others for the GPU.
For example, I want to take a model like Llama and try using the CPU for operations with low arithmetic intensity like RMSNorm and see what the impact is.

Thanks!

[Advait Kawale (Advait251206)]

Hi Marco,

ONNX Runtime can partition a model across execution providers, but the standard session API generally does not provide a simple per-node switch such as “run this RMSNorm node on CPU and the remaining nodes on CUDA.”

When you register multiple execution providers, ORT assigns supported nodes or subgraphs according to provider priority and capability. For example, with CUDA followed by CPU, CUDA gets the opportunity to claim supported nodes, while nodes not assigned to CUDA can run on CPU. This is primarily capability-based placement, not a user-defined list of individual nodes.

For your Llama/RMSNorm experiment, a few approaches are possible:

• Inspect the existing partition first. Enable ORT profiling and verbose logging to see which nodes are assigned to each execution provider. This will help establish whether RMSNorm is already running on the CPU or being fused into a larger GPU subgraph.
• Use a custom execution provider or graph transformation if you need precise node-level placement within one session. This is more involved than configuring the standard CUDA and CPU providers.
• Split the graph into separate subgraphs/sessions at the desired boundary and explicitly run one part on CPU and another on GPU. This gives more control, but may introduce host/device copies and synchronization overhead that outweigh the CPU savings.
• Compare end-to-end latency, not just operator time. RMSNorm has relatively low arithmetic intensity, but moving activations between CPU and GPU can be expensive. GPU kernel launch and synchronization costs also matter, so profile the complete inference path.

The relevant starting points are the [Execution Provider documentation](https://onnxruntime.ai/docs/execution-providers/) and [ONNX Runtime profiling guide](https://onnxruntime.ai/docs/performance/tune-performance/profiling-tools.html).

If you need exact per-node control, it would help to know which execution provider you’re using (CUDA, TensorRT, ROCm, etc.) and whether you’re using the standard ONNX Runtime session or an ORT-optimized Llama model.

If this answer helped or pointed you in the right direction, I'd appreciate it if you could mark it as the accepted answer so it's easier for others with the same issue to find.

Also, if you found my contribution useful, I'd appreciate it if you could check out my GitHub profile, follow me, and star any repositories you find interesting.

GitHub: https://github.com/Advait251206

[jinwoo (dev-jinwoohong)]

ORT 1.30.0 has a more direct option for this: session.name_based_layer_assignment. It matches substrings of ONNX node names, so with a CUDA-enabled build you can try:

import onnxruntime as ort

opts = ort.SessionOptions()
opts.enable_profiling = True
opts.add_session_config_entry(
    "session.name_based_layer_assignment",
    "cpu(input_layernorm, post_attention_layernorm)",
)
session = ort.InferenceSession(
    "model.onnx",
    sess_options=opts,
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
)

Replace those patterns with the actual names in your exported graph. They match node names, not operator types. Unmatched nodes still follow normal EP capability/priority assignment, so CUDA can handle the rest where supported.

The requested CPU operators still need matching CPU kernels, and fusions can change the graph you're measuring. After running inference, call session.end_profiling() to get the profiling JSON path. Check the node placement there and compare end-to-end latency, including the extra CPU/GPU transfers. This gives you an experiment within one session without first splitting the model.