Fix w8a8 lightop restriction

Signed-off-by: wenjh <wenjh@sugon.com>

transformer_engine/pytorch/cpp_extensions/gemm.py

@@ -46,17 +46,18 @@ __all__ = [
     "batchgemm",
 ]
 
-def w8a8_block_int8_matmul_wgrad_batched_native(A_list, B_list, As_list, Bs_list, C_list, accumulate, output_dtype=torch.float16):
+def w8a8_block_int8_matmul_wgrad_batched_native(A_list, B_list, As_list, Bs_list, C_list, accumulate, out_dtype=torch.float16):
     for i in range(len(C_list)):
         assert C_list[i] is not None
-        if get_device_compute_capability() >= (9, 3) and blockwise_fp8_block_len == 128:
+        if get_device_compute_capability() >= (9, 3) and blockwise_fp8_block_len == 128  and ((out_dtype is torch.bfloat16) or (out_dtype is torch.float16)):
             C_list[i] = lightop.gemm_w8a8_wgrad_asm(
-                A_list[i], B_list[i], As_list[i], Bs_list[i], C_list[i], accumulate, blockwise_fp8_block_len, output_dtype, "TN"
+                A_list[i], B_list[i], As_list[i], Bs_list[i], C_list[i], accumulate, blockwise_fp8_block_len, out_dtype, "TN"
             )
         else:
+            warnings.warn("Lightop is not available. Using default implementation for w8a8.")
             C_list[i], _ = w8a8_block_int8_matmul_wgrad(
                 A_list[i], B_list[i], As_list[i], Bs_list[i], C_list[i], accumulate, blockwise_fp8_block_len,
-                output_dtype,
+                out_dtype,
                 None
             )
     return C_list
@@ -75,7 +76,7 @@ def w8a8_int8_general_gemm(
         qw_data = (A._rowwise_data.view(dtype=torch.int8))
         ref_scales_x = B._rowwise_scale_inv
         ref_scales_w = A._rowwise_scale_inv
-        if get_device_compute_capability() >= (9, 3) and blockwise_fp8_block_len == 128:
+        if get_device_compute_capability() >= (9, 3) and blockwise_fp8_block_len == 128 and ((out_dtype is torch.bfloat16) or (out_dtype is torch.float16)):
             y = lightop.gemm_w8a8_asm(qx_data, qw_data, ref_scales_x, ref_scales_w, [blockwise_fp8_block_len, blockwise_fp8_block_len], out_dtype, 'TN')
         else:
             warnings.warn("Lightop is not available. Using default implementation for w8a8.")
@@ -84,7 +85,7 @@ def w8a8_int8_general_gemm(
     elif layout == "NN":
         assert accumulate is False, "Accumulate not supported in w8a8_general_gemm with NN layout"
         assert out is None, "Output tensor not supported in w8a8_general_gemm with NN layout"
-        if get_device_compute_capability() >= (9, 3) and blockwise_fp8_block_len == 128:
+        if get_device_compute_capability() >= (9, 3) and blockwise_fp8_block_len == 128 and ((out_dtype is torch.bfloat16) or (out_dtype is torch.float16)):
             qdout_data = (B._rowwise_data.view(dtype=torch.int8))
             qw_data = (A._rowwise_data.view(dtype=torch.int8))
             ref_scales_dout = B._rowwise_scale_inv
@@ -103,7 +104,7 @@ def w8a8_int8_general_gemm(
         qx_data = (A._columnwise_data.view(dtype=torch.int8))
         ref_scales_dout = B._columnwise_scale_inv
         ref_scales_x = A._columnwise_scale_inv
-        if get_device_compute_capability() >= (9, 3) and blockwise_fp8_block_len == 128:
+        if get_device_compute_capability() >= (9, 3) and blockwise_fp8_block_len == 128 and ((out_dtype is torch.bfloat16) or (out_dtype is torch.float16)):
             out = lightop.gemm_w8a8_wgrad_asm(qdout_data, qx_data, ref_scales_dout, ref_scales_x, out, accumulate, [blockwise_fp8_block_len, blockwise_fp8_block_len], out_dtype, 'TN')
         else:
             warnings.warn("Lightop is not available. Using default implementation for w8a8.")