trasmuon-pytorch 0.1.0__tar.gz

trasmuon_pytorch-0.1.0/.gitignore

@@ -0,0 +1,11 @@
+__pycache__/
+*.py[cod]
+.venv/
+venv/
+dist/
+build/
+*.egg-info/
+.DS_Store
+.env
+.env.*
+!.env.example

trasmuon_pytorch-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) TrasMuon Authors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE

trasmuon_pytorch-0.1.0/NOTICE.md

@@ -0,0 +1,56 @@
+# Notice
+
+This repository is a trimmed research code release for the TrasMuon project and paper at <https://arxiv.org/abs/2602.13498>.
+
+The repository remains MIT licensed under `LICENSE`. It also retains material adapted from the following upstream sources.
+
+## Modded-NanoGPT
+
+Source: <https://github.com/KellerJordan/modded-nanogpt>
+
+MIT License
+
+Copyright (c) 2024 Keller Jordan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+## Moonlight
+
+Source: <https://github.com/MoonshotAI/Moonlight>
+
+The MIT License (MIT)
+Copyright (c) 2025 Moonshot AI
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

trasmuon_pytorch-0.1.0/PKG-INFO

@@ -0,0 +1,105 @@
+Metadata-Version: 2.5
+Name: trasmuon-pytorch
+Version: 0.1.0
+Summary: Trust-region adaptive scaling for orthogonalized momentum optimizers in PyTorch
+Project-URL: Homepage, https://github.com/zangjiucheng/TrasMuon
+Project-URL: Source, https://github.com/zangjiucheng/TrasMuon
+Project-URL: Issues, https://github.com/zangjiucheng/TrasMuon/issues
+Project-URL: Paper, https://arxiv.org/abs/2602.13498
+Author: TrasMuon Authors
+License-Expression: MIT
+License-File: LICENSE
+License-File: NOTICE.md
+Keywords: deep-learning,muon,optimizer,pytorch
+Classifier: Development Status :: 3 - Alpha
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Requires-Dist: torch>=2.5.1
+Requires-Dist: triton
+Description-Content-Type: text/markdown
+
+# TrasMuon
+
+[![arXiv](https://img.shields.io/badge/arXiv-2602.13498-b31b1b.svg)](https://arxiv.org/abs/2602.13498)
+
+Optimizer implementation for **TrasMuon: Trust-Region Adaptive Scaling for Orthogonalized Momentum Optimizers**. TrasMuon combines Muon with adaptive scaling, trust-region column gating, and schedule-free smoothing.
+
+![TrasMuon algorithm: adaptive scaling, trust-region column gating, and schedule-free smoothing](https://raw.githubusercontent.com/zangjiucheng/TrasMuon/main/assets/trasmuon-algorithm.png)
+
+## Installation
+
+Requires Python ≥ 3.10, PyTorch ≥ 2.5.1, and Triton. Use Linux with CUDA-enabled PyTorch for GPU execution.
+
+Install from GitHub:
+
+```bash
+pip install git+https://github.com/zangjiucheng/TrasMuon.git
+```
+
+After the first PyPI release, install with:
+
+```bash
+pip install trasmuon-pytorch
+```
+
+## Usage
+
+```python
+from trasmuon import TrasMuon
+
+optimizer = TrasMuon(model.parameters(), lr=1e-3, weight_decay=0.01)
+
+optimizer.zero_grad(set_to_none=True)
+loss = loss_fn(model(inputs), targets)
+loss.backward()
+optimizer.step()
+```
+
+Two-dimensional parameters use TrasMuon; other parameter shapes use the optimizer's AdamW fallback. See [`TrasMuon.__init__`](https://github.com/zangjiucheng/TrasMuon/blob/main/trasmuon/optimizer.py) for available options.
+
+## Files
+
+- [`trasmuon/optimizer.py`](https://github.com/zangjiucheng/TrasMuon/blob/main/trasmuon/optimizer.py): optimizer implementation.
+- [`trasmuon/newton_schulz_triton.py`](https://github.com/zangjiucheng/TrasMuon/blob/main/trasmuon/newton_schulz_triton.py): required Newton–Schulz kernel dependency.
+- [`trasmuon/__init__.py`](https://github.com/zangjiucheng/TrasMuon/blob/main/trasmuon/__init__.py): exports `TrasMuon`.
+
+## Publishing
+
+Build and validate a release locally:
+
+```bash
+python -m pip install build twine
+python -m build
+python -m twine check --strict dist/*
+```
+
+For the first release, add a [pending Trusted Publisher](https://pypi.org/manage/account/publishing/) on PyPI with:
+
+- Project: `trasmuon-pytorch`
+- Owner / repository: `zangjiucheng` / `TrasMuon`
+- Workflow: `publish.yml`
+- Environment: `pypi`
+
+Create the matching `pypi` environment in the GitHub repository settings. Update the version in `pyproject.toml`, then publish a GitHub Release with the matching tag (initially `v0.1.0`). The release workflow builds and validates both distributions, then publishes them through Trusted Publishing. No PyPI API token is needed. See the [PyPI setup guide](https://docs.pypi.org/trusted-publishers/creating-a-project-through-oidc/).
+
+## Citation
+
+If you use this code, cite the TrasMuon paper:
+
+```bibtex
+@misc{cheng2026trasmuon,
+  title = {TrasMuon: Trust-Region Adaptive Scaling for Orthogonalized Momentum Optimizers},
+  author = {Peng Cheng and Jiucheng Zang and Qingnan Li and Liheng Ma and Yufei Cui and Yingxue Zhang and Boxing Chen and Ming Jian and Wen Tong},
+  year = {2026},
+  eprint = {2602.13498},
+  archivePrefix = {arXiv},
+  primaryClass = {cs.LG},
+  url = {https://arxiv.org/abs/2602.13498}
+}
+```
+
+## License
+
+This repository is released under the MIT license. See [LICENSE](https://github.com/zangjiucheng/TrasMuon/blob/main/LICENSE) and [NOTICE.md](https://github.com/zangjiucheng/TrasMuon/blob/main/NOTICE.md).

trasmuon_pytorch-0.1.0/README.md

@@ -0,0 +1,83 @@
+# TrasMuon
+
+[![arXiv](https://img.shields.io/badge/arXiv-2602.13498-b31b1b.svg)](https://arxiv.org/abs/2602.13498)
+
+Optimizer implementation for **TrasMuon: Trust-Region Adaptive Scaling for Orthogonalized Momentum Optimizers**. TrasMuon combines Muon with adaptive scaling, trust-region column gating, and schedule-free smoothing.
+
+![TrasMuon algorithm: adaptive scaling, trust-region column gating, and schedule-free smoothing](https://raw.githubusercontent.com/zangjiucheng/TrasMuon/main/assets/trasmuon-algorithm.png)
+
+## Installation
+
+Requires Python ≥ 3.10, PyTorch ≥ 2.5.1, and Triton. Use Linux with CUDA-enabled PyTorch for GPU execution.
+
+Install from GitHub:
+
+```bash
+pip install git+https://github.com/zangjiucheng/TrasMuon.git
+```
+
+After the first PyPI release, install with:
+
+```bash
+pip install trasmuon-pytorch
+```
+
+## Usage
+
+```python
+from trasmuon import TrasMuon
+
+optimizer = TrasMuon(model.parameters(), lr=1e-3, weight_decay=0.01)
+
+optimizer.zero_grad(set_to_none=True)
+loss = loss_fn(model(inputs), targets)
+loss.backward()
+optimizer.step()
+```
+
+Two-dimensional parameters use TrasMuon; other parameter shapes use the optimizer's AdamW fallback. See [`TrasMuon.__init__`](https://github.com/zangjiucheng/TrasMuon/blob/main/trasmuon/optimizer.py) for available options.
+
+## Files
+
+- [`trasmuon/optimizer.py`](https://github.com/zangjiucheng/TrasMuon/blob/main/trasmuon/optimizer.py): optimizer implementation.
+- [`trasmuon/newton_schulz_triton.py`](https://github.com/zangjiucheng/TrasMuon/blob/main/trasmuon/newton_schulz_triton.py): required Newton–Schulz kernel dependency.
+- [`trasmuon/__init__.py`](https://github.com/zangjiucheng/TrasMuon/blob/main/trasmuon/__init__.py): exports `TrasMuon`.
+
+## Publishing
+
+Build and validate a release locally:
+
+```bash
+python -m pip install build twine
+python -m build
+python -m twine check --strict dist/*
+```
+
+For the first release, add a [pending Trusted Publisher](https://pypi.org/manage/account/publishing/) on PyPI with:
+
+- Project: `trasmuon-pytorch`
+- Owner / repository: `zangjiucheng` / `TrasMuon`
+- Workflow: `publish.yml`
+- Environment: `pypi`
+
+Create the matching `pypi` environment in the GitHub repository settings. Update the version in `pyproject.toml`, then publish a GitHub Release with the matching tag (initially `v0.1.0`). The release workflow builds and validates both distributions, then publishes them through Trusted Publishing. No PyPI API token is needed. See the [PyPI setup guide](https://docs.pypi.org/trusted-publishers/creating-a-project-through-oidc/).
+
+## Citation
+
+If you use this code, cite the TrasMuon paper:
+
+```bibtex
+@misc{cheng2026trasmuon,
+  title = {TrasMuon: Trust-Region Adaptive Scaling for Orthogonalized Momentum Optimizers},
+  author = {Peng Cheng and Jiucheng Zang and Qingnan Li and Liheng Ma and Yufei Cui and Yingxue Zhang and Boxing Chen and Ming Jian and Wen Tong},
+  year = {2026},
+  eprint = {2602.13498},
+  archivePrefix = {arXiv},
+  primaryClass = {cs.LG},
+  url = {https://arxiv.org/abs/2602.13498}
+}
+```
+
+## License
+
+This repository is released under the MIT license. See [LICENSE](https://github.com/zangjiucheng/TrasMuon/blob/main/LICENSE) and [NOTICE.md](https://github.com/zangjiucheng/TrasMuon/blob/main/NOTICE.md).

trasmuon_pytorch-0.1.0/pyproject.toml

@@ -0,0 +1,33 @@
+[build-system]
+requires = ["hatchling>=1.26"]
+build-backend = "hatchling.build"
+
+[project]
+name = "trasmuon-pytorch"
+version = "0.1.0"
+description = "Trust-region adaptive scaling for orthogonalized momentum optimizers in PyTorch"
+readme = "README.md"
+requires-python = ">=3.10"
+license = "MIT"
+license-files = ["LICENSE", "NOTICE.md"]
+authors = [{ name = "TrasMuon Authors" }]
+keywords = ["pytorch", "optimizer", "muon", "deep-learning"]
+classifiers = [
+  "Development Status :: 3 - Alpha",
+  "Programming Language :: Python :: 3",
+  "Operating System :: POSIX :: Linux",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = ["torch>=2.5.1", "triton"]
+
+[project.urls]
+Homepage = "https://github.com/zangjiucheng/TrasMuon"
+Source = "https://github.com/zangjiucheng/TrasMuon"
+Issues = "https://github.com/zangjiucheng/TrasMuon/issues"
+Paper = "https://arxiv.org/abs/2602.13498"
+
+[tool.hatch.build.targets.wheel]
+packages = ["trasmuon"]
+
+[tool.hatch.build.targets.sdist]
+include = ["/trasmuon", "/assets", "/README.md", "/LICENSE", "/NOTICE.md", "/pyproject.toml"]

trasmuon_pytorch-0.1.0/trasmuon/__init__.py

@@ -0,0 +1,5 @@
+from .optimizer import TrasMuon
+
+__all__ = [
+    "TrasMuon",
+]

trasmuon_pytorch-0.1.0/trasmuon/newton_schulz_triton.py

@@ -0,0 +1,374 @@
+import torch
+import triton
+import triton.language as tl
+from torch import Tensor
+
+
+def _get_autotune_configs():
+    return [
+        triton.Config(
+            {
+                "BLOCK_SIZE_M": bm,
+                "BLOCK_SIZE_N": bn,
+                "BLOCK_SIZE_K": bk,
+                "GROUP_SIZE_M": 8,
+                "LOWER_UPPER": 1,
+            },
+            num_stages=stages,
+            num_warps=warps,
+        )
+        for bm in [64, 128]
+        for bn in [64, 128, 256]
+        for bk in [64, 128]
+        for stages, warps in [(3, 4), (3, 8), (4, 4)]
+        if bm // bn <= 2 and bn // bm <= 2
+    ]
+
+
+@triton.jit
+def _pid_to_block(
+    pid,
+    M,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+):
+    """
+    Helper function to map Triton program ID to (batch, row, col) of the output matrix.
+    """
+    # Split output matrix into blocks of size (BLOCK_SIZE_M, BLOCK_SIZE_N)
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(M, BLOCK_SIZE_N)
+
+    # Map PID to a single matrix in batch
+    batch_idx = pid // (num_pid_m * num_pid_n)
+    pid = pid % (num_pid_m * num_pid_n)
+
+    # Map PID to 2D grid of blocks
+    pid_m = pid // num_pid_n
+    pid_n = pid % num_pid_n
+    pid_m, pid_n = tl.swizzle2d(pid_m, pid_n, num_pid_m, num_pid_n, GROUP_SIZE_M)
+
+    m_idx = pid_m * BLOCK_SIZE_M
+    n_idx = pid_n * BLOCK_SIZE_N
+
+    return batch_idx, m_idx, n_idx
+
+
+@triton.autotune(
+    configs=_get_autotune_configs(),
+    key=["M", "K", "a_stride_r", "a_stride_c", "c_stride_r", "c_stride_c"],
+)
+@triton.jit
+def ns_line_1_kernel(
+    A_ptr,
+    C_ptr,
+    M,
+    K,
+    a_stride_b,
+    a_stride_r,
+    a_stride_c,
+    c_stride_b,
+    c_stride_r,
+    c_stride_c,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    LOWER_UPPER: tl.constexpr,
+):
+    """
+    Input A has shape (M, K)
+    Output C has shape (M, M)
+    Compute C = A @ A.T
+    """
+
+    pid = tl.program_id(axis=0)
+    batch_idx, m_idx, n_idx = _pid_to_block(
+        pid, M, BLOCK_SIZE_M, BLOCK_SIZE_N, GROUP_SIZE_M
+    )
+
+    # Skip blocks that don't need to be computed
+    skip_block_below_diag = (LOWER_UPPER == 0) and (n_idx + BLOCK_SIZE_N <= m_idx)
+    skip_block_above_diag = (LOWER_UPPER != 0) and (m_idx + BLOCK_SIZE_M <= n_idx)
+    if skip_block_below_diag or skip_block_above_diag:
+        return
+
+    # Index into one matrix of batch
+    A_ptr += batch_idx * a_stride_b
+    C_ptr += batch_idx * c_stride_b
+
+    # Create pointer arrays for A and A.T
+    offs_m = (m_idx + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_n = (n_idx + tl.arange(0, BLOCK_SIZE_N)) % M
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    a_ptrs = A_ptr + (offs_m[:, None] * a_stride_r + offs_k[None, :] * a_stride_c)
+    at_ptrs = A_ptr + (offs_k[:, None] * a_stride_c + offs_n[None, :] * a_stride_r)
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+
+    # Accumulate over blocks of K
+    for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
+        at = tl.load(at_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
+        accumulator = tl.dot(a, at, accumulator)
+        a_ptrs += BLOCK_SIZE_K * a_stride_c
+        at_ptrs += BLOCK_SIZE_K * a_stride_c
+
+    out_dtype = C_ptr.dtype.element_ty
+    output = accumulator.to(out_dtype)
+
+    # Store block of C
+    offs_cm = m_idx + tl.arange(0, BLOCK_SIZE_M)
+    offs_cn = n_idx + tl.arange(0, BLOCK_SIZE_N)
+    c_ptrs = C_ptr + (offs_cm[:, None] * c_stride_r + offs_cn[None, :] * c_stride_c)
+    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M)
+    tl.store(c_ptrs, output, mask=c_mask)
+
+    # Store block of C mirrored across the diagonal
+    c_ptrs_t = C_ptr + (offs_cn[:, None] * c_stride_r + offs_cm[None, :] * c_stride_c)
+    c_mask_t = (offs_cn[:, None] < M) & (offs_cm[None, :] < M)
+    tl.store(c_ptrs_t, output.T, mask=c_mask_t)
+
+
+def ns_line_1(A: Tensor, *, out: Tensor = None):
+    """
+    Launch Triton kernel to compute C = A @ A.T
+    """
+    if A.ndim > 3 or A.ndim < 2:
+        raise ValueError(f"Input tensor must be 2D or 3D, but got {A.ndim}D tensor.")
+
+    M, K = A.shape[-2:]
+
+    if out is None:
+        out = torch.empty((*A.shape[:-1], M), device=A.device, dtype=A.dtype)
+    assert out.size(-2) == M, "Output matrix has incorrect shape"
+    assert out.size(-1) == M, "Output matrix has incorrect shape"
+
+    batch_size = A.size(0) if A.ndim == 3 else 1
+    input_batch_stride = A.stride(0) if A.ndim == 3 else 0
+    output_batch_stride = out.stride(0) if out.ndim == 3 else 0
+
+    grid = lambda meta: (
+        batch_size
+        * triton.cdiv(M, meta["BLOCK_SIZE_M"])
+        * triton.cdiv(M, meta["BLOCK_SIZE_N"]),
+    )
+    ns_line_1_kernel[grid](
+        A_ptr=A,
+        C_ptr=out,
+        M=M,
+        K=K,
+        a_stride_b=input_batch_stride,
+        a_stride_r=A.stride(-2),
+        a_stride_c=A.stride(-1),
+        c_stride_b=output_batch_stride,
+        c_stride_r=out.stride(-2),
+        c_stride_c=out.stride(-1),
+    )
+
+    return out
+
+
+@triton.autotune(
+    configs=_get_autotune_configs(),
+    key=["M", "a_stride_r", "a_stride_c", "c_stride_r", "c_stride_c"],
+)
+@triton.jit
+def ns_line_2_kernel(
+    A_ptr,
+    C_ptr,
+    M,
+    a_stride_b,
+    a_stride_r,
+    a_stride_c,
+    c_stride_b,
+    c_stride_r,
+    c_stride_c,
+    alpha,
+    beta,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    LOWER_UPPER: tl.constexpr,
+):
+    """
+    Input A is square matrix with shape (M, M)
+    Output C has shape (M, M)
+    Compute C = alpha * A @ A.T + beta * A
+    """
+
+    pid = tl.program_id(axis=0)
+    batch_idx, m_idx, n_idx = _pid_to_block(
+        pid, M, BLOCK_SIZE_M, BLOCK_SIZE_N, GROUP_SIZE_M
+    )
+
+    # Skip blocks that don't need to be computed
+    skip_block_below_diag = (LOWER_UPPER == 0) and (n_idx + BLOCK_SIZE_N <= m_idx)
+    skip_block_above_diag = (LOWER_UPPER != 0) and (m_idx + BLOCK_SIZE_M <= n_idx)
+    if skip_block_below_diag or skip_block_above_diag:
+        return
+
+    # Index into one matrix of batch
+    A_ptr += batch_idx * a_stride_b
+    C_ptr += batch_idx * c_stride_b
+
+    # Create pointer arrays for A and A.T
+    offs_m = (m_idx + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_n = (n_idx + tl.arange(0, BLOCK_SIZE_N)) % M
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    a_ptrs = A_ptr + (offs_m[:, None] * a_stride_r + offs_k[None, :] * a_stride_c)
+    at_ptrs = A_ptr + (offs_k[:, None] * a_stride_c + offs_n[None, :] * a_stride_r)
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+
+    # Accumulate over blocks of K
+    for k in tl.range(0, tl.cdiv(M, BLOCK_SIZE_K)):
+        a = tl.load(a_ptrs, mask=offs_k[None, :] < M - k * BLOCK_SIZE_K, other=0.0)
+        at = tl.load(at_ptrs, mask=offs_k[:, None] < M - k * BLOCK_SIZE_K, other=0.0)
+        accumulator = tl.dot(a, at, accumulator)
+        a_ptrs += BLOCK_SIZE_K * a_stride_c
+        at_ptrs += BLOCK_SIZE_K * a_stride_c
+
+    # Load block of A to add (corresponds to the current block of C)
+    offs_am = m_idx + tl.arange(0, BLOCK_SIZE_M)
+    offs_an = n_idx + tl.arange(0, BLOCK_SIZE_N)
+    a_add_ptrs = A_ptr + (offs_am[:, None] * a_stride_r + offs_an[None, :] * a_stride_c)
+    a_add_mask = (offs_am[:, None] < M) & (offs_an[None, :] < M)
+    a_add = tl.load(a_add_ptrs, mask=a_add_mask, other=0.0).to(tl.float32)
+
+    # Apply alpha and beta
+    accumulator *= alpha
+    accumulator += a_add * beta
+
+    out_dtype = C_ptr.dtype.element_ty
+    output = accumulator.to(out_dtype)
+
+    # Store block of C
+    offs_cm = m_idx + tl.arange(0, BLOCK_SIZE_M)
+    offs_cn = n_idx + tl.arange(0, BLOCK_SIZE_N)
+    c_ptrs = C_ptr + (offs_cm[:, None] * c_stride_r + offs_cn[None, :] * c_stride_c)
+    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M)
+    tl.store(c_ptrs, output, mask=c_mask)
+
+    # Store block of C mirrored across the diagonal
+    c_ptrs_t = C_ptr + (offs_cn[:, None] * c_stride_r + offs_cm[None, :] * c_stride_c)
+    c_mask_t = (offs_cn[:, None] < M) & (offs_cm[None, :] < M)
+    tl.store(c_ptrs_t, output.T, mask=c_mask_t)
+
+
+def ns_line_2(A: Tensor, alpha: float, beta: float, *, out: Tensor = None):
+    """
+    Launch Triton kernel to compute C = alpha * A @ A.T + beta * A
+    """
+    if A.ndim > 3 or A.ndim < 2:
+        raise ValueError(f"Input tensor must be 2D or 3D, but got {A.ndim}D tensor.")
+
+    M, K = A.shape[-2:]
+    if M != K:
+        raise ValueError(
+            f"Input must be symmetric square matrix, but got shape {A.shape}"
+        )
+
+    if out is None:
+        out = torch.empty((*A.shape[:-1], M), device=A.device, dtype=A.dtype)
+    assert out.size(-2) == M, "Output matrix has incorrect shape"
+    assert out.size(-1) == M, "Output matrix has incorrect shape"
+
+    batch_size = A.size(0) if A.ndim == 3 else 1
+    input_batch_stride = A.stride(0) if A.ndim == 3 else 0
+    output_batch_stride = out.stride(0) if out.ndim == 3 else 0
+
+    grid = lambda meta: (
+        batch_size
+        * triton.cdiv(M, meta["BLOCK_SIZE_M"])
+        * triton.cdiv(M, meta["BLOCK_SIZE_N"]),
+    )
+    ns_line_2_kernel[grid](
+        A_ptr=A,
+        C_ptr=out,
+        M=M,
+        a_stride_b=input_batch_stride,
+        a_stride_r=A.stride(-2),
+        a_stride_c=A.stride(-1),
+        c_stride_b=output_batch_stride,
+        c_stride_r=out.stride(-2),
+        c_stride_c=out.stride(-1),
+        alpha=alpha,
+        beta=beta,
+    )
+
+    return out
+
+
+@torch.compile(dynamic=False, fullgraph=True)
+def zeropower_via_newtonschulz5(G: Tensor, epsilon: float = 1e-7):
+    """
+    Reference implementation of Newton-Schulz without Triton.
+    """
+    # Newton-Schulz constants
+    ns_consts = [
+        (4.0848, -6.8946, 2.9270),
+        (3.9505, -6.3029, 2.6377),
+        (3.7418, -5.5913, 2.3037),
+        (2.8769, -3.1427, 1.2046),
+        (2.8366, -3.0525, 1.2012),
+    ]
+
+    X = G.to(dtype=torch.bfloat16)
+    if G.size(-2) > G.size(-1):
+        X = X.mT
+
+    # Ensure spectral norm is at most 1
+    X = X / (X.norm(dim=(-2, -1), keepdim=True) + epsilon)
+
+    for a, b, c in ns_consts:
+        A = X @ X.mT
+        B = b * A + c * (A @ A)
+        X = a * X + B @ X
+
+    if G.size(-2) > G.size(-1):
+        X = X.mT
+    return X
+
+
+@torch.compile(dynamic=False, fullgraph=True)
+def newton_schulz_triton(G: Tensor, epsilon: float = 1e-7):
+    """
+    Triton implementation of Newton-Schulz iteration
+    """
+    # Newton-Schulz constants
+    ns_consts = [
+        (4.0848, -6.8946, 2.9270),
+        (3.9505, -6.3029, 2.6377),
+        (3.7418, -5.5913, 2.3037),
+        (2.8769, -3.1427, 1.2046),
+        (2.8366, -3.0525, 1.2012),
+    ]
+
+    X = G.to(dtype=torch.bfloat16)
+    if G.size(-2) > G.size(-1):
+        X = X.mT
+
+    # Ensure spectral norm is at most 1
+    X = X / (X.norm(dim=(-2, -1), keepdim=True) + epsilon)
+
+    # Allocate buffers
+    X = X.contiguous()
+    A = torch.empty((*X.shape[:-1], X.size(-2)), device=X.device, dtype=X.dtype)
+    B = torch.empty_like(A)
+    C = torch.empty_like(X)
+
+    ns_line_3 = torch.baddbmm if X.ndim > 2 else torch.addmm
+
+    # Perform the NS iterations
+    for a, b, c in ns_consts:
+        ns_line_1(X, out=A)  # A = X @ X.mT
+        ns_line_2(A, alpha=c, beta=b, out=B)  # B = b * A + c * A @ A
+        ns_line_3(X, B, X, beta=a, out=C)  # C = a * X + B @ X
+        X, C = C, X  # Swap references to avoid unnecessary copies
+
+    if G.size(-2) > G.size(-1):
+        X = X.mT
+    return X