Fix 3.9 Python syntax (#9018)

Author: zpcore

test/test_multi_queries_paged_attention_kernel.py

@@ -5,6 +5,7 @@
 from torch_xla.experimental.pallas_kernels.multi_queries_paged_attention_kernel import paged_attention
 import jax.numpy as jnp
 import numpy as np
+from typing import Optional
 
 jax.config.parse_flags_with_absl()
 
@@ -45,7 +46,7 @@ def _ref_jax_extended_paged_attention(
     lengths,  # [batch_size], the effective kv_length.
     page_indices,  # [batch_size, pages_per_sequence]
     effective_q_lens,  # [batch_size] the effective q_length
-    attn_logits_soft_cap: float | None = None,
+    attn_logits_soft_cap: Optional[float] = None,
 ):
   batch_size, query_len, num_query_heads, head_size = q.shape
   num_kv_heads, total_num_pages, page_size, _ = k_pages.shape

torch_xla/experimental/custom_kernel.py

@@ -1086,7 +1086,7 @@ def _multi_queries_paged_attention_nonkernel(
     lengths,  # seq_lengths, [batch_size]. nb batch_size = len(seq_lens), the effective kv_length.
     page_indices,  # [batch_size, pages_per_sequence]
     effective_q_lens,  # [batch_size], the effective q_length
-    attn_logits_soft_cap: float | None = None,
+    attn_logits_soft_cap: Optional[float] = None,
 ) -> torch.Tensor:  # [batch_size, query_len, num_heads, head_dim]
   batch_size, query_len, num_query_heads, head_size = q.shape
   num_kv_heads, total_num_pages, page_size, _ = k_pages.shape
@@ -1154,7 +1154,7 @@ def multi_queries_paged_attention(
     num_kv_pages_per_compute_block,
     num_queries_per_compute_block,
     use_kernel=True,
-    attn_logits_soft_cap: float | None = None,
+    attn_logits_soft_cap: Optional[float] = None,
 ):  # [batch_size, query_len, num_heads, head_dim]:
   assert len(q.shape) == 4, "q should have 4 dimensions."
   if not use_kernel:
@@ -1672,8 +1672,8 @@ def paged_attention_xla(q: torch.Tensor,
                         lengths: torch.Tensor,
                         page_indices: torch.Tensor,
                         pages_per_compute_block: int,
-                        megacore_mode: str | None = None,
-                        attn_logits_soft_cap: float | None = None):
+                        megacore_mode: Optional[str] = None,
+                        attn_logits_soft_cap: Optional[float] = None):
   return paged_attention(q, k_pages, v_pages, lengths, page_indices,
                          pages_per_compute_block, megacore_mode,
                          attn_logits_soft_cap)
@@ -1686,8 +1686,8 @@ def paged_attention_non_xla(q: torch.Tensor,
                             lengths: torch.Tensor,
                             page_indices: torch.Tensor,
                             pages_per_compute_block: int,
-                            megacore_mode: str | None = None,
-                            attn_logits_soft_cap: float | None = None):
+                            megacore_mode: Optional[str] = None,
+                            attn_logits_soft_cap: Optional[float] = None):
   return non_xla_attetion(q, k_pages, v_pages, "paged")
 
 
@@ -1698,17 +1698,17 @@ def paged_attention_non_xla(q: torch.Tensor,
 
 
 @impl(XLA_LIB, "multi_queries_paged_attention", "XLA")
-def multi_queries_paged_attention_xla(q: torch.Tensor,
-                                      k_pages: torch.Tensor,
-                                      v_pages: torch.Tensor,
-                                      lengths: torch.Tensor,
-                                      page_indices: torch.Tensor,
-                                      effective_q_lens: torch.Tensor,
-                                      num_kv_pages_per_compute_block: int,
-                                      num_queries_per_compute_block: int,
-                                      use_kernel: bool,
-                                      attn_logits_soft_cap: float |
-                                      None = None):
+def multi_queries_paged_attention_xla(
+    q: torch.Tensor,
+    k_pages: torch.Tensor,
+    v_pages: torch.Tensor,
+    lengths: torch.Tensor,
+    page_indices: torch.Tensor,
+    effective_q_lens: torch.Tensor,
+    num_kv_pages_per_compute_block: int,
+    num_queries_per_compute_block: int,
+    use_kernel: bool,
+    attn_logits_soft_cap: Optional[float] = None):
   return multi_queries_paged_attention(q, k_pages, v_pages, lengths,
                                        page_indices, effective_q_lens,
                                        num_kv_pages_per_compute_block,
@@ -1717,17 +1717,17 @@ def multi_queries_paged_attention_xla(q: torch.Tensor,
 
 
 @impl(XLA_LIB, "multi_queries_paged_attention", "CompositeExplicitAutograd")
-def multi_queries_paged_attention_non_xla(q: torch.Tensor,
-                                          k_pages: torch.Tensor,
-                                          v_pages: torch.Tensor,
-                                          lengths: torch.Tensor,
-                                          page_indices: torch.Tensor,
-                                          effective_q_lens: torch.Tensor,
-                                          num_kv_pages_per_compute_block: int,
-                                          num_queries_per_compute_block: int,
-                                          use_kernel: bool,
-                                          attn_logits_soft_cap: float |
-                                          None = None):
+def multi_queries_paged_attention_non_xla(
+    q: torch.Tensor,
+    k_pages: torch.Tensor,
+    v_pages: torch.Tensor,
+    lengths: torch.Tensor,
+    page_indices: torch.Tensor,
+    effective_q_lens: torch.Tensor,
+    num_kv_pages_per_compute_block: int,
+    num_queries_per_compute_block: int,
+    use_kernel: bool,
+    attn_logits_soft_cap: Optional[float] = None):
   return non_xla_attetion(q, k_pages, v_pages, "paged")
 
 
@@ -1751,8 +1751,8 @@ def ragged_paged_attention_xla(
     num_queries_per_block: int,
     use_kernel: bool,
     sm_scale: float = 1.0,
-    mask_value: float | None = None,
-    vmem_limit_bytes: int | None = None,
+    mask_value: Optional[float] = None,
+    vmem_limit_bytes: Optional[int] = None,
 ):
   return ragged_paged_attention(
       q,
@@ -1782,8 +1782,8 @@ def ragged_paged_attention_non_xla(q: torch.Tensor,
                                    num_queries_per_block: int,
                                    use_kernel: bool,
                                    sm_scale: float = 1.0,
-                                   mask_value: float | None = None,
-                                   vmem_limit_bytes: int | None = None):
+                                   mask_value: Optional[float] = None,
+                                   vmem_limit_bytes: Optional[int] = None):
   return non_xla_attetion(q, k_pages, v_pages, "paged")
 
 

torch_xla/experimental/pallas_kernels/multi_queries_paged_attention_kernel.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Sequence
 import functools
-from typing import Literal, cast
+from typing import Literal, cast, Union, Optional
 
 import jax
 from jax import lax
@@ -116,7 +116,7 @@ def _flash_attention(
     query_len: int,
     page_size: int,
     head_dim: int,
-    attn_logits_soft_cap: float | None,
+    attn_logits_soft_cap: Optional[float],
 ):
   b, kv_head_idx, q_blk_idx, kv_blk_idx = (
       pl.program_id(0),
@@ -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,
+    attn_logits_soft_cap: Optional[float],
 ):
   """Pallas kernel for paged attention."""
   b, kv_head_idx, q_blk_idx, kv_blk_idx = (
@@ -440,16 +440,16 @@ def prefetch_next_block():  # pylint: disable=unused-variable
 )
 def paged_attention(
     q: jax.Array,
-    k_pages: jax.Array | quantization_utils.QuantizedTensor,
-    v_pages: jax.Array | quantization_utils.QuantizedTensor,
+    k_pages: Union[jax.Array, quantization_utils.QuantizedTensor],
+    v_pages: Union[jax.Array, quantization_utils.QuantizedTensor],
     lengths: jax.Array,
     page_indices: jax.Array,
     effective_q_lens: jax.Array,
     *,
     mask_value: float = DEFAULT_MASK_VALUE,
     num_kv_pages_per_compute_block: int,
     num_queries_per_compute_block: int = 4,
-    attn_logits_soft_cap: float | None = None,
+    attn_logits_soft_cap: Optional[float] = None,
 ) -> jax.Array:
   """Paged grouped query attention.
 

torch_xla/experimental/pallas_kernels/ragged_paged_attention_v2.py

@@ -7,6 +7,8 @@
 """
 
 import functools
+from typing import Optional
+
 import jax
 from jax import lax
 from jax.experimental import pallas as pl
@@ -571,7 +573,7 @@ def ragged_paged_attention(
     mask_value: float = DEFAULT_MASK_VALUE,
     num_kv_pages_per_block: int = 16,
     num_queries_per_block: int = 128,
-    vmem_limit_bytes: int | None = None,
+    vmem_limit_bytes: Optional[int] = None,
 ):
   """Ragged paged attention that supports mixed prefill and decode.
 

torchax/examples/train_llama_torchtitan/splash_attn.py

@@ -1,4 +1,5 @@
 import functools
+from typing import Optional
 
 import jax
 import jax.numpy as jnp
@@ -16,8 +17,8 @@ def tpu_splash_attention(
     query: jax.Array,
     key: jax.Array,
     value: jax.Array,
-    decoder_segment_ids: jax.Array | None,
-    attn_logits_soft_cap: float | None = None,
+    decoder_segment_ids: Optional[jax.Array],
+    attn_logits_soft_cap: Optional[float] = None,
 ) -> jax.Array:
   """TPU Flash Attention."""
   if decoder_segment_ids is not None: