add todo ring attention

2024-10-29 14:08:53 +00:00 · 2024-10-29 14:08:53 +00:00 · 987a7c5c99
commit 987a7c5c99
parent 46af5b0425
1 changed files with 2 additions and 1 deletions
--- a/src/parallel/context_parallel.py
+++ b/src/parallel/context_parallel.py
@ -15,7 +15,8 @@ class RingAttentionFunc(torch.autograd.Function):
    @staticmethod
    def forward(ctx, q, k, v, sm_scale, is_causal):
        comm = ContextComms("comm")
-        #NOTE: Find a better to save these tensors without cloning
+        #TODO(fmom): add flash attention
+        #TODO(fmom): Find a better to save these tensors without cloning
        k_og = k.clone()
        v_og = v.clone()
        out, lse = None, None