[pallas-tpu] Switch float16 to HALF_DTYPE in examples for TPU compatibility (#1648)

Author: v0i0

examples/attention.py

@@ -166,7 +166,7 @@ def ref_attention(
 def main() -> None:
     """
     Main entry point that runs the attention kernel test with specific parameters.
-    Tests with batch size 2, 32 heads, 1024 sequence length, and 64-dimensional heads using float16.
+    Tests with batch size 2, 32 heads, 1024 sequence length, and 64-dimensional heads using half-precision.
     """
     test(2, 32, 1024, 64, HALF_DTYPE, device=DEVICE)
 

examples/flex_attention.py

@@ -272,7 +272,7 @@ def flex_fn(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
 def main() -> None:
     """
     Main entry point that runs the attention kernel test with specific parameters.
-    Tests with batch size 2, 32 heads, 1024 sequence length, and 64-dimensional heads using float16.
+    Tests with batch size 2, 32 heads, 1024 sequence length, and 64-dimensional heads using half-precision.
     """
     test(2, 32, 1024, 64, HALF_DTYPE)
     test(2, 4, 1024, 64, HALF_DTYPE, lambda score, *_: torch.tanh(score))

examples/fp8_gemm.py

@@ -3,7 +3,7 @@
 ====================================================
 This example demonstrates an FP8 GEMM kernel implemented in Helion. The kernel performs
 matrix multiplication on FP8 inputs, accumulating results in FP32 for accuracy, and
-outputs in FP16 format. It includes a reference PyTorch implementation using
+outputs in half-precision format. It includes a reference PyTorch implementation using
 torch._scaled_mm for correctness comparison, and a test function to validate the kernel.
 """
 
@@ -17,6 +17,7 @@
 
 import helion
 from helion._testing import DEVICE
+from helion._testing import HALF_DTYPE
 from helion._testing import run_example
 import helion.language as hl
 
@@ -39,13 +40,13 @@ def fp8_gemm(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         x (torch.Tensor): Input tensor of shape [m, k] in FP8 format.
         y (torch.Tensor): Input tensor of shape [k, n] in FP8 format.
     Returns:
-        torch.Tensor: Output tensor of shape [m, n] in FP16 format.
+        torch.Tensor: Output tensor of shape [m, n] in half-precision format.
     """
     m, k = x.size()
     k2, n = y.size()
     assert k == k2, f"size mismatch {k} != {k2}"
-    # Output is in FP16 to match tritonbench behavior
-    out = torch.empty([m, n], dtype=torch.float16, device=x.device)
+    # Output is in half-precision to match tritonbench behavior
+    out = torch.empty([m, n], dtype=HALF_DTYPE, device=x.device)
     for tile_m, tile_n in hl.tile([m, n]):
         # Accumulate in FP32 for accuracy
         acc = hl.zeros([tile_m, tile_n], dtype=torch.float32)
@@ -55,7 +56,7 @@ def fp8_gemm(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
             y_tile = y[tile_k, tile_n]
             # Use hl.dot for FP8 GEMM
             acc = hl.dot(x_tile, y_tile, acc=acc)
-        out[tile_m, tile_n] = acc.to(torch.float16)
+        out[tile_m, tile_n] = acc.to(HALF_DTYPE)
     return out
 
 
@@ -69,14 +70,14 @@ def reference_fp8_gemm_pytorch(
         x_fp8 (torch.Tensor): Input tensor in FP8 format.
         y_fp8 (torch.Tensor): Input tensor in FP8 format.
     Returns:
-        torch.Tensor: Output tensor in FP16 format.
+        torch.Tensor: Output tensor in half-precision format.
     """
     # torch._scaled_mm requires column-major for second operand
     y_fp8_t = y_fp8.T.contiguous().T
     scale_a = torch.tensor(1.0, device=x_fp8.device)
     scale_b = torch.tensor(1.0, device=x_fp8.device)
     return torch._scaled_mm(
-        x_fp8, y_fp8_t, scale_a, scale_b, use_fast_accum=False, out_dtype=torch.float16
+        x_fp8, y_fp8_t, scale_a, scale_b, use_fast_accum=False, out_dtype=HALF_DTYPE
     )
 
 
@@ -97,7 +98,7 @@ def fp8_gemm_tritonbench(
         scale_a (torch.Tensor): Scale factor for tensor a (unused in our implementation).
         scale_b (torch.Tensor): Scale factor for tensor b (unused in our implementation).
     Returns:
-        Callable that returns output tensor in FP16 format.
+        Callable that returns output tensor in half-precision format.
     """
     return lambda: fp8_gemm(a, b)
 

examples/layer_norm.py

@@ -2,7 +2,7 @@
 Helion Layer Normalization Forward and Backward Example
 =======================================================
 This example demonstrates a Helion kernel implementation of 1D layer normalization
-with both forward and backward passes using FP16 inputs and compares it against
+with both forward and backward passes using half-precision inputs and compares it against
 PyTorch's built-in layer_norm function.
 """
 
@@ -33,14 +33,14 @@ def layer_norm_fwd(
     """
     Performs 1D layer normalization on the input tensor using Helion.
     Args:
-        x (torch.Tensor): Input tensor of shape [batch_size, dim], expected to be FP16.
+        x (torch.Tensor): Input tensor of shape [batch_size, dim], expected to be half-precision.
         normalized_shape (list[int]): List containing the dimension to normalize over (should be length 1).
         weight (torch.Tensor): Learnable scale parameter of shape [dim].
         bias (torch.Tensor | None): Optional learnable bias parameter of shape [dim].
         eps (float, optional): Small value added to variance for numerical stability. Default is 1e-5.
     Returns:
         tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-            - The layer-normalized output tensor of shape [batch_size, dim], in FP16.
+            - The layer-normalized output tensor of shape [batch_size, dim], in half-precision.
             - Mean tensor of shape [batch_size], in FP32.
             - Reciprocal standard deviation tensor of shape [batch_size], in FP32.
     """

examples/mamba2_chunk_scan.py

@@ -260,7 +260,7 @@ def fn(*args: int) -> torch.Tensor:
 def main() -> None:
     """
     Main entry point that runs the attention kernel test with specific parameters.
-    Tests with batch size 2, 32 heads, 1024 sequence length, and 64-dimensional heads using float16.
+    Tests with batch size 2, 32 heads, 1024 sequence length, and 64-dimensional heads using half-precision.
     """
     test("zzzzzzz", 8, 80, 1, 4096, 256, 64, 128)
     test("zrzzzzr", 8, 80, 1, 4096, 256, 64, 128)  # D * x

test/test_examples.py

@@ -529,7 +529,7 @@ def test_bf16xint16(self):
             fn_name="_int16xbf16_gemm",
         )
 
-    @xfailIfPallas("Mosaic: Invalid vector type for load with f16 tiling")
+    @xfailIfPallas("Mosaic: offset not aligned to sublanes")
     def test_rms_norm_fwd(self):
         args = (
             torch.randn([128, 256], device=DEVICE, dtype=HALF_DTYPE),