Enable SDPA without kv cache

Summary: Sdpa custom op has been decoupled from kv cache by kimishpatel. Update the llama definition so that the sdpa op is applied both with and without kv cache.

Differential Revision: D70593177

Author: Martin Yuan

examples/models/llama/attention.py

@@ -209,13 +209,13 @@ def __init__(self, args: ModelArgs, layer_id: int, rope: Rope):
                 self.head_dim,
                 args.enable_dynamic_shape,
             )
-            self.SDPA = SDPA(
-                dim=self.n_local_heads * self.head_dim,
-                head_dim=self.head_dim,
-                n_rep=self.n_rep,
-                max_context_len=self.max_context_len,
-                enable_dynamic_shape=args.enable_dynamic_shape,
-            )
+        self.SDPA = SDPA(
+            dim=self.n_local_heads * self.head_dim,
+            head_dim=self.head_dim,
+            n_rep=self.n_rep,
+            max_context_len=self.max_context_len,
+            enable_dynamic_shape=args.enable_dynamic_shape,
+        )
 
     def forward(
         self,
@@ -244,21 +244,5 @@ def forward(
         if self.use_kv_cache:
             assert input_pos is not None
             k, v = self.kv_cache.update(input_pos, k, v)
-            output = self.SDPA(input_pos, q, k, v, bsz, seqlen, self.mask)
-            return self.wo(output), None
-
-        # grouped multiquery attention: expand out keys and values
-        k = k.repeat_interleave(self.n_rep, dim=1)
-        v = v.repeat_interleave(self.n_rep, dim=1)
-
-        assert hasattr(self, "mask")
-
-        mask = self.mask[:seqlen, :seqlen]
-
-        output = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0)
-
-        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
-
-        output = self.wo(output)
-
-        return output, None
+        output = self.SDPA(input_pos, q, k, v, bsz, seqlen, self.mask)
+        return self.wo(output), None