[BENCH] fixed dependencies (#6436)

Author: ptillet

bench/triton_bench/matmul_ogs_details/_common.py

@@ -1,4 +1,5 @@
-from ki.safe_import import tl, triton
+import triton
+import triton.language as tl
 
 # -----------------------------------------------------------------------------
 #                                  Utilities

bench/triton_bench/matmul_ogs_details/opt_flags.py

@@ -1,9 +1,9 @@
 from dataclasses import dataclass
+import triton
 
 import torch
 
-from ki.meta import cuda_capability_geq
-from ki.safe_import import triton
+from triton_bench.meta import cuda_capability_geq
 
 from . import opt_flags_amd, opt_flags_nvidia
 

bench/triton_bench/matmul_ogs_details/opt_flags_amd.py

@@ -1,6 +1,5 @@
 import torch
-
-from ki.safe_import import triton
+import triton
 
 
 def compute_block_nk(n, block_m, grid_m, num_xcds, lhs_dtype, rhs_dtype, microscaling_ctx):

bench/triton_bench/matmul_ogs_details/opt_flags_nvidia.py

@@ -1,7 +1,6 @@
 import torch
-
-from ki.meta import cuda_capability_geq
-from ki.safe_import import triton
+import triton
+from triton_bench.meta import cuda_capability_geq
 
 
 def compute_grid_size(routing_data, m, n, block_m, block_n):

bench/triton_bench/mxfp.py

@@ -1,11 +1,9 @@
 from enum import Enum
-
+import triton
+import triton.language as tl
 import torch
 import torch.nn.functional as F
 
-from ki.meta import is_float8_dtype
-from ki.safe_import import tl, triton
-
 # -----------------------------------------------------------------------------
 #                      Dequantization / Quantization Utilities
 # -----------------------------------------------------------------------------
@@ -476,7 +474,7 @@ def upcast_from_mxfp(tensor: torch.Tensor, scale: torch.Tensor, dtype: torch.dty
         assert -ndim <= swizzle_axis < ndim, f"Invalid swizzle axis {swizzle_axis=}"
         swizzle_axis = swizzle_axis if swizzle_axis >= 0 else swizzle_axis + ndim
 
-    multiplier = 1 if is_float8_dtype(tensor.dtype) else 2
+    multiplier = 1 if "float8" in str(tensor.dtype) else 2
     logical_quant_dim_shape = tensor.shape[axis] * multiplier
     assert tensor.ndim == scale.ndim, (f"Weight and scale must have the same number of dimensions. "
                                        f"Got {tensor.ndim=} and {scale.ndim=}")
@@ -560,7 +558,7 @@ def downcast_to_mxfp_torch(src_tensor: torch.Tensor, out_quant_type: torch.dtype
         assert -ndim <= swizzle_axis < ndim, f"Invalid swizzle axis {swizzle_axis=}"
         swizzle_axis = swizzle_axis if swizzle_axis >= 0 else swizzle_axis + ndim
     is_fp4 = out_quant_type == torch.uint8
-    is_fp8 = is_float8_dtype(out_quant_type)
+    is_fp8 = "float8" in str(out_quant_type)
     assert is_fp4 or is_fp8, f"Invalid input tensor dtype {out_quant_type}"
 
     device = src_tensor.device

bench/triton_bench/testing.py

@@ -6,7 +6,7 @@
 
 import torch
 
-from ki.meta import MAX_FINITE_FLOAT8E4B8, MAX_FINITE_FLOAT8E4NV, MAX_FINITE_FLOAT8E5
+from triton_bench.meta import MAX_FINITE_FLOAT8E4B8, MAX_FINITE_FLOAT8E4NV, MAX_FINITE_FLOAT8E5
 
 
 def assert_equal(ref, tri):
@@ -18,12 +18,7 @@ def assert_equal(ref, tri):
 
 def assert_close(ref, tri, maxtol=None, rmstol=None, description="--", verbose=True):
     if tri.dtype.itemsize == 1:
-        # TODO:
-        # switch to ref.to(tri.dtype) when Triton does
-        # RTNE on A100
-        from ki.tritium import type
-
-        ref_as_type = type(ref, tri.dtype)
+        ref_as_type = ref.to(tri.dtype)
         if ref.dtype == tri.dtype:
             assert torch.all(ref_as_type == tri)
             return