Address comments

apple · Jan 13, 2025 · 24071f5 · 24071f5
1 parent 9850e2e
commit 24071f5
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/axlearn/common/flash_attention/gpu_attention.py b/axlearn/common/flash_attention/gpu_attention.py
@@ -83,6 +83,9 @@ def _mha_forward_kernel(
 
     See also `_mha_backward_kernel` for the backward pass.
 
+    Note: the kernel name is used to do string matching for rematerialization in `remat.py`. Be
+    careful when renaming this.
+
     Args:
         q_ref: Input query ref.
         k_ref: Input key ref.