[cuDNN] cuDNN frontend for LayerNorm RMSNorm #159682
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/159682
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ❌ 2 New Failures, 2 Unrelated FailuresAs of commit 6030457 with merge base aeb5321 ( NEW FAILURES - The following jobs have failed:
UNSTABLE - The following jobs are marked as unstable, possibly due to flakiness on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
|
FYI: @eqy had a stale PR doing it |
|
haha yes this is based on that |
| }; | ||
|
|
||
| void setLayerNormParams(LayerNormParams& params, const Tensor& X, int64_t M, int64_t N) { | ||
| memset(¶ms, 0, sizeof(params)); |
| {Y_fe, Y->data_ptr()}}; | ||
| variant_pack = std::move(variant_pack_); | ||
| auto result = std::make_tuple(layernorm_graph, X_fe, mean_fe, inv_variance_fe, scale_fe, bias_fe, Y_fe); | ||
| layernorm_forward_graph_cache.update(key, result); |
There was a problem hiding this comment.
Can result not be moved here too?
| {DX_fe, dX->data_ptr()}}; | ||
| variant_pack = std::move(variant_pack_); | ||
| auto result = std::make_tuple(layernorm_graph, X_fe, DY_fe, mean_fe, inv_variance_fe, scale_fe, dscale_fe, dbias_fe, DX_fe); | ||
| layernorm_backward_graph_cache.update(key, result); |
| }; | ||
|
|
||
| void setRMSNormParams(RMSNormParams& params, const Tensor& X, int64_t M, int64_t N) { | ||
| memset(¶ms, 0, sizeof(params)); |
| return &(it->second); | ||
| } | ||
|
|
||
| void update(const KeyType& key, T& results) { |
There was a problem hiding this comment.
Should use perfect forwarding, no?
| {DX_fe, dX->data_ptr()}}; | ||
| variant_pack = std::move(variant_pack_); | ||
| auto result = std::make_tuple(rmsnorm_graph, X_fe, DY_fe, inv_variance_fe, scale_fe, dscale_fe, DX_fe); | ||
| rmsnorm_backward_graph_cache.update(key, result); |
|
@pytorchbot rebase |
|
@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here |
|
Successfully rebased |
d4144f0 to
6030457
Compare
cuDNN performance as of 9.10 is still not great
if it ever becomes good, this PR could streamline process of adding experimental cudnn backend for layernorm and rmsnorm
Follows examples at: https://github.com/NVIDIA/cudnn-frontend/tree/main/samples/cpp/norm