Enable flash attention by default (#690)

Author: rsuderman

sharktank/sharktank/layers/configs/llm_configs.py

@@ -167,7 +167,7 @@ class LlamaModelConfig:
     tensor_parallelism_size: int = 1
 
     # Which attention kernel to use.
-    attention_kernel: str = "decomposed"
+    attention_kernel: str = "torch"
 
     # Indicates if running with HuggingFace implementation and ensures
     # numerical equivalency to HuggingFace's LLaMa if true (by modifying

sharktank/sharktank/layers/paged_llama_attention_block.py

@@ -216,14 +216,12 @@ def repeat_kv(x: torch.Tensor) -> torch.Tensor:
                 attn_weights, values
             )  # (bs, heads, slen, head_dim)
         else:
-            is_causal = True
-            attention_mask = None
             attn_output = ops.scaled_dot_product_attention(
                 q=xq,  # [bs, ..., sl, dim]
                 k=keys,  # [bs, ..., sl, dim]
                 v=values,  # [bs, ..., sl, dim]
                 a=attention_mask,  # [bs, ..., sl, sl]
-                is_causal=is_causal,  # assumes causal masking when true
+                is_causal=False,  # assumes causal masking when true
                 scale=None,  # defaults to 1/sqrt(dim)
             )
 

sharktank/sharktank/utils/cli.py

@@ -66,7 +66,7 @@ def add_model_options(parser: argparse.ArgumentParser):
     parser.add_argument(
         "--attention-kernel",
         type=str,
-        default="decomposed",
+        default="torch",
         choices=["decomposed", "torch"],
     )
     parser.add_argument(