Skip to content

[WebGPU] Support QuickGelu activation in the im2col-matmul path - #32560

Merged
Jiajia Qin (qjia7) merged 1 commit into
microsoft:mainfrom
xhcao:webgpu-im2col-quickgelu
Sep 14, 2026
Merged

Jiajia Qin (qjia7) merged 1 commit into
microsoft:mainfrom
xhcao:webgpu-im2col-quickgelu

Conversation

@xhcao

@xhcao xhcao commented Sep 11, 2026

Copy link
Copy Markdown
Contributor

Description

Add QuickGelu to the activation epilogue of the WebGPU im2col-matmul Conv
path (Im2ColMatMulProgram).

Motivation and Context

yolov8*-fp16 and yolov8*-pose-fp16 used to run Conv through the
im2col-matmul path. After Conv+QuickGelu fusion landed, IsActivationSupported()
rejected QuickGelu, so CanApplyIm2ColMatMulProgram() returned false and these
models silently fell back to the slower generic Conv path — the fusion made
them regress.

This restores the im2col path for those models, and they now also keep the
gain from having the activation fused into the epilogue.

Copilot AI balanced review requested due to automatic review settings September 11, 2026 07:52
@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

@xhcao

xhcao commented Sep 11, 2026

Copy link
Copy Markdown
Contributor Author

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟢 Approval recommended

The implementation handles both QuickGelu variants consistently and includes targeted parity coverage.

Pull request overview

Adds QuickGelu support to the WebGPU im2col-matmul Conv epilogue, restoring the optimized path for fused models.

Changes:

  • Supports QuickGelu with unit and arbitrary alpha values.
  • Updates generated WGSL template snapshots.
  • Adds parity tests covering both alpha variants.

No actionable issues found.

File summaries
File Description
onnxruntime/core/providers/webgpu/nn/im2col_matmul.cc Enables and configures QuickGelu.
onnxruntime/core/providers/webgpu/nn/im2col_matmul.h Retains full activation configuration.
onnxruntime/core/providers/webgpu/nn/im2col_matmul.wgsl.template Implements QuickGelu epilogues.
onnxruntime/test/optimizer/graph_transform_test.cc Adds QuickGelu parity tests.
tools/python/wgsl_template/test/in_tree_golden/static-cpp/index.h Updates generated parameters.
tools/python/wgsl_template/test/in_tree_golden/static-cpp/index_impl.h Updates generated hashes.
tools/python/wgsl_template/test/in_tree_golden/static-cpp/string_table.h Adds generated shader strings.
tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/nn/im2col_matmul.h Updates generated template implementation.
tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/oihw_to_ohwi.h Adjusts generated string references.
tools/python/wgsl_template/test/in_tree_golden/static-cpp/generated/tensor/pad.h Adjusts generated string references.
tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/index.h Updates literal-template parameters.
tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/index_impl.h Updates generated hash.
tools/python/wgsl_template/test/in_tree_golden/static-cpp-literal/generated/nn/im2col_matmul.h Updates literal template implementation.
Review details
  • Files reviewed: 9/13 changed files
  • Comments generated: 0
  • Review effort level: Balanced

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

@jchen10

Copy link
Copy Markdown
Contributor

LGTM, thanks!

@qjia7 Jiajia Qin (qjia7) left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review frame

  • Problem/feature validity: Validated. At base 29f53b51176d931283cf55eed0ef68a4db8495d2, IsActivationSupported() rejects ActivationKind::QuickGelu, so an otherwise eligible fused FP16 NHWC Conv cannot select Im2ColMatMulProgram and falls back to the generic Conv path.
  • Risk/scope: Normal. The change widens one Intel Xe-2/Xe-3-only WebGPU dispatch path and adds two generated shader variants, while preserving the existing eligibility, weight-layout, fallback, and cache boundaries.
  • Direction gate: Pass. The smallest owner solution is to admit QuickGelu in the existing im2col selector, carry the existing Activation state into the program, specialize only the alpha-equals-one source variant, and keep arbitrary alpha as a runtime uniform. This PR does that without duplicating Conv dispatch policy or introducing another kernel abstraction.

Confirmed findings

None.

Follow-up opportunities

F1: Consolidate QuickGelu onto the uniform-backed formula

The alpha-equals-one specialization saves one multiply and one uniform read per output, but it also creates a separate shader and pipeline-cache variant and requires each QuickGelu consumer to preserve parallel source-generation logic. Unless measurements show a material benefit from this specialization, prefer always supplying activation_param_0 and using the canonical x * sigmoid(alpha * x) formula.

Please make this a centralized follow-up rather than changing only Im2ColMatMulProgram: make QuickGelu always contribute one uniform in GetActivationUsedUniformCount(), remove the HasUnitQuickGeluAlpha() branch from GetActivationSnippet(), remove QuickGeluUnitAlpha from Activation::CacheKey(), remove the corresponding template parameter and branch here, update the cache-key/snippet tests, and regenerate the WGSL goldens. A local-only change would be incorrect because AppendActivationUniformsData() currently omits the alpha uniform when alpha is exactly 1.

Clarifications

None.

Test coverage

The two tests correctly construct production fusion patterns, require a fused QuickGelu Conv, compare against the unfused graph, and attempt to observe Im2ColMatMul through profiling. They cover the two source-specializing cases: arbitrary alpha uses activation_param_0, while alpha=1 omits both the multiply and uniform. At the reviewed head, however, available WebGPU CI skips both cases before the changed program executes.

WGSL template generation is covered: all 87 local template/parser/golden tests pass, matching the green Linux and Windows WGSL Template CI jobs. Relevant WebGPU build-and-test jobs pass. The lone failing Windows CUDA test job reports two numerical assertion failures in transformers/test_paged_attention_int4.py::TestPagedAttentionInt4::test_xqa_large_attention_scale_and_k_scale; it does not exercise WebGPU or the changed files.

Verdict

Approve. The performance regression is valid, the implementation direction is appropriate, and the shader specialization, uniform layout, cache identity, and generated artifacts are internally consistent. I found no confirmed code blocker, compatibility issue, or outstanding clarification. Available CI cannot exercise the Intel Xe-2/Xe-3-only route, but the tests correctly detect and report that limitation rather than producing false coverage. F1 is non-blocking follow-up work: consolidate QuickGelu centrally onto the uniform-backed formula unless measurements justify retaining the alpha-1 shader specialization. There is no cleanup-only feedback.

@qjia7

Copy link
Copy Markdown
Contributor

xhcao, please resolve the conflicts. Thanks.

@xhcao

xhcao commented Sep 14, 2026

Copy link
Copy Markdown
Contributor Author

xhcao, please resolve the conflicts. Thanks.

Jiajia Qin (@qjia7) Done, thanks

@qjia7
Jiajia Qin (qjia7) enabled auto-merge (squash) September 14, 2026 03:41
@qjia7
Jiajia Qin (qjia7) merged commit 8d85527 into microsoft:main Sep 14, 2026
90 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants