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

CUDA graphs bake in memory addresses at capture time. When cpu_offload
frees optimizer state storage via resize_(0) between steps, the
CUDA graph replay accesses stale addresses, causing illegal memory access.

Remove triton.cudagraphs option and cudagraph_mark_step_begin() from
Newton-Schulz compile. torch.compile is preserved for kernel optimization.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

torch-ext/optimizer/newton_schulz.py

@@ -213,10 +213,8 @@ 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()
 
 
@@ -229,8 +227,6 @@ 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()