Makes fallback float8 1x128 by 128x128 gemm output bfloat16

test/kernel/test_blockwise_triton.py

@@ -66,6 +66,7 @@ def test_blockwise_fp8_gemm(M, N, K, dtype):
     A_q, A_s = fp8_blockwise_act_quant(A, dtype=dtype)
     B_q, B_s = fp8_blockwise_weight_quant(B, dtype=dtype)
     C_q = blockwise_fp8_gemm(A_q, A_s, B_q, B_s)
+    assert C_q.dtype == torch.bfloat16, "unsupported"
     error = torch.linalg.vector_norm(C - C_q) / torch.linalg.vector_norm(C)
     print(f"Relative Error: {error.item():.6f}")
 

torchao/kernel/blockwise_quantization.py

@@ -92,7 +92,7 @@ def blockwise_fp8_gemm(
         M = a.numel() // K
         N = b.size(0)
         M_BUCKET = math.ceil(math.log2(M))
-        c = a.new_empty(*a.size()[:-1], N, dtype=torch.get_default_dtype())
+        c = a.new_empty(*a.size()[:-1], N, dtype=torch.bfloat16)
         grid = lambda META: (
             triton.cdiv(M, META["BLOCK_SIZE_M"]),
             triton.cdiv(N, META["BLOCK_SIZE_N"]),
@@ -105,7 +105,7 @@ def blockwise_fp8_gemm(
     @blockwise_fp8_gemm.register_fake
     def _(a, a_s, b, b_s, block_size=128):
         N = b.size(0)
-        c = a.new_empty(*a.size()[:-1], N, dtype=torch.get_default_dtype())
+        c = a.new_empty(*a.size()[:-1], N, dtype=torch.bfloat16)
         return c
 
     @triton.jit