force is_causal for ndc

Author: dan-garvey

sharktank/sharktank/layers/paged_llama_attention_block.py

@@ -170,7 +170,9 @@ def repeat_kv(x: torch.Tensor) -> torch.Tensor:
                 attn_weights, values
             )  # (bs, heads, slen, head_dim)
         else:
-            is_causal = attention_mask is None and batch_seq_len == 1
+            # Use the builtin attention mask when not decomposed
+            is_causal = True
+            attention_mask = None
             attn_output = torch.nn.functional.scaled_dot_product_attention(
                 query=xq,  # [bs, ..., sl, dim]
                 key=keys,  # [bs, ..., sl, dim]