expand saved attention ops to incl cudnn attention (#274)

xmfan · web-flow · commit 29281606432b · 2025-12-08T17:01:27.000-08:00
stack-info: PR: #274, branch: xmfan/stack/25
diff --git a/autoparallel/activation_checkpointing.py b/autoparallel/activation_checkpointing.py
@@ -454,5 +454,6 @@ def ac_joint_pass(
         torch.ops.aten.mm.default,
         torch.ops.aten._scaled_dot_product_efficient_attention.default,
         torch.ops.aten._scaled_dot_product_flash_attention.default,
+        torch.ops.aten._scaled_dot_product_cudnn_attention.default,
     }
     _apply_ac_policy(graph, save_list=save_list)

Original file line number	Diff line number	Diff line change
`@@ -454,5 +454,6 @@ def ac_joint_pass(`
`454`	`454`	`torch.ops.aten.mm.default,`
`455`	`455`	`torch.ops.aten._scaled_dot_product_efficient_attention.default,`
`456`	`456`	`torch.ops.aten._scaled_dot_product_flash_attention.default,`
	`457`	`+ torch.ops.aten._scaled_dot_product_cudnn_attention.default,`
`457`	`458`	`}`
`458`	`459`	`_apply_ac_policy(graph, save_list=save_list)`