[WebGPU] Accelerated model loading using DirectStorage APIs

[Sushanth Rajasankar (sushraja-msft)]

Description

This change adds experimental accelerated external-weight loading for the native Windows WebGPU execution provider.

External initializers can now be loaded directly from disk into exact tensor-sized D3D12 default-heap resources using DirectStorage. These resources are subsequently imported into Dawn as WebGPU buffers, avoiding the existing disk → CPU/mmap → BufferManager::Upload path.

The implementation includes:

• One batched DirectStorage submission for all external initializers.
• Requests split at a maximum of 64 MiB.
• Four concurrent D3D12 resource-allocation workers.
• Overlap between resource allocation and DirectStorage initialization.
• One completion signal and wait for the complete batch.
• Support for multiple external-data files and non-contiguous file ranges.
• Validation of file bounds, tensor sizes, DirectStorage status, D3D12 operations, and Dawn import/access results.
• Lifetime management for D3D12 resources, Dawn shared memory, WebGPU buffers, fences, and access state.
• An extended external-data-loader lifecycle that allows ORT to prepare and finalize an initializer batch without affecting other execution providers.
• Early DXGI adapter selection and Dawn LUID pinning for pipelined modes, allowing weight loading to overlap WebGPU device initialization.
• Build-time gating so DirectStorage dependencies are only included for supported native Windows WebGPU builds.

The WebGPU provider exposes one policy option:

"weightLoadAcceleration": "preferred-pipelined"

Supported values are:

• off: Use the existing initializer loading path.
• preferred: Attempt accelerated disk-to-GPU loading and fall back to ordinary loading.
• preferred-pipelined: Additionally attempt to overlap loading with WebGPU device initialization. If early pipelining is unavailable, continue non-pipelined where possible.
• required: Require accelerated disk-to-GPU loading, but not pipelining.
• required-pipelined: Require both accelerated loading and the pipelined initialization path.

Cancellation is always propagated and is never converted into a successful fallback.

Motivation and Context

Large models commonly store most of their parameters as external ONNX initializer data. On the native WebGPU path, these weights were previously read or mapped into CPU memory and then copied into Dawn-created GPU buffers.

For multi-gigabyte models, this causes several avoidable startup costs:

• CPU-side materialization of data that is only consumed by the GPU.
• Additional memory pressure and address-space activity.
• A separate GPU upload for every initializer.
• Serialized weight loading and WebGPU device initialization.
• CPU, disk, and driver work on the critical path before session creation can finish.

DirectStorage and Metal I/O provide platform-specific mechanisms for loading file data directly into device-local resources. The new weightLoadAcceleration policy describes this behavior independently of the platform implementation, allowing additional native WebGPU backends to provide equivalent functionality in the future.

The pipelined modes select the target DXGI adapter early, create its D3D12 device, and force Dawn to select the same adapter using its LUID. ORT can then parse the model and begin loading validated external initializer ranges while Dawn completes adapter and device initialization.

The final import path verifies that DirectStorage and Dawn resolved to the same D3D12 device before exposing any loaded initializer to ORT. Preferred modes discard incomplete resources and fall back safely; required modes report the underlying error.

Observed Measurements

Measurements were collected using an approximately 2-billion-parameter-class decoder-only model with:

• 1,291,602,752 bytes of external initializer data.
• 454 external tensors.
• 457 DirectStorage requests.
• Runtime graph optimization disabled after offline optimization.
• Five or six fresh processes per reported median, with no inference warmup for cold first-token measurements.

Session creation

Metric	off	preferred	preferred-pipelined
Median session creation	1.592 s	0.749 s	0.609 s
Observed range	1.425–1.685 s	0.673–0.796 s	0.559–0.642 s
Improvement relative to off	—	843 ms / 53.0%	983 ms / 61.7%

Pipelining reduced median session creation by approximately 141 ms, or 18.8%, relative to non-pipelined accelerated loading.

The off measurements were collected in a separate six-run batch. Filesystem caching, GPU power state, and driver initialization can affect comparisons between batches.

Cold first token

Using a one-token prompt and generating one token in five fresh processes:

Metric	off	preferred	preferred-pipelined
Session median plus prompt TTFT median	1.844 s	0.976 s	0.805 s
Median prompt-to-first-token time	251.8 ms	226.5 ms	196.0 ms
Median complete one-token process time	3.014 s	2.106 s	2.023 s
Prompt-to-first-token range	214.3–278.3 ms	195.1–242.5 ms	191.1–206.5 ms
Complete process-time range	2.722–3.966 s	2.077–2.134 s	1.869–2.085 s
Dedicated GPU memory	1.31 GiB	1.25 GiB	1.25 GiB

Compared with off, preferred-pipelined reduced the combined median session-creation and prompt-processing interval by approximately 1.039 seconds, or 56.4%.

The session-plus-TTFT values combine medians from separate session and inference benchmark processes. Complete process time additionally includes executable startup, configuration and tokenizer loading, and process teardown.

Representative pipelined load

Phase	Time
Early DXGI selection and D3D12 device creation	199.7 ms
ORT work before external loading begins	approximately 48.5 ms
Remaining Dawn initialization after early selection	approximately 136.0 ms
DirectStorage initialization	132.0 ms
D3D12 allocation	128.0 ms
Overlapped preparation	132.1 ms
Request enqueue	0.30 ms
File-to-GPU transfer	194.2 ms
Dawn import and access	0.67 ms
Initializer handoff and remaining session finalization	approximately 44.7 ms
Complete ORT session creation	627.5 ms

Approximately 87.5 ms of Dawn initialization overlapped DirectStorage preparation and transfer in this representative trace. LUID pinning also avoids Dawn’s ordinary multi-adapter discovery path, so the total improvement is not solely attributable to the overlap interval.

[Copilot]

🟡 Changes recommended

Build-definition propagation, required-pipelined semantics, cleanup, empty tensors, and coverage have unresolved issues.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds experimental DirectStorage-based external-weight loading for native Windows WebGPU.

Changes:

• Adds batched disk-to-GPU loading and Dawn resource import.
• Extends external-data loader lifecycle and asynchronous WebGPU initialization.
• Adds configuration, packaging, benchmark support, and tests.

File summaries

File	Description
tools/nuget/generate_nuspec_for_native_nuget.py	Packages DirectStorage runtimes.
onnxruntime/test/providers/webgpu/webgpu_context_test.cc	Tests acceleration modes.
onnxruntime/test/perftest/ort_test_session.cc	Passes WebGPU benchmark options.
onnxruntime/test/framework/external_data_loader_test.cc	Tests loader batch lifecycle.
onnxruntime/core/session/inference_session.h	Tracks preload state.
onnxruntime/core/session/inference_session.cc	Starts external-data preloading.
onnxruntime/core/providers/webgpu/webgpu_provider_options.h	Defines acceleration policies.
onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc	Parses the new option.
onnxruntime/core/providers/webgpu/webgpu_execution_provider.h	Adds DirectStorage state.
onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc	Integrates loader and allocator.
onnxruntime/core/providers/webgpu/webgpu_context.h	Adds asynchronous initialization state.
onnxruntime/core/providers/webgpu/webgpu_context.cc	Implements pipelined device initialization.
onnxruntime/core/providers/webgpu/external_data_loader.h	Updates loader signature.
onnxruntime/core/providers/webgpu/external_data_loader.cc	Adapts WebAssembly loader.
onnxruntime/core/providers/webgpu/direct_storage_external_data_loader.h	Declares DirectStorage loading APIs.
onnxruntime/core/providers/webgpu/direct_storage_external_data_loader.cc	Implements batched loading and import.
onnxruntime/core/providers/webgpu/data_transfer.h	Defers buffer-manager lookup.
onnxruntime/core/providers/webgpu/data_transfer.cc	Uses deferred lookup.
onnxruntime/core/providers/webgpu/allocator.h	Defers UMA detection.
onnxruntime/core/providers/webgpu/allocator.cc	Initializes allocation mode lazily.
onnxruntime/core/providers/js/external_data_loader.h	Updates JS loader signature.
onnxruntime/core/providers/js/external_data_loader.cc	Adapts JS loader implementation.
onnxruntime/core/framework/tensorprotoutils.h	Adds external-data preparation API.
onnxruntime/core/framework/tensorprotoutils.cc	Validates and prepares external tensors.
onnxruntime/core/framework/session_state_utils.cc	Batches initializer preparation/loading.
onnxruntime/core/framework/external_data_loader.h	Extends loader lifecycle interface.
onnxruntime/core/framework/external_data_loader.cc	Provides default lifecycle hooks.
onnxruntime/core/framework/external_data_loader_manager.h	Adds preload and batch management.
onnxruntime/core/framework/external_data_loader_manager.cc	Coordinates loader lifecycles.
cmake/onnxruntime_providers_webgpu.cmake	Configures DirectStorage dependencies.
cmake/CMakeLists.txt	Adds DirectStorage build options.

Review details

• Files reviewed: 31/31 changed files
• Comments generated: 6
• Review effort level: Balanced

---

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

[Copilot]

🟡 Changes recommended

Cancellation propagation, queue sizing, preload filtering, and initialization portability contain blocking defects.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

• Files reviewed: 31/31 changed files
• Comments generated: 5
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

Cancellation can be lost, a lazy callback can outlive its provider, and the feature lacks an enabled CI build.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

• Files reviewed: 31/31 changed files
• Comments generated: 3
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

Preloading breaks caller-supplied external initializer overrides, and packaged feature availability is inconsistent across Windows Python wheels.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

• Files reviewed: 32/32 changed files
• Comments generated: 2
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

The Windows wheel omits required runtime DLLs, and unresolved loader compatibility, boolean normalization, and cancellation defects remain.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (4)

Previously missed (2) — in code that hasn't changed since the last review.

cmake/onnxruntime_providers_webgpu.cmake:16

• This gate still permits onnxruntime_USE_EP_API_ADAPTERS, but the OrtEp adapter has no external-data-loader callback/bridge: the wrapped WebGPU provider's GetExternalDataLoader() is never registered with InferenceSession. In that build, every acceleration mode is ineffective and even required silently uses ordinary loading. Either implement an EP-API bridge or reject adapter/plugin builds here.
  onnxruntime/core/providers/webgpu/direct_storage_external_data_loader.cc:728
• The zero-request preload fast path never evaluates is_cancelled, so an all-empty external-initializer model can report successful preload even when cancellation is already requested. This contradicts the stated cancellation guarantee; continue initialization, but return MODEL_LOAD_CANCELED when the callback is true.

This issue also appears in the following locations of the same file:

• line 807
• line 860

onnxruntime/core/providers/webgpu/direct_storage_external_data_loader.cc:819

• The zero-request final batch likewise ignores the current cancellation callback. If there was no non-empty preload (for example, all external tensors are empty), cancellation requested after preparation is converted into success. Check is_cancelled in this branch and return MODEL_LOAD_CANCELED through fail_or_fallback.

  if (batch.request_count == 0) {
    common::Status preload_status = common::Status::OK();
    if (impl_->preload_future.valid()) {
      preload_status = impl_->preload_future.get();
    }
    impl_->preload_batch.reset();
    impl_->context.ContinueInitialize();
    if (!preload_status.IsOK()) {
      return fail_or_fallback(preload_status);
    }
    impl_->context.WaitForInitializeComplete();
    batch.finalized = true;
    return common::Status::OK();

onnxruntime/core/providers/webgpu/direct_storage_external_data_loader.cc:864

• After the DirectStorage transfer completes, cancellation is no longer polled while this potentially long Dawn initialization wait and the subsequent per-tensor import run. In pipelined mode, a cancellation arriving during that interval is therefore returned as a successful load. Recheck the callback after the wait (and during a large import loop) and propagate MODEL_LOAD_CANCELED via fail_or_fallback.

  impl_->context.WaitForInitializeComplete();

  const auto import_start = Clock::now();
  const auto import_status = [&]() -> common::Status {
    ORT_RETURN_IF_NOT(

• Files reviewed: 32/32 changed files
• Comments generated: 2
• Review effort level: Balanced