re-apply

Author: oraluben

examples/attention_sink/example_gqa_sink_bwd_bhsd.py

@@ -6,7 +6,6 @@
 from tilelang.profiler import do_bench
 import tilelang.language as T
 import argparse
-from typing import Optional
 
 
 def get_bwd_configs():
@@ -405,7 +404,7 @@ def ref_program(query: torch.Tensor,
                 key: torch.Tensor,
                 value: torch.Tensor,
                 sinks: torch.Tensor,
-                sliding_window: Optional[int] = None,
+                sliding_window: int | None = None,
                 dtype: torch.dtype = torch.float16) -> torch.Tensor:
 
     key = key.transpose(1, 2).contiguous()

examples/attention_sink/example_gqa_sink_fwd_bhsd_wgmma_pipelined.py

@@ -13,7 +13,6 @@
 import triton
 import triton.language as tl
 from triton.tools.tensor_descriptor import TensorDescriptor
-from typing import Optional
 
 
 def get_configs():
@@ -213,7 +212,7 @@ def ref_program(query: torch.Tensor,
                 key: torch.Tensor,
                 value: torch.Tensor,
                 sinks: torch.Tensor,
-                sliding_window: Optional[int] = None,
+                sliding_window: int | None = None,
                 dtype: torch.dtype = torch.float16) -> torch.Tensor:
 
     key = key.transpose(1, 2).contiguous()

examples/attention_sink/example_mha_sink_bwd_bhsd.py

@@ -6,7 +6,6 @@
 from tilelang.profiler import do_bench
 import tilelang.language as T
 import argparse
-from typing import Optional
 
 
 def get_bwd_configs():
@@ -401,7 +400,7 @@ def ref_program(query: torch.Tensor,
                 key: torch.Tensor,
                 value: torch.Tensor,
                 sinks: torch.Tensor,
-                sliding_window: Optional[int] = None,
+                sliding_window: int | None = None,
                 dtype: torch.dtype = torch.float16) -> torch.Tensor:
 
     query = query.transpose(1, 2).contiguous().unsqueeze(

examples/attention_sink/example_mha_sink_fwd_bhsd.py

@@ -9,7 +9,6 @@
 from tilelang.layout import make_swizzled_layout
 import itertools
 import argparse
-from typing import Optional
 
 
 def get_configs():
@@ -193,7 +192,7 @@ def ref_program(query: torch.Tensor,
                 key: torch.Tensor,
                 value: torch.Tensor,
                 sinks: torch.Tensor,
-                sliding_window: Optional[int] = None,
+                sliding_window: int | None = None,
                 dtype: torch.dtype = torch.float16) -> torch.Tensor:
 
     query = query.transpose(1, 2).contiguous().unsqueeze(
@@ -306,8 +305,8 @@ def main(batch: int = 1,
 
         latency = do_bench(
             lambda: ref_program(Q, K, V, sinks, window_size, dtype=torch_dtype), warmup=500)
-        print("Ref: {:.2f} ms".format(latency))
-        print("Ref: {:.2f} TFlops".format(total_flops / latency * 1e-9))
+        print(f"Ref: {latency:.2f} ms")
+        print(f"Ref: {total_flops / latency * 1e-9:.2f} TFlops")
         latency = do_bench(lambda: kernel(Q, K, V, sinks), warmup=500)
         print(f"Tilelang: {latency:.2f} ms")
         print(f"Tilelang: {total_flops / latency * 1e-9:.2f} TFlops")

examples/attention_sink/example_mha_sink_fwd_bhsd_wgmma_pipelined.py

@@ -13,7 +13,6 @@
 import triton
 import triton.language as tl
 from triton.tools.tensor_descriptor import TensorDescriptor
-from typing import Optional
 
 
 def get_configs():
@@ -206,7 +205,7 @@ def ref_program(query: torch.Tensor,
                 key: torch.Tensor,
                 value: torch.Tensor,
                 sinks: torch.Tensor,
-                sliding_window: Optional[int] = None,
+                sliding_window: int | None = None,
                 dtype: torch.dtype = torch.float16) -> torch.Tensor:
 
     query = query.transpose(1, 2).contiguous().unsqueeze(