Thanks to visit codestin.com
Credit goes to github.com

Skip to content

torch.compile Pass For Replacing Layer Norm With CuNN Version#1892

Merged
crcrpar merged 1 commit into
NVIDIA:masterfrom
alpha0422:alpha0422/torchsched-layernorm
Apr 4, 2025
Merged

torch.compile Pass For Replacing Layer Norm With CuNN Version#1892
crcrpar merged 1 commit into
NVIDIA:masterfrom
alpha0422:alpha0422/torchsched-layernorm

Conversation

@alpha0422
Copy link
Copy Markdown
Contributor

@alpha0422 alpha0422 commented Apr 2, 2025

This PR implements a customized torch.compile pass called pre_grad_custom_pass, and registers a replacement for PyTorch layer norm with its CuDNN version for better performance automatically.

@alpha0422 alpha0422 force-pushed the alpha0422/torchsched-layernorm branch from 1940ae2 to 7580bd9 Compare April 2, 2025 09:06
Comment thread apex/contrib/torchsched/passes/pre_grad_passes.py
Comment thread apex/contrib/torchsched/ops/layer_norm.py
@alpha0422 alpha0422 force-pushed the alpha0422/torchsched-layernorm branch from 7580bd9 to 87b6ada Compare April 3, 2025 07:53
@alpha0422 alpha0422 force-pushed the alpha0422/torchsched-layernorm branch from 87b6ada to 74ac4e5 Compare April 3, 2025 14:10
@crcrpar crcrpar added the contrib label Apr 4, 2025
@crcrpar crcrpar merged commit 379937e into NVIDIA:master Apr 4, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants