Update cudnn-frontend to 1.0.3 to fix cuDNN v9 SDPA NaNs (#650)

* Update cudnn frontend to 1.0.3 to fix cudnn v9 Nans Signed-off-by: Charlene Yang <[email protected]> * make d_out contiguous for bwd Signed-off-by: Charlene Yang <[email protected]> * remove cudnnDestroy to let torch handle it Signed-off-by: Charlene Yang <[email protected]> * Update transformer_engine/pytorch/attention.py Co-authored-by: Tim Moon <[email protected]> Signed-off-by: cyanguwa <[email protected]> * Update transformer_engine/pytorch/attention.py Co-authored-by: Tim Moon <[email protected]> Signed-off-by: cyanguwa <[email protected]> * Update transformer_engine/pytorch/attention.py Co-authored-by: Tim Moon <[email protected]> Signed-off-by: cyanguwa <[email protected]> --------- Signed-off-by: Charlene Yang <[email protected]> Signed-off-by: cyanguwa <[email protected]> Co-authored-by: Tim Moon <[email protected]>
NVIDIA · Feb 3, 2024 · 5b90b7f · 5b90b7f
1 parent df9c29e
commit 5b90b7f
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 6 deletions.
diff --git a/3rdparty/cudnn-frontend b/3rdparty/cudnn-frontend
diff --git a/transformer_engine/common/fused_attn/utils.h b/transformer_engine/common/fused_attn/utils.h
@@ -152,11 +152,6 @@ class cudnnExecutionPlanManager {
     }
 
     ~cudnnExecutionPlanManager() {
-        static thread_local std::once_flag flag;
-        std::call_once(flag, [&] {
-                        if (handle_ != nullptr) {
-                          cudnnDestroy(handle_);
-                        }});
     }
 
  private:

diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py
@@ -1733,6 +1733,7 @@ def forward(ctx, is_training, max_seqlen, cu_seqlens, qkv, qkv_dtype, attn_bias,
 
     @staticmethod
     def backward(ctx, d_out):
+        d_out = d_out.contiguous()
         qkv, out, cu_seqlens = ctx.saved_tensors
         if not ctx.aux_ctx_tensors[0].is_contiguous():
             ctx.aux_ctx_tensors[0] = ctx.aux_ctx_tensors[0].contiguous()
@@ -1802,6 +1803,7 @@ def forward(ctx, is_training, max_seqlen_q, max_seqlen_kv, cu_seqlens_q, cu_seql
 
     @staticmethod
     def backward(ctx, d_out):
+        d_out = d_out.contiguous()
         q, kv, out, cu_seqlens_q, cu_seqlens_kv = ctx.saved_tensors
         if not ctx.aux_ctx_tensors[0].is_contiguous():
             ctx.aux_ctx_tensors[0] = ctx.aux_ctx_tensors[0].contiguous()
@@ -1883,6 +1885,7 @@ def forward(ctx, is_training, max_seqlen_q, max_seqlen_kv, cu_seqlens_q, cu_seql
 
     @staticmethod
     def backward(ctx, d_out):
+        d_out = d_out.contiguous()
         q, k, v, out, cu_seqlens_q, cu_seqlens_kv = ctx.saved_tensors
         if not ctx.aux_ctx_tensors[0].is_contiguous():
             ctx.aux_ctx_tensors[0] = ctx.aux_ctx_tensors[0].contiguous()