rename splitk code to not mention float8, try 2

Differential Revision: D59977582

Pull Request resolved: pytorch#529

Author: vkuzo

test/dtypes/test_fp8.py renamed to 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)