[WebGPU] Support QuickGelu activation in the im2col-matmul path - #32560
Conversation
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
There was a problem hiding this comment.
🟢 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.
|
LGTM, thanks! |
Jiajia Qin (qjia7)
left a comment
There was a problem hiding this comment.
Review frame
- Problem/feature validity: Validated. At base
29f53b51176d931283cf55eed0ef68a4db8495d2,IsActivationSupported()rejectsActivationKind::QuickGelu, so an otherwise eligible fused FP16 NHWC Conv cannot selectIm2ColMatMulProgramand 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
Activationstate 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.
|
xhcao, please resolve the conflicts. Thanks. |
895e205 to
f43e275
Compare
Jiajia Qin (@qjia7) Done, thanks |
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 thesemodels 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.