Fix dtype issue for TMA kernel (pytorch#3072)

Summary:
Pull Request resolved: pytorch#3072

X-link: facebookresearch/FBGEMM#166

Using original tensor type instead of the TMA descriptor type for casting.

Reviewed By: choutim

Differential Revision: D62187577

fbshipit-source-id: c27ad275821e46f93bc8a644d222fce348c064da

fbgemm_gpu/experimental/gemm/triton_gemm/fp8_gemm.py

@@ -720,6 +720,7 @@ def _kernel_matmul_fp8_row_tma_persistent(
     stride_cm,
     stride_cn,
     dot_out_dtype: tl.constexpr,
+    c_dtype: tl.constexpr,
     allow_tf32: tl.constexpr,
     fp8_fast_accum: tl.constexpr,
     BLOCK_M: tl.constexpr,
@@ -839,8 +840,7 @@ def _kernel_matmul_fp8_row_tma_persistent(
                 )
                 acc += bias[None, :]
 
-            acc = acc.to(C_ptr.dtype.element_ty)
-
+            acc = acc.to(c_dtype)
             tl._experimental_descriptor_store(C_ptr, acc, [offs_am, offs_bn])
             acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=dot_out_dtype)
 
@@ -972,8 +972,8 @@ def matmul_fp8_row(
     # Reinterpret inputs into proper triton fp8 dtype.
     a_tl = convert_fp8_type(a, tl_dtype)
     b_tl = convert_fp8_type(b, tl_dtype)
-    M, N, K, m_key, n_key, k_key, c, dot_out_dtype_triton, device = prep_matmul(
-        a_tl, b_tl, dot_out_dtype
+    M, N, K, m_key, n_key, k_key, c, c_dtype_triton, dot_out_dtype_triton, device = (
+        prep_matmul(a_tl, b_tl, dot_out_dtype)
     )
 
     output_shape = a_shape[:-1] + (N,)
@@ -1086,7 +1086,6 @@ def persistent_grid_tma(META):
             desc_a,
             desc_b,
             desc_c,
-            # c,
             M,
             N,
             K,
@@ -1103,6 +1102,7 @@ def persistent_grid_tma(META):
             c.stride(0),
             c.stride(1),
             dot_out_dtype=dot_out_dtype_triton,
+            c_dtype=c_dtype_triton,
             allow_tf32=allow_tf32,
             fp8_fast_accum=fp8_fast_accum,
             GROUP_M=8,
@@ -1637,7 +1637,7 @@ def matmul_fp8_block(
     a_tl = convert_fp8_type(a, tl_dtype)
     b_tl = convert_fp8_type(b, tl_dtype)
 
-    M, N, K, m_key, n_key, k_key, c, dot_out_dtype_triton, device = prep_matmul(
+    M, N, K, m_key, n_key, k_key, c, _, dot_out_dtype_triton, device = prep_matmul(
         a_tl, b_tl, dot_out_dtype
     )
 
@@ -1765,7 +1765,7 @@ def get_matmul_tune(M: int, N: int, K: int) -> Tuple[int, int, int]:
 
 def prep_matmul(
     a: TensorWrapper, b: TensorWrapper, dot_out_dtype: Optional[torch.dtype]
-) -> Tuple[int, int, int, int, int, int, torch.Tensor, str, torch.device]:
+) -> Tuple[int, int, int, int, int, int, torch.Tensor, str, str, torch.device]:
     """
     Shared bookkeeping for a @ b.T matmul.
 
@@ -1808,6 +1808,7 @@ def prep_matmul(
         tl.float8e4b8,
     ]
     c_dtype = torch.bfloat16
+    c_dtype_triton = tl.bfloat16
 
     c = torch.empty((M, N), device=device, dtype=c_dtype)
     if dot_out_dtype is None:
@@ -1823,7 +1824,7 @@ def prep_matmul(
         else:
             dot_out_dtype_triton = tl.int32
 
-    return M, N, K, m_key, n_key, k_key, c, dot_out_dtype_triton, device
+    return M, N, K, m_key, n_key, k_key, c, c_dtype_triton, dot_out_dtype_triton, device
 
 
 @triton.autotune(