FP8 Blockwise Training: triton_op for dense model (#3402)

added triton_op to two gemm kernels

Author: agolajko

torchao/prototype/blockwise_fp8_training/kernels.py

@@ -96,14 +96,15 @@ def triton_fp8_gemm_1x128_128x128_kernel(
     tl.store(c_ptrs, c, mask=c_mask)
 
 
+@triton_op("torchao::triton_fp8_gemm_1x128_128x128", mutates_args={})
 def triton_fp8_gemm_1x128_128x128(
     a: torch.Tensor,  # (M, K)
     b: torch.Tensor,  # (K, N)
     a_s: torch.Tensor,  # (M, K // block_size)
     b_s: torch.Tensor,  # (K // block_size, N // block_size)
     block_size: int = 128,
     out_dtype: torch.dtype = torch.float32,
-):
+) -> torch.Tensor:
     # 'a' must be in row-major layout, 'b' must be in column-major layout
     assert _is_row_major(a), "a must be row-major"
     assert _is_column_major(b), "b must be column-major"
@@ -214,14 +215,15 @@ def triton_fp8_gemm_1x128_128x1_kernel(
     tl.store(c_ptrs, c, mask=c_mask)
 
 
+@triton_op("torchao::triton_fp8_gemm_1x128_128x1", mutates_args={})
 def triton_fp8_gemm_1x128_128x1(
     a: torch.Tensor,  # (M, K)
     b: torch.Tensor,  # (K, N)
     a_s: torch.Tensor,  # (M, K // block_size) reciprocals of scales
     b_s: torch.Tensor,  # (K // block_size, N) reciprocals of scales
     block_size: int = 128,
     out_dtype: torch.dtype = torch.float32,
-):
+) -> torch.Tensor:
     # 'a' must be in row-major layout, 'b' must be in column-major layout
     assert _is_row_major(a), "a must be row-major"
     assert _is_column_major(b), "b must be column-major"