From 2a4b02fff9ee7c5edf128dfc9261d8296bbbd013 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Wed, 28 Oct 2026 16:30:24 -0700
Subject: [PATCH 1/7] Forward compiled stream recording to tensor storage

---
 sdm/tensor/categorical.py         |  4 ++
 sdm/tensor/columnar.py            |  5 +-
 sdm/tensor/mixin.py               | 12 ++++-
 sdm/tensor/nullable.py            |  4 ++
 sdm/tensor/table.py               |  5 +-
 sdm/tensor/var_len.py             |  4 ++
 test/tensor/test_record_stream.py | 80 +++++++++++++++++++++++++++++++
 7 files changed, 111 insertions(+), 3 deletions(-)

diff --git a/sdm/tensor/categorical.py b/sdm/tensor/categorical.py
index 2f7e56d249..c89cd51ab8 100644
--- a/sdm/tensor/categorical.py
+++ b/sdm/tensor/categorical.py
@@ -21,6 +21,7 @@
     to_arrow,
     to_cudf,
 )
+from sdm.tensor.mixin import _record_stream_op
 
 if TYPE_CHECKING:
     import cudf
@@ -463,6 +464,9 @@ def __torch_dispatch__(  # type: ignore
         if (handler := cls.HANDLED_FUNCTIONS.get(func)) is not None:
             return handler(*args, **(kwargs or {}))
 
+        if str(func) == "streams.record_stream.default":
+            return _record_stream_op(func, *args, **(kwargs or {}))
+
         # Operate on vanilla tensors for all non-handled functions:
         args = pytree.tree_map_only(CategoricalTensor, lambda x: x._code, args)
         kwargs = pytree.tree_map_only(
diff --git a/sdm/tensor/columnar.py b/sdm/tensor/columnar.py
index c1a18da813..3324fbae5d 100644
--- a/sdm/tensor/columnar.py
+++ b/sdm/tensor/columnar.py
@@ -16,7 +16,7 @@
 
 from sdm.tensor import NullableTensor, StringTensor, VarLenTensor
 from sdm.tensor.io import arrow_as_tensor, to_arrow, to_cudf
-from sdm.tensor.mixin import _resolve_device
+from sdm.tensor.mixin import _record_stream_op, _resolve_device
 
 if TYPE_CHECKING:
     import cudf
@@ -360,6 +360,9 @@ def __torch_dispatch__(  # type: ignore
         if (handler := cls.HANDLED_FUNCTIONS.get(func)) is not None:
             return handler(*args, **(kwargs or {}))
 
+        if str(func) == "streams.record_stream.default":
+            return _record_stream_op(func, *args, **(kwargs or {}))
+
         raise NotImplementedError(
             f"'{func}' is not supported for {cls.__name__!r}"
         )
diff --git a/sdm/tensor/mixin.py b/sdm/tensor/mixin.py
index 78f91cf6ec..1f9f258406 100644
--- a/sdm/tensor/mixin.py
+++ b/sdm/tensor/mixin.py
@@ -3,7 +3,7 @@
 
 import abc
 from collections.abc import Callable, Iterator
-from typing import Self
+from typing import Any, Self
 
 import torch
 from torch import Tensor
@@ -122,3 +122,13 @@ def _resolve_device(device: torch.device | str | None) -> torch.device | None:
     if device.type == "mps" and device.index is None:
         return torch.device("mps", 0)
     return device
+
+
+def _record_stream_op(
+    op: Callable[[Tensor, int], None], tensor: Any, stream_index: int
+) -> None:
+    # Dynamo represents streams as graph-owned indices. Forward the builtin
+    # operation to backing tensors without resolving the stream in Python.
+    attrs, _ = tensor.__tensor_flatten__()
+    for name in attrs:
+        op(getattr(tensor, name), stream_index)
diff --git a/sdm/tensor/nullable.py b/sdm/tensor/nullable.py
index fc51c8ec25..962bcc789b 100644
--- a/sdm/tensor/nullable.py
+++ b/sdm/tensor/nullable.py
@@ -20,6 +20,7 @@
     to_arrow,
     to_cudf,
 )
+from sdm.tensor.mixin import _record_stream_op
 
 if TYPE_CHECKING:
     import cudf
@@ -385,6 +386,9 @@ def __torch_dispatch__(  # type: ignore
         if (handler := cls.HANDLED_FUNCTIONS.get(func)) is not None:
             return handler(*args, **(kwargs or {}))
 
+        if str(func) == "streams.record_stream.default":
+            return _record_stream_op(func, *args, **(kwargs or {}))
+
         raise NotImplementedError(
             f"'{func}' is not supported for {cls.__name__!r}"
         )
diff --git a/sdm/tensor/table.py b/sdm/tensor/table.py
index 996fc8b0e6..081753e182 100644
--- a/sdm/tensor/table.py
+++ b/sdm/tensor/table.py
@@ -21,7 +21,7 @@
 from sdm import NaT, Stype, StypeLike
 from sdm.tensor import CategoricalTensor, ColumnarTensor, StringTensor
 from sdm.tensor.io import arrow_as_tensor, to_arrow, to_cudf
-from sdm.tensor.mixin import _resolve_device
+from sdm.tensor.mixin import _record_stream_op, _resolve_device
 
 if TYPE_CHECKING:
     import cudf
@@ -1201,6 +1201,9 @@ def __torch_dispatch__(  # type: ignore
         if func in cls.HANDLED_FUNCTIONS:
             return cls.HANDLED_FUNCTIONS[func](*args, **(kwargs or {}))
 
+        if str(func) == "streams.record_stream.default":
+            return _record_stream_op(func, *args, **(kwargs or {}))
+
         raise NotImplementedError(
             f"'{func}' is not supported for {cls.__name__!r}"
         )
diff --git a/sdm/tensor/var_len.py b/sdm/tensor/var_len.py
index a3c439b9c4..1e88aae00a 100644
--- a/sdm/tensor/var_len.py
+++ b/sdm/tensor/var_len.py
@@ -21,6 +21,7 @@
     combine_arrow_chunks,
     to_arrow,
 )
+from sdm.tensor.mixin import _record_stream_op
 
 aten = torch.ops.aten
 
@@ -628,6 +629,9 @@ def __torch_dispatch__(  # type: ignore
             with enable_reentrant_dispatch():  # Record autograd in `_data`.
                 return handler(*args, **(kwargs or {}))
 
+        if str(func) == "streams.record_stream.default":
+            return _record_stream_op(func, *args, **(kwargs or {}))
+
         raise NotImplementedError(
             f"'{func}' is not supported for {cls.__name__!r}"
         )
diff --git a/test/tensor/test_record_stream.py b/test/tensor/test_record_stream.py
index 4f93eeaaae..99982145c8 100644
--- a/test/tensor/test_record_stream.py
+++ b/test/tensor/test_record_stream.py
@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import pytest
 import torch
 
 from sdm import (
@@ -32,3 +33,82 @@ def test_record_stream() -> None:
 
     for tensor in tensors:
         assert tensor.record_stream(stream) is None
+
+
+def _stream_sources(kind: str) -> tuple[torch.Tensor, ...]:
+    data = torch.arange(4096, device="cuda")
+    valid = data.remainder(3).ne(0)
+    if kind == "categorical":
+        return data.remainder(4096).view(-1, 1), data
+    if kind in ("var_len", "string"):
+        offset = torch.arange(4097, device="cuda")
+        if kind == "string":
+            data = data.remainder(26).add(97).to(torch.uint8)
+        return data, offset, valid
+    return data, valid
+
+
+def _stream_tensor(
+    kind: str, sources: tuple[torch.Tensor, ...]
+) -> torch.Tensor:
+    if kind == "categorical":
+        return CategoricalTensor(sources[0], (sources[1],))
+    if kind == "columnar":
+        return ColumnarTensor((sources[0], NullableTensor(*sources)))
+    if kind == "nullable":
+        return NullableTensor(*sources)
+    if kind in ("var_len", "string"):
+        cls = StringTensor if kind == "string" else VarLenTensor
+        return cls(*sources, size=(4096,))
+    assert kind == "table"
+    return TableTensor(
+        numerical=NullableTensor(
+            sources[0].view(-1, 1), sources[1].view(-1, 1)
+        )
+    )
+
+
+def _record_stream_tensor(
+    kind: str,
+    sources: tuple[torch.Tensor, ...],
+    stream: torch.cuda.Stream,
+) -> None:
+    _stream_tensor(kind, sources).record_stream(stream)
+
+
+@onlyCUDA
+@pytest.mark.usefixtures("compilation")
+@pytest.mark.parametrize(
+    "kind",
+    ["categorical", "columnar", "nullable", "var_len", "string", "table"],
+)
+def test_record_stream_fullgraph_lifetime(kind: str) -> None:
+    if not hasattr(torch.ops.streams, "record_stream"):
+        pytest.skip(
+            "graph-owned stream recording requires streams.record_stream"
+        )
+    sources = _stream_sources(kind)
+    expected = tuple(source.cpu() for source in sources)
+    outputs = tuple(
+        torch.empty_like(value, pin_memory=True) for value in expected
+    )
+    specs = tuple((source.shape, source.dtype) for source in sources)
+    stream = torch.cuda.Stream()
+
+    compiled_record = torch.compile(_record_stream_tensor, fullgraph=True)
+    assert compiled_record(kind, sources, stream) is None
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        # Keep copies pending while the allocation stream reuses storage.
+        torch.cuda._sleep(50_000_000)
+        for output, source in zip(outputs, sources):
+            output.copy_(source, non_blocking=True)
+    del sources, source
+    replacements = [
+        torch.full(shape, 1, dtype=dtype, device="cuda")
+        for shape, dtype in specs
+    ]
+    stream.synchronize()
+    for output, value in zip(outputs, expected):
+        torch.testing.assert_close(output, value)
+    del replacements

From ce997f509300dcf81281c8842772489c4e51c587 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Wed, 28 Oct 2026 16:52:46 -0700
Subject: [PATCH 2/7] Simplify compiled stream lifetime test setup

---
 test/tensor/test_record_stream.py | 46 +++++++++++++------------------
 1 file changed, 19 insertions(+), 27 deletions(-)

diff --git a/test/tensor/test_record_stream.py b/test/tensor/test_record_stream.py
index 99982145c8..8b69d5e158 100644
--- a/test/tensor/test_record_stream.py
+++ b/test/tensor/test_record_stream.py
@@ -48,32 +48,28 @@ def _stream_sources(kind: str) -> tuple[torch.Tensor, ...]:
     return data, valid
 
 
-def _stream_tensor(
-    kind: str, sources: tuple[torch.Tensor, ...]
-) -> torch.Tensor:
-    if kind == "categorical":
-        return CategoricalTensor(sources[0], (sources[1],))
-    if kind == "columnar":
-        return ColumnarTensor((sources[0], NullableTensor(*sources)))
-    if kind == "nullable":
-        return NullableTensor(*sources)
-    if kind in ("var_len", "string"):
-        cls = StringTensor if kind == "string" else VarLenTensor
-        return cls(*sources, size=(4096,))
-    assert kind == "table"
-    return TableTensor(
-        numerical=NullableTensor(
-            sources[0].view(-1, 1), sources[1].view(-1, 1)
-        )
-    )
-
-
 def _record_stream_tensor(
     kind: str,
     sources: tuple[torch.Tensor, ...],
     stream: torch.cuda.Stream,
 ) -> None:
-    _stream_tensor(kind, sources).record_stream(stream)
+    if kind == "categorical":
+        tensor = CategoricalTensor(sources[0], (sources[1],))
+    elif kind == "columnar":
+        tensor = ColumnarTensor((sources[0], NullableTensor(*sources)))
+    elif kind == "nullable":
+        tensor = NullableTensor(*sources)
+    elif kind in ("var_len", "string"):
+        cls = StringTensor if kind == "string" else VarLenTensor
+        tensor = cls(*sources, size=(4096,))
+    else:
+        assert kind == "table"
+        tensor = TableTensor(
+            numerical=NullableTensor(
+                sources[0].view(-1, 1), sources[1].view(-1, 1)
+            )
+        )
+    tensor.record_stream(stream)
 
 
 @onlyCUDA
@@ -84,15 +80,12 @@ def _record_stream_tensor(
 )
 def test_record_stream_fullgraph_lifetime(kind: str) -> None:
     if not hasattr(torch.ops.streams, "record_stream"):
-        pytest.skip(
-            "graph-owned stream recording requires streams.record_stream"
-        )
+        pytest.skip("requires torch.ops.streams.record_stream")
     sources = _stream_sources(kind)
     expected = tuple(source.cpu() for source in sources)
     outputs = tuple(
         torch.empty_like(value, pin_memory=True) for value in expected
     )
-    specs = tuple((source.shape, source.dtype) for source in sources)
     stream = torch.cuda.Stream()
 
     compiled_record = torch.compile(_record_stream_tensor, fullgraph=True)
@@ -105,8 +98,7 @@ def test_record_stream_fullgraph_lifetime(kind: str) -> None:
             output.copy_(source, non_blocking=True)
     del sources, source
     replacements = [
-        torch.full(shape, 1, dtype=dtype, device="cuda")
-        for shape, dtype in specs
+        torch.full_like(value, 1, device="cuda") for value in expected
     ]
     stream.synchronize()
     for output, value in zip(outputs, expected):

From 79afba4cf0b24f33c3888e67652da9d8a806ad12 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Thu, 29 Oct 2026 09:42:24 -0700
Subject: [PATCH 3/7] Remove added compiled stream tests

---
 test/tensor/test_record_stream.py | 72 -------------------------------
 1 file changed, 72 deletions(-)

diff --git a/test/tensor/test_record_stream.py b/test/tensor/test_record_stream.py
index 8b69d5e158..4f93eeaaae 100644
--- a/test/tensor/test_record_stream.py
+++ b/test/tensor/test_record_stream.py
@@ -1,7 +1,6 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-import pytest
 import torch
 
 from sdm import (
@@ -33,74 +32,3 @@ def test_record_stream() -> None:
 
     for tensor in tensors:
         assert tensor.record_stream(stream) is None
-
-
-def _stream_sources(kind: str) -> tuple[torch.Tensor, ...]:
-    data = torch.arange(4096, device="cuda")
-    valid = data.remainder(3).ne(0)
-    if kind == "categorical":
-        return data.remainder(4096).view(-1, 1), data
-    if kind in ("var_len", "string"):
-        offset = torch.arange(4097, device="cuda")
-        if kind == "string":
-            data = data.remainder(26).add(97).to(torch.uint8)
-        return data, offset, valid
-    return data, valid
-
-
-def _record_stream_tensor(
-    kind: str,
-    sources: tuple[torch.Tensor, ...],
-    stream: torch.cuda.Stream,
-) -> None:
-    if kind == "categorical":
-        tensor = CategoricalTensor(sources[0], (sources[1],))
-    elif kind == "columnar":
-        tensor = ColumnarTensor((sources[0], NullableTensor(*sources)))
-    elif kind == "nullable":
-        tensor = NullableTensor(*sources)
-    elif kind in ("var_len", "string"):
-        cls = StringTensor if kind == "string" else VarLenTensor
-        tensor = cls(*sources, size=(4096,))
-    else:
-        assert kind == "table"
-        tensor = TableTensor(
-            numerical=NullableTensor(
-                sources[0].view(-1, 1), sources[1].view(-1, 1)
-            )
-        )
-    tensor.record_stream(stream)
-
-
-@onlyCUDA
-@pytest.mark.usefixtures("compilation")
-@pytest.mark.parametrize(
-    "kind",
-    ["categorical", "columnar", "nullable", "var_len", "string", "table"],
-)
-def test_record_stream_fullgraph_lifetime(kind: str) -> None:
-    if not hasattr(torch.ops.streams, "record_stream"):
-        pytest.skip("requires torch.ops.streams.record_stream")
-    sources = _stream_sources(kind)
-    expected = tuple(source.cpu() for source in sources)
-    outputs = tuple(
-        torch.empty_like(value, pin_memory=True) for value in expected
-    )
-    stream = torch.cuda.Stream()
-
-    compiled_record = torch.compile(_record_stream_tensor, fullgraph=True)
-    assert compiled_record(kind, sources, stream) is None
-    stream.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(stream):
-        # Keep copies pending while the allocation stream reuses storage.
-        torch.cuda._sleep(50_000_000)
-        for output, source in zip(outputs, sources):
-            output.copy_(source, non_blocking=True)
-    del sources, source
-    replacements = [
-        torch.full_like(value, 1, device="cuda") for value in expected
-    ]
-    stream.synchronize()
-    for output, value in zip(outputs, expected):
-        torch.testing.assert_close(output, value)
-    del replacements

From c22de11a6a8c4211601111af1017086f80ea9be8 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Thu, 29 Oct 2026 11:10:52 -0700
Subject: [PATCH 4/7] Clarify compiled stream recording helper

---
 sdm/tensor/categorical.py | 4 ++--
 sdm/tensor/columnar.py    | 4 ++--
 sdm/tensor/mixin.py       | 6 +++---
 sdm/tensor/nullable.py    | 4 ++--
 sdm/tensor/table.py       | 4 ++--
 sdm/tensor/var_len.py     | 4 ++--
 6 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/sdm/tensor/categorical.py b/sdm/tensor/categorical.py
index c89cd51ab8..58f29e7ffa 100644
--- a/sdm/tensor/categorical.py
+++ b/sdm/tensor/categorical.py
@@ -21,7 +21,7 @@
     to_arrow,
     to_cudf,
 )
-from sdm.tensor.mixin import _record_stream_op
+from sdm.tensor.mixin import _record_compiled_stream
 
 if TYPE_CHECKING:
     import cudf
@@ -465,7 +465,7 @@ def __torch_dispatch__(  # type: ignore
             return handler(*args, **(kwargs or {}))
 
         if str(func) == "streams.record_stream.default":
-            return _record_stream_op(func, *args, **(kwargs or {}))
+            return _record_compiled_stream(func, *args, **(kwargs or {}))
 
         # Operate on vanilla tensors for all non-handled functions:
         args = pytree.tree_map_only(CategoricalTensor, lambda x: x._code, args)
diff --git a/sdm/tensor/columnar.py b/sdm/tensor/columnar.py
index 3324fbae5d..5c3f13cb10 100644
--- a/sdm/tensor/columnar.py
+++ b/sdm/tensor/columnar.py
@@ -16,7 +16,7 @@
 
 from sdm.tensor import NullableTensor, StringTensor, VarLenTensor
 from sdm.tensor.io import arrow_as_tensor, to_arrow, to_cudf
-from sdm.tensor.mixin import _record_stream_op, _resolve_device
+from sdm.tensor.mixin import _record_compiled_stream, _resolve_device
 
 if TYPE_CHECKING:
     import cudf
@@ -361,7 +361,7 @@ def __torch_dispatch__(  # type: ignore
             return handler(*args, **(kwargs or {}))
 
         if str(func) == "streams.record_stream.default":
-            return _record_stream_op(func, *args, **(kwargs or {}))
+            return _record_compiled_stream(func, *args, **(kwargs or {}))
 
         raise NotImplementedError(
             f"'{func}' is not supported for {cls.__name__!r}"
diff --git a/sdm/tensor/mixin.py b/sdm/tensor/mixin.py
index 1f9f258406..fe5b7756d2 100644
--- a/sdm/tensor/mixin.py
+++ b/sdm/tensor/mixin.py
@@ -124,11 +124,11 @@ def _resolve_device(device: torch.device | str | None) -> torch.device | None:
     return device
 
 
-def _record_stream_op(
+def _record_compiled_stream(
     op: Callable[[Tensor, int], None], tensor: Any, stream_index: int
 ) -> None:
-    # Dynamo represents streams as graph-owned indices. Forward the builtin
-    # operation to backing tensors without resolving the stream in Python.
+    """Forward compiled stream recording to each backing tensor."""
+    # Keep Dynamo's stream index opaque instead of resolving it in Python.
     attrs, _ = tensor.__tensor_flatten__()
     for name in attrs:
         op(getattr(tensor, name), stream_index)
diff --git a/sdm/tensor/nullable.py b/sdm/tensor/nullable.py
index 962bcc789b..29db95b61c 100644
--- a/sdm/tensor/nullable.py
+++ b/sdm/tensor/nullable.py
@@ -20,7 +20,7 @@
     to_arrow,
     to_cudf,
 )
-from sdm.tensor.mixin import _record_stream_op
+from sdm.tensor.mixin import _record_compiled_stream
 
 if TYPE_CHECKING:
     import cudf
@@ -387,7 +387,7 @@ def __torch_dispatch__(  # type: ignore
             return handler(*args, **(kwargs or {}))
 
         if str(func) == "streams.record_stream.default":
-            return _record_stream_op(func, *args, **(kwargs or {}))
+            return _record_compiled_stream(func, *args, **(kwargs or {}))
 
         raise NotImplementedError(
             f"'{func}' is not supported for {cls.__name__!r}"
diff --git a/sdm/tensor/table.py b/sdm/tensor/table.py
index 081753e182..05d61c5240 100644
--- a/sdm/tensor/table.py
+++ b/sdm/tensor/table.py
@@ -21,7 +21,7 @@
 from sdm import NaT, Stype, StypeLike
 from sdm.tensor import CategoricalTensor, ColumnarTensor, StringTensor
 from sdm.tensor.io import arrow_as_tensor, to_arrow, to_cudf
-from sdm.tensor.mixin import _record_stream_op, _resolve_device
+from sdm.tensor.mixin import _record_compiled_stream, _resolve_device
 
 if TYPE_CHECKING:
     import cudf
@@ -1202,7 +1202,7 @@ def __torch_dispatch__(  # type: ignore
             return cls.HANDLED_FUNCTIONS[func](*args, **(kwargs or {}))
 
         if str(func) == "streams.record_stream.default":
-            return _record_stream_op(func, *args, **(kwargs or {}))
+            return _record_compiled_stream(func, *args, **(kwargs or {}))
 
         raise NotImplementedError(
             f"'{func}' is not supported for {cls.__name__!r}"
diff --git a/sdm/tensor/var_len.py b/sdm/tensor/var_len.py
index 2ca32a8396..832eeef0bd 100644
--- a/sdm/tensor/var_len.py
+++ b/sdm/tensor/var_len.py
@@ -21,7 +21,7 @@
     combine_arrow_chunks,
     to_arrow,
 )
-from sdm.tensor.mixin import _record_stream_op
+from sdm.tensor.mixin import _record_compiled_stream
 
 aten = torch.ops.aten
 
@@ -634,7 +634,7 @@ def __torch_dispatch__(  # type: ignore
                 return handler(*args, **(kwargs or {}))
 
         if str(func) == "streams.record_stream.default":
-            return _record_stream_op(func, *args, **(kwargs or {}))
+            return _record_compiled_stream(func, *args, **(kwargs or {}))
 
         raise NotImplementedError(
             f"'{func}' is not supported for {cls.__name__!r}"

From 30c13a4caec8b3a809b8c2b1a68b643ecf8057dc Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Thu, 29 Oct 2026 11:18:02 -0700
Subject: [PATCH 5/7] Call compiled stream recording directly on stored tensors

---
 sdm/tensor/categorical.py |  2 +-
 sdm/tensor/columnar.py    |  2 +-
 sdm/tensor/mixin.py       | 14 ++++++--------
 sdm/tensor/nullable.py    |  2 +-
 sdm/tensor/table.py       |  2 +-
 sdm/tensor/var_len.py     |  2 +-
 6 files changed, 11 insertions(+), 13 deletions(-)

diff --git a/sdm/tensor/categorical.py b/sdm/tensor/categorical.py
index 58f29e7ffa..09c409a2f5 100644
--- a/sdm/tensor/categorical.py
+++ b/sdm/tensor/categorical.py
@@ -465,7 +465,7 @@ def __torch_dispatch__(  # type: ignore
             return handler(*args, **(kwargs or {}))
 
         if str(func) == "streams.record_stream.default":
-            return _record_compiled_stream(func, *args, **(kwargs or {}))
+            return _record_compiled_stream(*args, **(kwargs or {}))
 
         # Operate on vanilla tensors for all non-handled functions:
         args = pytree.tree_map_only(CategoricalTensor, lambda x: x._code, args)
diff --git a/sdm/tensor/columnar.py b/sdm/tensor/columnar.py
index 5c3f13cb10..9653b1873e 100644
--- a/sdm/tensor/columnar.py
+++ b/sdm/tensor/columnar.py
@@ -361,7 +361,7 @@ def __torch_dispatch__(  # type: ignore
             return handler(*args, **(kwargs or {}))
 
         if str(func) == "streams.record_stream.default":
-            return _record_compiled_stream(func, *args, **(kwargs or {}))
+            return _record_compiled_stream(*args, **(kwargs or {}))
 
         raise NotImplementedError(
             f"'{func}' is not supported for {cls.__name__!r}"
diff --git a/sdm/tensor/mixin.py b/sdm/tensor/mixin.py
index fe5b7756d2..5d213f34af 100644
--- a/sdm/tensor/mixin.py
+++ b/sdm/tensor/mixin.py
@@ -124,11 +124,9 @@ def _resolve_device(device: torch.device | str | None) -> torch.device | None:
     return device
 
 
-def _record_compiled_stream(
-    op: Callable[[Tensor, int], None], tensor: Any, stream_index: int
-) -> None:
-    """Forward compiled stream recording to each backing tensor."""
-    # Keep Dynamo's stream index opaque instead of resolving it in Python.
-    attrs, _ = tensor.__tensor_flatten__()
-    for name in attrs:
-        op(getattr(tensor, name), stream_index)
+def _record_compiled_stream(tensor: Any, stream_index: int) -> None:
+    """Keep each stored tensor's memory alive while the stream uses it."""
+    tensor_names, _ = tensor.__tensor_flatten__()
+    for tensor_name in tensor_names:
+        stored_tensor = getattr(tensor, tensor_name)
+        torch.ops.streams.record_stream.default(stored_tensor, stream_index)
diff --git a/sdm/tensor/nullable.py b/sdm/tensor/nullable.py
index 29db95b61c..2ed562ba18 100644
--- a/sdm/tensor/nullable.py
+++ b/sdm/tensor/nullable.py
@@ -387,7 +387,7 @@ def __torch_dispatch__(  # type: ignore
             return handler(*args, **(kwargs or {}))
 
         if str(func) == "streams.record_stream.default":
-            return _record_compiled_stream(func, *args, **(kwargs or {}))
+            return _record_compiled_stream(*args, **(kwargs or {}))
 
         raise NotImplementedError(
             f"'{func}' is not supported for {cls.__name__!r}"
diff --git a/sdm/tensor/table.py b/sdm/tensor/table.py
index 05d61c5240..37344a4ce4 100644
--- a/sdm/tensor/table.py
+++ b/sdm/tensor/table.py
@@ -1202,7 +1202,7 @@ def __torch_dispatch__(  # type: ignore
             return cls.HANDLED_FUNCTIONS[func](*args, **(kwargs or {}))
 
         if str(func) == "streams.record_stream.default":
-            return _record_compiled_stream(func, *args, **(kwargs or {}))
+            return _record_compiled_stream(*args, **(kwargs or {}))
 
         raise NotImplementedError(
             f"'{func}' is not supported for {cls.__name__!r}"
diff --git a/sdm/tensor/var_len.py b/sdm/tensor/var_len.py
index 832eeef0bd..a31438a978 100644
--- a/sdm/tensor/var_len.py
+++ b/sdm/tensor/var_len.py
@@ -634,7 +634,7 @@ def __torch_dispatch__(  # type: ignore
                 return handler(*args, **(kwargs or {}))
 
         if str(func) == "streams.record_stream.default":
-            return _record_compiled_stream(func, *args, **(kwargs or {}))
+            return _record_compiled_stream(*args, **(kwargs or {}))
 
         raise NotImplementedError(
             f"'{func}' is not supported for {cls.__name__!r}"

From 0a4ea05ecb7a59fc90e29b0dc7c6ba05bfd2be89 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Thu, 29 Oct 2026 11:33:43 -0700
Subject: [PATCH 6/7] Document compiler version required for stream recording

---
 sdm/tensor/mixin.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/sdm/tensor/mixin.py b/sdm/tensor/mixin.py
index 5d213f34af..ed48bf2f48 100644
--- a/sdm/tensor/mixin.py
+++ b/sdm/tensor/mixin.py
@@ -125,7 +125,11 @@ def _resolve_device(device: torch.device | str | None) -> torch.device | None:
 
 
 def _record_compiled_stream(tensor: Any, stream_index: int) -> None:
-    """Keep each stored tensor's memory alive while the stream uses it."""
+    """Keep each stored tensor's memory alive while the stream uses it.
+
+    Capturing ``record_stream()`` inside a compiled graph requires PyTorch 2.12
+    or later.
+    """
     tensor_names, _ = tensor.__tensor_flatten__()
     for tensor_name in tensor_names:
         stored_tensor = getattr(tensor, tensor_name)

From 4fde84631683b7710cd2b6c874c69349b21cd308 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Thu, 29 Oct 2026 11:38:16 -0700
Subject: [PATCH 7/7] Remove compiled stream version note

---
 sdm/tensor/mixin.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/sdm/tensor/mixin.py b/sdm/tensor/mixin.py
index ed48bf2f48..5d213f34af 100644
--- a/sdm/tensor/mixin.py
+++ b/sdm/tensor/mixin.py
@@ -125,11 +125,7 @@ def _resolve_device(device: torch.device | str | None) -> torch.device | None:
 
 
 def _record_compiled_stream(tensor: Any, stream_index: int) -> None:
-    """Keep each stored tensor's memory alive while the stream uses it.
-
-    Capturing ``record_stream()`` inside a compiled graph requires PyTorch 2.12
-    or later.
-    """
+    """Keep each stored tensor's memory alive while the stream uses it."""
     tensor_names, _ = tensor.__tensor_flatten__()
     for tensor_name in tensor_names:
         stored_tensor = getattr(tensor, tensor_name)