Skip to content

Add VAE graph optimization surgeries - #1260

Draft
xieofxie wants to merge 7 commits into
mainfrom
hualxie/more_optims_vae_draft
Draft

Add VAE graph optimization surgeries#1260
xieofxie wants to merge 7 commits into
mainfrom
hualxie/more_optims_vae_draft

Conversation

@xieofxie

Copy link
Copy Markdown
Contributor

Summary

  • add four opt-in graph surgeries for redundant L2 normalization nodes, generated nearest-neighbor GatherND paths, exact SiLU topology, and scaled MatMul
  • run proof-dependent rewrites before ORT graph processing while emitting QuickGelu afterward so contrib operators survive the full pipeline
  • use conservative topology, shape, type, constant, and consumer checks, including BF16 scalar decoding and element-type matching
  • add generated-model regression coverage and document the Wan 2.1 VAE analysis, mathematical equivalence, CUDA constraints, examples, and end-to-end operator counts in vae_opt.md

Behavior

All new capabilities are disabled by default. QuickGelu and FusedMatMul are constrained to CUDA. The scaled MatMul rewrite is algebraically equivalent but may change floating-point rounding when scaling is incorporated into GEMM.

from typing import Any, ClassVar

import numpy as np
import onnx
dimensions.append(dimension.dim_param)
else:
dimensions.append(None)
shapes[value.name] = tuple(dimensions)
from typing import TYPE_CHECKING

import numpy as np
import onnx
hualxie added 3 commits July 31, 2026 14:40
…_to_optimize

# Conflicts:
#	src/winml/modelkit/commands/optimize.py
#	src/winml/modelkit/optim/analysis.py
#	tests/unit/commands/test_optimize_cli.py
#	tests/unit/optim/test_analysis.py
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.

2 participants