WebGPU: Graph capture mode should handle INT64 transparently without per-op opt-in

[feich-ms]

Summary

When enable_graph_capture=true, the WebGPU EP requires all nodes to be partitioned to the WebGPU provider — any CPU fallback breaks graph capture. However, INT64 support is currently gated behind a per-op allowlist (enable_int64), which only covers Cast, Unsqueeze, Expand, and Range. Any model that uses INT64 inputs with other ops (ReduceSum, Sub, Equal, Where, etc.) fails at runtime with:

"This session cannot use the graph capture feature as requested by the user as all compute graph nodes have not been partitioned to the WebGpuExecutionProvider"

This forces exporters to work around the limitation by inserting explicit INT32 casts throughout the model graph, which adds complexity that doesn't belong in the model layer.

Current behavior

• webgpu_execution_provider.cc:RegisterKernels(bool enable_graph_capture, bool enable_int64) uses an allowlist approach — only Cast, Unsqueeze, Expand, Range get INT64 type registration when enable_int64=true.
• When graph capture is enabled, enable_int64 is always set to true (line 531-532), but this only helps the 4 ops above.
• Any other op encountering INT64 tensors falls back to CPU, which is fatal under graph capture.

Problem

This design doesn't scale. Every new model architecture that introduces INT64 usage in a new op will hit this wall, requiring either:

1. A new ORT PR to add that op to the enable_int64 allowlist, or
2. The exporter to insert INT32 workarounds throughout the graph

Both are whack-a-mole approaches that accumulate tech debt over time.

Proposed solutions (in order of preference)

Option A: Transparent INT64→INT32 lowering by the EP (best)

When graph capture is enabled, the EP's graph partitioner could auto-insert Cast(INT64→INT32) / Cast(INT32→INT64) nodes around ops that don't natively support INT64, rather than falling back to CPU. The model stays clean INT64 throughout; the EP handles type adaptation internally.

This is the cleanest solution because:

• Model exporters don't need to think about EP type limitations
• Kernel authors don't need to explicitly register INT64 for every op
• New ops automatically work without code changes

Option B: Default-on INT64 registration for all ops under graph capture

Flip the enable_int64 logic from an allowlist to a denylist: when enable_graph_capture=true, ALL ops register INT64 support by default, with an explicit denylist only for ops where INT64 is genuinely impossible (e.g., ops that use WebGPU shader features incompatible with 64-bit integers).

Option C (short-term): Extend the allowlist to common ops

Add INT64 type registration for ReduceSum, Sub, Equal, and Where. The pattern already exists — these ops use GetOpTypeConstraints() which accepts enable_int64. This is ~150 lines of boilerplate but doesn't solve the architectural issue.

Concrete example

Exporting a Gemma 4 model for WebGPU graph capture. The attention mask logic uses:

• ReduceSum(attention_mask, axis=1) — attention_mask is INT64
• Sub(sum_result, 1) — arithmetic on the sum
• Equal(input_ids, pad_token_id) — comparison
• Where(mask, embed_a, embed_b) — conditional selection

All of these fail under graph capture because their INT64 kernels aren't registered. The current workaround is to branch all export logic on enable_graph_capture and emit INT32 variants, but this pollutes the model exporter with EP-specific concerns.

Environment

• ORT version: main branch (tested June 2026)
• EP: WebGPU with enable_graph_capture=true
• Relevant code: onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc lines 460-534