remove ssyrk which is not used

Signed-off-by: Hao Wu <skyw@nvidia.com>

Author: skyw

emerging_optimizers/orthogonalized_optimizers/muon.py

@@ -82,18 +82,19 @@ def __init__(
         if num_ns_steps < 1:
             raise ValueError(f"num_ns_steps must be at least 1, got {num_ns_steps}")
 
-        if torch.cuda.is_available():
-            sm_version = torch.cuda.get_device_capability()
-        else:
-            sm_version = (0, 0)
-        if not triton_kernels.HAS_TRITON_340:  # type: ignore[attr-defined]
-            logging.error("Triton 3.4.0 or higher is required for use_syrk to be True.")
-            use_syrk = False
-        elif sm_version not in ((8, 0), (9, 0), (10, 0), (10, 3)):
-            logging.error(
-                f"Correctness of Triton kernel on SM {sm_version} cannot be guaranteed. Setting use_syrk to False."
-            )
-            use_syrk = False
+        if use_syrk:
+            if torch.cuda.is_available():
+                sm_version = torch.cuda.get_device_capability()
+            else:
+                sm_version = (0, 0)
+            if not triton_kernels.HAS_TRITON_340:  # type: ignore[attr-defined]
+                logging.error("Triton 3.4.0 or higher is required for use_syrk to be True.")
+                use_syrk = False
+            elif sm_version not in ((8, 0), (9, 0), (10, 0), (10, 3)):
+                logging.error(
+                    f"Correctness of Triton kernel on SM {sm_version} cannot be guaranteed. Setting use_syrk to False."
+                )
+                use_syrk = False
         orthogonalize_fn = partial(
             newton_schulz, steps=num_ns_steps, coefficient_type=coefficient_type, use_syrk=use_syrk
         )

emerging_optimizers/triton_kernels/syrk.py

@@ -26,130 +26,7 @@
     HAS_TRITON_340 = False
 
 
-__all__ = ["ssyrk", "tsyrk_ex", "HAS_TRITON_340"]
-
-
-@triton.jit
-def cvt_tf32_rn(x: tl.tensor) -> tl.tensor:
-    return tl.inline_asm_elementwise("cvt.rna.tf32.f32 $0, $1;", "=r, r", [x], dtype=tl.float32, is_pure=True, pack=1)
-
-
-@triton.autotune(
-    configs=[
-        triton.Config({"TILE_N": tn, "TILE_K": tk}, num_warps=nw, num_stages=ns)
-        for tn in (64, 128)
-        for tk in (16, 32, 64)
-        for nw in (4, 8)
-        for ns in (3, 4)
-    ],
-    key=["N", "K", "ALLOW_TF32"],
-)
-@triton.jit
-def syrk_op_n_simple_kernel(
-    c_ptr,
-    a_ptr,
-    N: tl.constexpr,
-    K: tl.constexpr,
-    STRIDE_N: tl.constexpr,
-    STRIDE_K: tl.constexpr,
-    ALLOW_TF32: tl.constexpr,
-    TILE_N: tl.constexpr,
-    TILE_K: tl.constexpr,
-):
-    # receives tensor of shape (N, K)
-    # computes A * A^T (-> produces NxN)
-
-    pid_row = tl.program_id(0)
-    pid_col = tl.program_id(1)
-
-    IS_BELOW_DIAG = pid_row < pid_col
-    IS_ABOVE_DIAG = pid_row > pid_col
-
-    if IS_ABOVE_DIAG:
-        return
-
-    offs_row = pid_row * TILE_N + tl.arange(0, TILE_N)
-    offs_col = pid_col * TILE_N + tl.arange(0, TILE_N)
-    offs_k = tl.arange(0, TILE_K)
-
-    mask_row = offs_row < N
-    mask_col = offs_col < N
-
-    a_ptrs_x = a_ptr + offs_row[:, None] * STRIDE_N + offs_k[None, :] * STRIDE_K
-    a_ptrs_y = a_ptr + offs_col[None, :] * STRIDE_N + offs_k[:, None] * STRIDE_K
-
-    acc = tl.zeros((TILE_N, TILE_N), dtype=tl.float32)
-
-    num_tiles_k = tl.cdiv(K, TILE_K)
-    for k in range(0, num_tiles_k):
-        mask_k = offs_k < K - k * TILE_K
-        mask_x = mask_row[:, None] & mask_k[None, :]
-        mask_y = mask_col[None, :] & mask_k[:, None]
-        x = tl.load(a_ptrs_x, mask=mask_x, other=0.0)
-        y = tl.load(a_ptrs_y, mask=mask_y, other=0.0)
-
-        if ALLOW_TF32 == 0:
-            acc = tl.dot(x, y, acc=acc, input_precision="ieee")
-        elif ALLOW_TF32 == 1:
-            x = cvt_tf32_rn(x)
-            y = cvt_tf32_rn(y)
-            acc = tl.dot(x, y, acc=acc, input_precision="tf32")
-        else:
-            tl.static_assert(False, "Unsupported precision.")
-
-        a_ptrs_x += TILE_K * STRIDE_K
-        a_ptrs_y += TILE_K * STRIDE_K
-
-    # store diagonal or below diagonal values
-    c_ptrs = c_ptr + offs_row[:, None] * N + offs_col[None, :]
-    mask_c = mask_row[:, None] & mask_col[None, :]
-    tl.store(c_ptrs, acc, mask=mask_c)
-
-    # store replicated values above diagonal
-    if IS_BELOW_DIAG:
-        c_ptrs_diag = c_ptr + offs_col[None, :] * N + offs_row[:, None]
-        tl.store(c_ptrs_diag, acc, mask=mask_c)
-
-
-def ssyrk(a: torch.Tensor, trans: bool = False) -> torch.Tensor:
-    """Triton implementation of BLAS ssyrk operation.
-
-    Note:
-        This function assumes row major layout of the input tensor.
-
-    TODO(mstadler): Add support for alpha, beta and c.
-
-    Args:
-        a: Input tensor of shape (N, K) or (K, N)
-        trans: Whether to compute A * A^T (trans=False) or A^T * A (trans=True)
-
-    Returns:
-        Output tensor of shape (N, N)
-    """
-    assert a.dim() == 2, "Input tensor must be 2D"
-    N, K = a.shape
-    if trans:
-        raise NotImplementedError("Transpose is not supported yet.")
-
-    STRIDE_N = a.stride(0)
-    STRIDE_K = a.stride(1)
-
-    if (fp32_matmul_prec := torch.get_float32_matmul_precision()) == "highest":
-        ALLOW_TF32 = 0
-    elif fp32_matmul_prec == "high":
-        ALLOW_TF32 = 1
-    else:
-        raise ValueError(f"Unsupported precision {fp32_matmul_prec}, only 'highest' and 'high' are supported.")
-
-    c = torch.empty((N, N), dtype=a.dtype, device=a.device)
-
-    def grid(META):
-        return (triton.cdiv(N, META["TILE_N"]), triton.cdiv(N, META["TILE_N"]))
-
-    if not trans:
-        syrk_op_n_simple_kernel[grid](c, a, N, K, STRIDE_N, STRIDE_K, ALLOW_TF32)
-
-    return c
+__all__ = ["tsyrk_ex", "HAS_TRITON_340"]
 
 
 def prune_invalid_configs(configs: list[triton.Config], named_args: dict, **kwargs) -> list[triton.Config]:

tests/test_triton_kernels.py

@@ -15,81 +15,7 @@
 import torch
 from absl.testing import absltest, parameterized
 
-from emerging_optimizers import triton_kernels, utils
-
-
-class SsyrkTest(parameterized.TestCase):
-    @parameterized.product(
-        ({"n": 5, "k": 7, "rtol": 1e-5}, {"n": 17, "k": 23, "rtol": 1e-3}, {"n": 127, "k": 255, "rtol": 2e-2}),
-        ({"trans": False},),
-    )
-    def test_ssyrk_fp32_close_to_matmul(self, n: int, k: int, trans: bool, rtol: float):
-        a = torch.randn(n, k, device="cuda")
-        a_warmup = torch.randn_like(a, device=a.device)
-        if trans:
-            ref = a.T @ a
-        else:
-            ref = a @ a.T
-        with utils.fp32_matmul_precision("highest"):
-            # warmup the triton kernel to avoid the wrong result from the first run.
-            _ = triton_kernels.ssyrk(a_warmup, trans=trans)
-            c = triton_kernels.ssyrk(a, trans=trans)
-        torch.testing.assert_close(c, ref, atol=0, rtol=rtol)
-
-    @absltest.skipIf(torch.cuda.get_device_capability()[0] < 8, "TF32 needs compute capability >8.0")
-    @parameterized.product(
-        ({"n": 5, "k": 7, "rtol": 0.05}, {"n": 17, "k": 23, "rtol": 0.5}),
-        ({"trans": False},),
-    )
-    def test_ssyrk_tf32_not_far_from_matmul(self, n: int, k: int, trans: bool, rtol: float):
-        a = torch.randn(n, k, device="cuda")
-        a_warmup = torch.randn_like(a, device=a.device)
-        if trans:
-            ref = a.T @ a
-        else:
-            ref = a @ a.T
-        with utils.fp32_matmul_precision("high"):
-            # warmup the triton kernel to avoid the wrong result from the first run.
-            _ = triton_kernels.ssyrk(a_warmup, trans=trans)
-            c = triton_kernels.ssyrk(a, trans=trans)
-        torch.testing.assert_close(c, ref, atol=0, rtol=rtol)
-
-
-class SsyrkIntegerInputTest(parameterized.TestCase):
-    @parameterized.product(
-        ({"n": 5, "k": 7}, {"n": 17, "k": 23}, {"n": 127, "k": 255}),
-        ({"trans": False},),
-    )
-    def test_ssyrk_fp32_match_matmul(self, n: int, k: int, trans: bool):
-        a = torch.randint(-10, 10, (n, k), device="cuda", dtype=torch.float32)
-        a_warmup = torch.randint_like(a, -10, 10, device=a.device, dtype=torch.float32)
-        if trans:
-            ref = a.T @ a
-        else:
-            ref = a @ a.T
-        with utils.fp32_matmul_precision("highest"):
-            # warmup the triton kernel to avoid the wrong result from the first run.
-            _ = triton_kernels.ssyrk(a_warmup, trans=trans)
-            c = triton_kernels.ssyrk(a, trans=trans)
-        torch.testing.assert_close(c, ref, atol=0, rtol=0)
-
-    @absltest.skipIf(torch.cuda.get_device_capability()[0] < 8, "TF32 needs compute capability >8.0")
-    @parameterized.product(
-        ({"n": 5, "k": 7}, {"n": 17, "k": 23}, {"n": 127, "k": 255}),
-        ({"trans": False},),
-    )
-    def test_ssyrk_tf32_match_matmul(self, n: int, k: int, trans: bool):
-        a = torch.randint(-10, 10, (n, k), device="cuda", dtype=torch.float32)
-        a_warmup = torch.randint_like(a, -10, 10, device=a.device, dtype=torch.float32)
-        if trans:
-            ref = a.T @ a
-        else:
-            ref = a @ a.T
-        with utils.fp32_matmul_precision("high"):
-            # warmup the triton kernel to avoid the wrong result from the first run.
-            _ = triton_kernels.ssyrk(a_warmup, trans=trans)
-            c = triton_kernels.ssyrk(a, trans=trans)
-        torch.testing.assert_close(c, ref, atol=0, rtol=0)
+from emerging_optimizers import triton_kernels
 
 
 class TsyrkTest(parameterized.TestCase):