rename splitk code to not mention float8, try 2

Differential Revision: D59977582 Pull Request resolved: #529
pytorch · Jul 24, 2024 · 10deaab · 10deaab
1 parent 5eca730
commit 10deaab
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 2 deletions.
diff --git a/test/dtypes/test_fp8.py → test/prototype/test_splitk.py b/test/dtypes/test_fp8.py → test/prototype/test_splitk.py
@@ -10,7 +10,7 @@
 from torchao.utils import TORCH_VERSION_AFTER_2_4
 
 try:
-    from torchao.prototype.fp8 import gemm_split_k, to_float8
+    from torchao.prototype.splitk import gemm_split_k, to_float8
     triton_available = True
 except ImportError:
     triton_available = False
@@ -38,7 +38,7 @@ def test_gemm_split_k(self):
         x_fp8, x_inv_s = to_float8(x, dtype=qdtype)
         w_fp8, w_inv_s = to_float8(w, dtype=qdtype)
 
-        y_torch, _ = torch._scaled_mm(x_fp8, w_fp8.t(), out_dtype=dtype, scale_a=x_inv_s, scale_b=w_inv_s)
+        y_torch = torch._scaled_mm(x_fp8, w_fp8.t(), out_dtype=dtype, scale_a=x_inv_s, scale_b=w_inv_s)
         y_triton = gemm_split_k(x_fp8, w_fp8.t(), scale_a=x_inv_s.item(), scale_b=w_inv_s.item())
         y_fp16 = torch.nn.functional.linear(x, w)
 

diff --git a/torchao/prototype/fp8/__init__.py → torchao/prototype/splitk/__init__.py b/torchao/prototype/fp8/__init__.py → torchao/prototype/splitk/__init__.py
diff --git a/torchao/prototype/fp8/splitk_gemm.py → torchao/prototype/splitk/splitk_gemm.py b/torchao/prototype/fp8/splitk_gemm.py → torchao/prototype/splitk/splitk_gemm.py