Revert "fix: disable CUDA graphs in Newton-Schulz for cpu_offload compatibility" (#29)

* Revert "fix: disable CUDA graphs in Newton-Schulz for cpu_offload compatibility"

* Add built binary [skip-build]

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

build/torch210-cxx11-cu126-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_2dce952_dirty
-ops = torch.ops._optimizer_2dce952_dirty
+from . import _optimizer_38f9b8e_dirty
+ops = torch.ops._optimizer_38f9b8e_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_2dce952_dirty::{op_name}"
+    return f"_optimizer_38f9b8e_dirty::{op_name}"

build/torch210-cxx11-cu126-x86_64-linux/{_optimizer_2dce952_dirty.abi3.so → _optimizer_38f9b8e_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a45eb1ca39a9122a7ec8171ecc1ed8d46b0bd6ec6fcf66e1f5635f2ba8b946ce
+oid sha256:cb6163428ce86500d61c2b765eecd7eb6f31c092066278e1d1af7a0848dc5126
 size 1940944

build/torch210-cxx11-cu126-x86_64-linux/newton_schulz.py

@@ -213,8 +213,10 @@ def zeropower_via_newtonschulz5(G, steps=5):
     if key not in _ns_per_shape:
         _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5,
                                            options={
+                                               "triton.cudagraphs": True,
                                                "shape_padding": False
                                            })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()
 
 
@@ -227,6 +229,8 @@ def zeropower_via_newtonschulz5_batched(G, steps=5):
         _ns_per_shape[key] = torch.compile(
             _zeropower_via_newtonschulz5_batched,
             options={
+                "triton.cudagraphs": True,
                 "shape_padding": False
             })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()

build/torch210-cxx11-cu128-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_2dce952_dirty
-ops = torch.ops._optimizer_2dce952_dirty
+from . import _optimizer_38f9b8e_dirty
+ops = torch.ops._optimizer_38f9b8e_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_2dce952_dirty::{op_name}"
+    return f"_optimizer_38f9b8e_dirty::{op_name}"

build/torch210-cxx11-cu128-x86_64-linux/{_optimizer_2dce952_dirty.abi3.so → _optimizer_38f9b8e_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:63fcd4e511d075431c2e7e2fefbb96dcced540e2c0ca1316faf29fdfd78e253a
+oid sha256:017323d479e8fbd3ed1f550f95fc4ba9f2e304dbe9351c0eaa75543ebe775e18
 size 2004144

build/torch210-cxx11-cu128-x86_64-linux/newton_schulz.py

@@ -213,8 +213,10 @@ def zeropower_via_newtonschulz5(G, steps=5):
     if key not in _ns_per_shape:
         _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5,
                                            options={
+                                               "triton.cudagraphs": True,
                                                "shape_padding": False
                                            })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()
 
 
@@ -227,6 +229,8 @@ def zeropower_via_newtonschulz5_batched(G, steps=5):
         _ns_per_shape[key] = torch.compile(
             _zeropower_via_newtonschulz5_batched,
             options={
+                "triton.cudagraphs": True,
                 "shape_padding": False
             })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()

build/torch210-cxx11-cu130-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_2dce952_dirty
-ops = torch.ops._optimizer_2dce952_dirty
+from . import _optimizer_38f9b8e_dirty
+ops = torch.ops._optimizer_38f9b8e_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_2dce952_dirty::{op_name}"
+    return f"_optimizer_38f9b8e_dirty::{op_name}"

build/torch210-cxx11-cu130-x86_64-linux/{_optimizer_2dce952_dirty.abi3.so → _optimizer_38f9b8e_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:62348f4dd6a65cfc98cf101dff522ad8ae88c56ee99fb7cfed2c33b3222262d8
+oid sha256:783a161f2d28e4244226c9d6e59ac33f74f7a79aad17c06e8ce027dd6182e03c
 size 2004728

build/torch210-cxx11-cu130-x86_64-linux/newton_schulz.py

@@ -213,8 +213,10 @@ def zeropower_via_newtonschulz5(G, steps=5):
     if key not in _ns_per_shape:
         _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5,
                                            options={
+                                               "triton.cudagraphs": True,
                                                "shape_padding": False
                                            })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()
 
 
@@ -227,6 +229,8 @@ def zeropower_via_newtonschulz5_batched(G, steps=5):
         _ns_per_shape[key] = torch.compile(
             _zeropower_via_newtonschulz5_batched,
             options={
+                "triton.cudagraphs": True,
                 "shape_padding": False
             })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()

build/torch210-cxx11-rocm70-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_2dce952_dirty
-ops = torch.ops._optimizer_2dce952_dirty
+from . import _optimizer_38f9b8e_dirty
+ops = torch.ops._optimizer_38f9b8e_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_2dce952_dirty::{op_name}"
+    return f"_optimizer_38f9b8e_dirty::{op_name}"

build/torch210-cxx11-rocm70-x86_64-linux/{_optimizer_2dce952_dirty.abi3.so → _optimizer_38f9b8e_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c9d1940e448de8c2290bb3142a1d41433731dba06c6d68738f974af26cb4f240
+oid sha256:8ec2fcc8a9dc8a1e4aa4e925eaee33613177873e474e8d627bf844dae80f5f8b
 size 1866400

build/torch210-cxx11-rocm70-x86_64-linux/newton_schulz.py

@@ -213,8 +213,10 @@ def zeropower_via_newtonschulz5(G, steps=5):
     if key not in _ns_per_shape:
         _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5,
                                            options={
+                                               "triton.cudagraphs": True,
                                                "shape_padding": False
                                            })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()
 
 
@@ -227,6 +229,8 @@ def zeropower_via_newtonschulz5_batched(G, steps=5):
         _ns_per_shape[key] = torch.compile(
             _zeropower_via_newtonschulz5_batched,
             options={
+                "triton.cudagraphs": True,
                 "shape_padding": False
             })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()

build/torch210-cxx11-rocm71-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_2dce952_dirty
-ops = torch.ops._optimizer_2dce952_dirty
+from . import _optimizer_38f9b8e_dirty
+ops = torch.ops._optimizer_38f9b8e_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_2dce952_dirty::{op_name}"
+    return f"_optimizer_38f9b8e_dirty::{op_name}"

build/torch210-cxx11-rocm71-x86_64-linux/{_optimizer_2dce952_dirty.abi3.so → _optimizer_38f9b8e_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1eb0a09bd1b20303563fe0cb3e8e2234dae40241e3d0d56502fc5a6ae442588f
+oid sha256:18373b2e448071735ce724008122f179dd814986925c9cf0fc03f32201b2b1fa
 size 1866112

build/torch210-cxx11-rocm71-x86_64-linux/newton_schulz.py

@@ -213,8 +213,10 @@ def zeropower_via_newtonschulz5(G, steps=5):
     if key not in _ns_per_shape:
         _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5,
                                            options={
+                                               "triton.cudagraphs": True,
                                                "shape_padding": False
                                            })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()
 
 
@@ -227,6 +229,8 @@ def zeropower_via_newtonschulz5_batched(G, steps=5):
         _ns_per_shape[key] = torch.compile(
             _zeropower_via_newtonschulz5_batched,
             options={
+                "triton.cudagraphs": True,
                 "shape_padding": False
             })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()

build/torch28-cxx11-cu126-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_2dce952_dirty
-ops = torch.ops._optimizer_2dce952_dirty
+from . import _optimizer_38f9b8e_dirty
+ops = torch.ops._optimizer_38f9b8e_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_2dce952_dirty::{op_name}"
+    return f"_optimizer_38f9b8e_dirty::{op_name}"

build/torch28-cxx11-cu126-x86_64-linux/{_optimizer_2dce952_dirty.abi3.so → _optimizer_38f9b8e_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:14140b8568b7a04120bf30a13f95a5dfa44255d7d107c0330b61ce92e79404d0
+oid sha256:d2db9c7fb764a1fae1872779bc9ffac2aff18d14a238111d6b8b53b7d3dfa0d3
 size 1936664

build/torch28-cxx11-cu126-x86_64-linux/newton_schulz.py

@@ -213,8 +213,10 @@ def zeropower_via_newtonschulz5(G, steps=5):
     if key not in _ns_per_shape:
         _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5,
                                            options={
+                                               "triton.cudagraphs": True,
                                                "shape_padding": False
                                            })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()
 
 
@@ -227,6 +229,8 @@ def zeropower_via_newtonschulz5_batched(G, steps=5):
         _ns_per_shape[key] = torch.compile(
             _zeropower_via_newtonschulz5_batched,
             options={
+                "triton.cudagraphs": True,
                 "shape_padding": False
             })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()

build/torch28-cxx11-cu128-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_2dce952_dirty
-ops = torch.ops._optimizer_2dce952_dirty
+from . import _optimizer_38f9b8e_dirty
+ops = torch.ops._optimizer_38f9b8e_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_2dce952_dirty::{op_name}"
+    return f"_optimizer_38f9b8e_dirty::{op_name}"

build/torch28-cxx11-cu128-x86_64-linux/{_optimizer_2dce952_dirty.abi3.so → _optimizer_38f9b8e_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b8a8e75d22105a30822529b34963a13a0a60a20a862bc7da971857f54008105c
+oid sha256:b1ddfe7e38a9213d5dede8052c81b78eca952aef122d4da919950ff504dc3908
 size 1999872

build/torch28-cxx11-cu128-x86_64-linux/newton_schulz.py

@@ -213,8 +213,10 @@ def zeropower_via_newtonschulz5(G, steps=5):
     if key not in _ns_per_shape:
         _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5,
                                            options={
+                                               "triton.cudagraphs": True,
                                                "shape_padding": False
                                            })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()
 
 
@@ -227,6 +229,8 @@ def zeropower_via_newtonschulz5_batched(G, steps=5):
         _ns_per_shape[key] = torch.compile(
             _zeropower_via_newtonschulz5_batched,
             options={
+                "triton.cudagraphs": True,
                 "shape_padding": False
             })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()

build/torch28-cxx11-cu129-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_2dce952_dirty
-ops = torch.ops._optimizer_2dce952_dirty
+from . import _optimizer_38f9b8e_dirty
+ops = torch.ops._optimizer_38f9b8e_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_2dce952_dirty::{op_name}"
+    return f"_optimizer_38f9b8e_dirty::{op_name}"

build/torch28-cxx11-cu129-x86_64-linux/{_optimizer_2dce952_dirty.abi3.so → _optimizer_38f9b8e_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:69697c5e77ccd05e26a3eca114732d905f9c3b8c68d335b13f129c56a710c7a3
+oid sha256:db68ba26f1b022f56a5ab4e6e0204bf26df8922750f32f21be0ad76e2674b717
 size 1999872

build/torch28-cxx11-cu129-x86_64-linux/newton_schulz.py

@@ -213,8 +213,10 @@ def zeropower_via_newtonschulz5(G, steps=5):
     if key not in _ns_per_shape:
         _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5,
                                            options={
+                                               "triton.cudagraphs": True,
                                                "shape_padding": False
                                            })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()
 
 
@@ -227,6 +229,8 @@ def zeropower_via_newtonschulz5_batched(G, steps=5):
         _ns_per_shape[key] = torch.compile(
             _zeropower_via_newtonschulz5_batched,
             options={
+                "triton.cudagraphs": True,
                 "shape_padding": False
             })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()

build/torch28-cxx11-rocm63-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_2dce952_dirty
-ops = torch.ops._optimizer_2dce952_dirty
+from . import _optimizer_38f9b8e_dirty
+ops = torch.ops._optimizer_38f9b8e_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_2dce952_dirty::{op_name}"
+    return f"_optimizer_38f9b8e_dirty::{op_name}"

build/torch28-cxx11-rocm63-x86_64-linux/{_optimizer_2dce952_dirty.abi3.so → _optimizer_38f9b8e_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:01a9dda5ac371798b9fd86b77b8cf7b572721a423ee69971910abe69a62d88fd
+oid sha256:5a6a8788f055b22d594330fc06487ae2c6eeb2b64e0ab0132b68036a78560cf6
 size 1865080

build/torch28-cxx11-rocm63-x86_64-linux/newton_schulz.py

@@ -213,8 +213,10 @@ def zeropower_via_newtonschulz5(G, steps=5):
     if key not in _ns_per_shape:
         _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5,
                                            options={
+                                               "triton.cudagraphs": True,
                                                "shape_padding": False
                                            })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()
 
 
@@ -227,6 +229,8 @@ def zeropower_via_newtonschulz5_batched(G, steps=5):
         _ns_per_shape[key] = torch.compile(
             _zeropower_via_newtonschulz5_batched,
             options={
+                "triton.cudagraphs": True,
                 "shape_padding": False
             })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()

build/torch28-cxx11-rocm64-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_2dce952_dirty
-ops = torch.ops._optimizer_2dce952_dirty
+from . import _optimizer_38f9b8e_dirty
+ops = torch.ops._optimizer_38f9b8e_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_2dce952_dirty::{op_name}"
+    return f"_optimizer_38f9b8e_dirty::{op_name}"

build/torch28-cxx11-rocm64-x86_64-linux/{_optimizer_2dce952_dirty.abi3.so → _optimizer_38f9b8e_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c1fb5220ced504ce27281ceb69fc21ce69de51baad5485a67973bb6be489aeb9
+oid sha256:6e5e17fd042010ec06456f5885603c4e38476981d43adb1cc99ea6dbe5f57c6f
 size 1865168

build/torch28-cxx11-rocm64-x86_64-linux/newton_schulz.py

@@ -213,8 +213,10 @@ def zeropower_via_newtonschulz5(G, steps=5):
     if key not in _ns_per_shape:
         _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5,
                                            options={
+                                               "triton.cudagraphs": True,
                                                "shape_padding": False
                                            })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()
 
 
@@ -227,6 +229,8 @@ def zeropower_via_newtonschulz5_batched(G, steps=5):
         _ns_per_shape[key] = torch.compile(
             _zeropower_via_newtonschulz5_batched,
             options={
+                "triton.cudagraphs": True,
                 "shape_padding": False
             })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()

build/torch29-cxx11-cu126-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_2dce952_dirty
-ops = torch.ops._optimizer_2dce952_dirty
+from . import _optimizer_38f9b8e_dirty
+ops = torch.ops._optimizer_38f9b8e_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_2dce952_dirty::{op_name}"
+    return f"_optimizer_38f9b8e_dirty::{op_name}"

build/torch29-cxx11-cu126-x86_64-linux/{_optimizer_2dce952_dirty.abi3.so → _optimizer_38f9b8e_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:930e785c6e2e19242fca7a43c1384693552777d4e0d07ba52558d1212bc158fd
+oid sha256:f6b4b64b9e80383e0e1e9d5482f39dc28256bf901211d60deda16d905f198e70
 size 1936664

build/torch29-cxx11-cu126-x86_64-linux/newton_schulz.py

@@ -213,8 +213,10 @@ def zeropower_via_newtonschulz5(G, steps=5):
     if key not in _ns_per_shape:
         _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5,
                                            options={
+                                               "triton.cudagraphs": True,
                                                "shape_padding": False
                                            })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()
 
 
@@ -227,6 +229,8 @@ def zeropower_via_newtonschulz5_batched(G, steps=5):
         _ns_per_shape[key] = torch.compile(
             _zeropower_via_newtonschulz5_batched,
             options={
+                "triton.cudagraphs": True,
                 "shape_padding": False
             })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()

build/torch29-cxx11-cu128-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_2dce952_dirty
-ops = torch.ops._optimizer_2dce952_dirty
+from . import _optimizer_38f9b8e_dirty
+ops = torch.ops._optimizer_38f9b8e_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_2dce952_dirty::{op_name}"
+    return f"_optimizer_38f9b8e_dirty::{op_name}"

build/torch29-cxx11-cu128-x86_64-linux/{_optimizer_2dce952_dirty.abi3.so → _optimizer_38f9b8e_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b49a73f9b90d9962a9bf859d85678dfd4f042762a561287098c30af2786779fd
+oid sha256:52bc357d5a5d09094e142d1bf87e9a2ba819a6770b8b34ab9469ceb2414ad29e
 size 1999872

build/torch29-cxx11-cu128-x86_64-linux/newton_schulz.py

@@ -213,8 +213,10 @@ def zeropower_via_newtonschulz5(G, steps=5):
     if key not in _ns_per_shape:
         _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5,
                                            options={
+                                               "triton.cudagraphs": True,
                                                "shape_padding": False
                                            })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()
 
 
@@ -227,6 +229,8 @@ def zeropower_via_newtonschulz5_batched(G, steps=5):
         _ns_per_shape[key] = torch.compile(
             _zeropower_via_newtonschulz5_batched,
             options={
+                "triton.cudagraphs": True,
                 "shape_padding": False
             })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()

build/torch29-cxx11-cu130-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_2dce952_dirty
-ops = torch.ops._optimizer_2dce952_dirty
+from . import _optimizer_38f9b8e_dirty
+ops = torch.ops._optimizer_38f9b8e_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_2dce952_dirty::{op_name}"
+    return f"_optimizer_38f9b8e_dirty::{op_name}"

build/torch29-cxx11-cu130-x86_64-linux/{_optimizer_2dce952_dirty.abi3.so → _optimizer_38f9b8e_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:818efe7b03956810294c938e701163f7c06285e72631a5fb78b4f82eba72d8a6
+oid sha256:faa0fc60fea48f7ba85933fb35d3c96afb6f5b357d4e28565e8112848fed890a
 size 2000456

build/torch29-cxx11-cu130-x86_64-linux/newton_schulz.py

@@ -213,8 +213,10 @@ def zeropower_via_newtonschulz5(G, steps=5):
     if key not in _ns_per_shape:
         _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5,
                                            options={
+                                               "triton.cudagraphs": True,
                                                "shape_padding": False
                                            })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()
 
 
@@ -227,6 +229,8 @@ def zeropower_via_newtonschulz5_batched(G, steps=5):
         _ns_per_shape[key] = torch.compile(
             _zeropower_via_newtonschulz5_batched,
             options={
+                "triton.cudagraphs": True,
                 "shape_padding": False
             })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()

build/torch29-cxx11-rocm63-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_2dce952_dirty
-ops = torch.ops._optimizer_2dce952_dirty
+from . import _optimizer_38f9b8e_dirty
+ops = torch.ops._optimizer_38f9b8e_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_2dce952_dirty::{op_name}"
+    return f"_optimizer_38f9b8e_dirty::{op_name}"

build/torch29-cxx11-rocm63-x86_64-linux/{_optimizer_2dce952_dirty.abi3.so → _optimizer_38f9b8e_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:51e9488c2003b603f0c9f09c03e7ddfb576074f2fdc6d94363e7dfd07f5a7111
+oid sha256:d8cf990229d1c3dc8378e74487af28ab48ffb91ec12869ddb0839d3b4cddc03e
 size 1865112

build/torch29-cxx11-rocm63-x86_64-linux/newton_schulz.py

@@ -213,8 +213,10 @@ def zeropower_via_newtonschulz5(G, steps=5):
     if key not in _ns_per_shape:
         _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5,
                                            options={
+                                               "triton.cudagraphs": True,
                                                "shape_padding": False
                                            })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()
 
 
@@ -227,6 +229,8 @@ def zeropower_via_newtonschulz5_batched(G, steps=5):
         _ns_per_shape[key] = torch.compile(
             _zeropower_via_newtonschulz5_batched,
             options={
+                "triton.cudagraphs": True,
                 "shape_padding": False
             })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()

build/torch29-cxx11-rocm64-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_2dce952_dirty
-ops = torch.ops._optimizer_2dce952_dirty
+from . import _optimizer_38f9b8e_dirty
+ops = torch.ops._optimizer_38f9b8e_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_2dce952_dirty::{op_name}"
+    return f"_optimizer_38f9b8e_dirty::{op_name}"

build/torch29-cxx11-rocm64-x86_64-linux/{_optimizer_2dce952_dirty.abi3.so → _optimizer_38f9b8e_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:887bca38f051d28d35af84921d78943d2eff7a3d86a48c0f29f0e908ae577d0c
+oid sha256:fbdc2be035c0380bdd9ea10a0f913ecf5b6be29d4d7d74e1bd4056143393f28d
 size 1865232

build/torch29-cxx11-rocm64-x86_64-linux/newton_schulz.py

@@ -213,8 +213,10 @@ def zeropower_via_newtonschulz5(G, steps=5):
     if key not in _ns_per_shape:
         _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5,
                                            options={
+                                               "triton.cudagraphs": True,
                                                "shape_padding": False
                                            })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()
 
 
@@ -227,6 +229,8 @@ def zeropower_via_newtonschulz5_batched(G, steps=5):
         _ns_per_shape[key] = torch.compile(
             _zeropower_via_newtonschulz5_batched,
             options={
+                "triton.cudagraphs": True,
                 "shape_padding": False
             })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()

torch-ext/optimizer/newton_schulz.py

@@ -213,8 +213,10 @@ def zeropower_via_newtonschulz5(G, steps=5):
     if key not in _ns_per_shape:
         _ns_per_shape[key] = torch.compile(_zeropower_via_newtonschulz5,
                                            options={
+                                               "triton.cudagraphs": True,
                                                "shape_padding": False
                                            })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()
 
 
@@ -227,6 +229,8 @@ def zeropower_via_newtonschulz5_batched(G, steps=5):
         _ns_per_shape[key] = torch.compile(
             _zeropower_via_newtonschulz5_batched,
             options={
+                "triton.cudagraphs": True,
                 "shape_padding": False
             })
+    torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()