Implements attn_logits_soft_cap and pass it through multi_queries_pag…

…ed_attention
pytorch · Jan 19, 2025 · 491dbdb · 491dbdb
1 parent 68cd431
commit 491dbdb
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torch_xla/experimental/pallas_kernels/multi_queries_paged_attention_kernel.py b/torch_xla/experimental/pallas_kernels/multi_queries_paged_attention_kernel.py
@@ -271,7 +271,7 @@ def paged_flash_attention_kernel(
     num_kv_pages_per_compute_block: int,
     mask_value: float,
     query_len: int,
-    attn_logits_soft_cap: float | None = None,
+    attn_logits_soft_cap: float | None,
 ):
   """Pallas kernel for paged attention."""
   b, kv_head_idx, q_blk_idx, kv_blk_idx = (