Fix:

- condition triton import in gemm
- linting

Author: the-tuning-machine

benchmarks/benchmark_blockwise_scaled_linear_triton.py

@@ -12,12 +12,10 @@
     fp8_blockwise_act_quant,
     fp8_blockwise_weight_quant,
 )
-
 from torchao.quantization.quant_api import (
-    _int8_symm_per_token_reduced_range_quant_cutlass,
     _int4_symm_per_token_quant_cutlass,
+    _int8_symm_per_token_reduced_range_quant_cutlass,
 )
-
 from torchao.utils import is_sm_at_least_89
 
 
@@ -44,7 +42,7 @@ def get_blockwise_problem(
     assert dtype in [
         torch.float8_e4m3fn,
         torch.float8_e5m2,
-    ], f"dtype must be torch.float8_e4m3fn or torch.float8_e5m2"
+    ], "dtype must be torch.float8_e4m3fn or torch.float8_e5m2"
     dtype_max = torch.finfo(dtype).max
     A = (dtype_max * (2 * torch.rand(m, k, device=device) - 1)).to(dtype)
     A_scale = torch.randn((m, k // block_size), dtype=torch.half, device=device)

test/prototype/test_blockwise_triton.py

@@ -7,7 +7,6 @@
     fp8_blockwise_weight_dequant,
     fp8_blockwise_weight_quant,
 )
-
 from torchao.utils import is_sm_at_least_89
 
 BLOCKWISE_SIZE_MNK = [

torchao/prototype/blockwise_fp8/__init__.py

@@ -2,8 +2,8 @@
 from .blockwise_linear import BlockwiseQuantLinear
 from .blockwise_quantization import (
     fp8_blockwise_act_quant,
-    fp8_blockwise_weight_quant,
     fp8_blockwise_weight_dequant,
+    fp8_blockwise_weight_quant,
 )
 
 __all__ = [

torchao/prototype/blockwise_fp8/blockwise_fp8_gemm_triton.py

@@ -1,7 +1,9 @@
 import torch
-import triton
-import triton.language as tl
-from triton import Config
+
+if torch.cuda.is_available():
+    import triton
+    import triton.language as tl
+    from triton import Config
 
 # Original implementation at https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/kernel.py
 

torchao/prototype/blockwise_fp8/blockwise_quantization.py

@@ -55,7 +55,7 @@ def fp8_blockwise_act_quant(
     assert dtype in [
         torch.float8_e4m3fn,
         torch.float8_e5m2,
-    ], f"dtype must be torch.float8_e4m3fn or torch.float8_e5m2"
+    ], "dtype must be torch.float8_e4m3fn or torch.float8_e5m2"
     y = torch.empty_like(x, dtype=dtype)
     s = x.new_empty(*x.size()[:-1], x.size(-1) // block_size, dtype=torch.float32)
     grid = lambda meta: (triton.cdiv(x.numel(), meta["BLOCK_SIZE"]),)
@@ -117,7 +117,7 @@ def fp8_blockwise_weight_quant(
     assert dtype in [
         torch.float8_e4m3fn,
         torch.float8_e5m2,
-    ], f"dtype must be torch.float8_e4m3fn or torch.float8_e5m2"
+    ], "dtype must be torch.float8_e4m3fn or torch.float8_e5m2"
     M, N = x.size()
     y = torch.empty_like(x, dtype=dtype)
     s = x.new_empty(M // block_size, N // block_size, dtype=torch.float32)