Fix an issue when piping attn_logits_soft_cap through in vllm. (#8600)

fenghuizhang · web-flow · commit 5b877beb6eb2 · 2025-01-22T13:06:58.000-08:00
diff --git a/torch_xla/experimental/custom_kernel.py b/torch_xla/experimental/custom_kernel.py
@@ -1080,7 +1080,8 @@ def flash_attention_non_xla(q: torch.Tensor,
 
 
 XLA_LIB.define(
-    "paged_attention(Tensor q, Tensor k_pages, Tensor v_pages, Tensor lengths, Tensor page_indices, int pages_per_compute_block, str megacore_mode=None, float attn_logits_soft_cap=None) -> Tensor",
+    "paged_attention(Tensor q, Tensor k_pages, Tensor v_pages, Tensor lengths, Tensor page_indices,"
+    " int pages_per_compute_block, str megacore_mode=None, float? attn_logits_soft_cap=None) -> Tensor",
 )
 
 

Original file line number	Diff line number	Diff line change
`@@ -1080,7 +1080,8 @@ def flash_attention_non_xla(q: torch.Tensor,`
`1080`	`1080`
`1081`	`1081`
`1082`	`1082`	`XLA_LIB.define(`
`1083`		`- "paged_attention(Tensor q, Tensor k_pages, Tensor v_pages, Tensor lengths, Tensor page_indices, int pages_per_compute_block, str megacore_mode=None, float attn_logits_soft_cap=None) -> Tensor",`
	`1083`	`+ "paged_attention(Tensor q, Tensor k_pages, Tensor v_pages, Tensor lengths, Tensor page_indices,"`
	`1084`	`+ " int pages_per_compute_block, str megacore_mode=None, float? attn_logits_soft_cap=None) -> Tensor",`
`1084`	`1085`	`)`
`1085`	`1086`
`1086`	`1087`