From 406cc9ec1414be1c6267b1ae2e4cadb8fac2f000 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Mon, 26 Oct 2026 13:43:34 -0700
Subject: [PATCH 1/4] Add native stable CUDA sorting for bounded UTF-8 strings

---
 sdm/tensor/string.py       |  30 ++++++++++
 test/tensor/test_string.py | 115 +++++++++++++++++++++++++++++++++++++
 2 files changed, 145 insertions(+)

diff --git a/sdm/tensor/string.py b/sdm/tensor/string.py
index 17b8da35f6..8123e69691 100644
--- a/sdm/tensor/string.py
+++ b/sdm/tensor/string.py
@@ -16,7 +16,9 @@
 
 from sdm._warnings import warn_once
 from sdm.tensor import VarLenTensor
+from sdm.tensor._string_keys import pack
 from sdm.tensor.io import arrow_as_tensor, combine_arrow_chunks
+from sdm.tensor.var_len import _select_entries
 
 if TYPE_CHECKING:
     import cudf
@@ -523,6 +525,10 @@ def _sort(
     if inp.dim() != 1:
         raise NotImplementedError("'sort' only supports one-dimensional input")
 
+    # Bound radix passes; longer strings retain the external sort.
+    if inp.is_cuda and inp._max_length <= 56:
+        return _sort_native(inp, descending=descending)
+
     backend: Literal["arrow", "cudf"] = "arrow"
     if inp.is_cuda:
         if importlib.util.find_spec("cudf") is not None:
@@ -553,3 +559,27 @@ def _sort(
             perm = torch.from_dlpack(perm_ser.astype("int64").to_cupy())
 
     return cast(StringTensor, inp[perm]), perm
+
+
+def _sort_native(
+    inp: StringTensor,
+    *,
+    descending: bool,
+) -> tuple[StringTensor, Tensor]:
+    words, valid = pack(inp)
+    if inp._valid is not None:
+        # Nulls sort last in both directions, including stable duplicate nulls.
+        sentinel = -1 if descending else torch.iinfo(torch.int64).max
+        words = torch.where(valid[:, None], words, sentinel)
+
+    # Least-significant word first preserves lexicographic ordering and ties.
+    perm = words[:, -1].argsort(descending=descending, stable=True)
+    for column in range(words.size(1) - 2, -1, -1):
+        order = words[perm, column].argsort(
+            descending=descending,
+            stable=True,
+        )
+        perm = perm[order]
+
+    out = _select_entries(inp, perm)
+    return cast(StringTensor, out), perm
diff --git a/test/tensor/test_string.py b/test/tensor/test_string.py
index d70c6a3882..662bf8d41b 100644
--- a/test/tensor/test_string.py
+++ b/test/tensor/test_string.py
@@ -8,6 +8,7 @@
 import torch
 
 from sdm import StringTensor
+from sdm.tensor.string import _sort_native
 from sdm.testing import onlyCUDA, withCUDA
 
 
@@ -382,3 +383,117 @@ def test_registers_as_a_pytorch_buffer() -> None:
     module.register_buffer("value", tensor)
 
     assert tuple(module.buffers()) == (tensor,)
+
+
+@pytest.mark.parametrize("descending", [False, True])
+@pytest.mark.parametrize(
+    "values",
+    [
+        [],
+        [None, None],
+        ["", "", None, "\0", "\0a", "a", "a", None],
+        ["é", "e", "中", "😀", "é", "\x7f", "\x80"],
+        ["a" * 7 + "b", "a" * 8, "a" * 7, "a" * 14, None],
+        ["a" * 55 + "b", "a" * 56, "a" * 55, "a" * 56, None],
+        ["a" * 57 + "z", "a" * 58, "a" * 57, "a" * 58, None, None],
+    ],
+)
+@withCUDA
+def test_sort_stable(
+    values: list[str | None],
+    descending: bool,
+    device: torch.device,
+) -> None:
+    # The sort indices are relative to the view, including a storage offset.
+    storage: list[str | None] = ["unused"]
+    for value in values:
+        storage.extend([value, "unused"])
+    tensor = StringTensor.from_list(storage, device=device)[1::2]
+    indices = sorted(
+        (i for i, value in enumerate(values) if value is not None),
+        key=lambda i: cast(str, values[i]).encode("utf-8"),
+        reverse=descending,
+    )
+    indices.extend(i for i, value in enumerate(values) if value is None)
+
+    out, perm = torch.sort(tensor, descending=descending, stable=True)
+    assert out.tolist() == [values[i] for i in indices]
+    assert perm.tolist() == indices
+    assert out.device == device
+
+
+@pytest.mark.parametrize("descending", [False, True])
+@withCUDA
+def test_sort_native_fullgraph(
+    compilation: None,
+    descending: bool,
+    device: torch.device,
+) -> None:
+    tensor = StringTensor.from_list(
+        [
+            "unused",
+            "é",
+            "unused",
+            "aaaaaaaa",
+            "unused",
+            None,
+            "unused",
+            "aaaaaaaa",
+            "unused",
+            "",
+            "unused",
+            "中",
+        ],
+        device=device,
+    )[1::2]
+
+    def sort(inp: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        if inp.is_cuda:
+            return torch.sort(inp, descending=descending, stable=True)
+        return _sort_native(cast(StringTensor, inp), descending=descending)
+
+    out, perm = torch.compile(sort, fullgraph=True)(tensor)
+    expected = [5, 0, 1, 3, 4, 2] if descending else [4, 1, 3, 0, 5, 2]
+    assert perm.tolist() == expected
+    values = tensor.tolist()
+    assert out.tolist() == [values[i] for i in expected]
+
+
+@withCUDA
+def test_sort_null_payload(device: torch.device) -> None:
+    # Null slots may contain arbitrary bytes and must retain their input order.
+    array = pa.Array.from_buffers(
+        pa.string(),
+        4,
+        [
+            pa.py_buffer(bytes([0b00000101])),
+            pa.py_buffer(
+                torch.tensor([0, 1, 4, 5, 8], dtype=torch.int32).numpy()
+            ),
+            pa.py_buffer(bytearray(b"bzzzazzz")),
+        ],
+    )
+    tensor = StringTensor.from_arrow(array, device=device)
+    out, perm = tensor.sort(stable=True)
+    assert out.tolist() == ["a", "b", None, None]
+    assert perm.tolist() == [2, 0, 1, 3]
+
+
+@withCUDA
+def test_sort_repeated_view(device: torch.device) -> None:
+    tensor = StringTensor.from_list(["abc"], device=device).expand(5)
+    out, perm = tensor.sort(stable=True)
+    assert out.tolist() == ["abc"] * 5
+    assert perm.tolist() == list(range(5))
+
+
+@pytest.mark.parametrize("values", [[], [None, None], ["", ""]])
+@withCUDA
+def test_sort_empty_payload(
+    values: list[str | None],
+    device: torch.device,
+) -> None:
+    tensor = StringTensor.from_list(values, device=device)
+    out, perm = tensor.sort(stable=True)
+    assert out.tolist() == values
+    assert perm.tolist() == list(range(len(values)))

From e27386a4006604bae74e19c568d6dca7d0505cb0 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Mon, 26 Oct 2026 16:42:21 -0700
Subject: [PATCH 2/4] Exercise native string sorting against stable UTF-8
 reference

---
 test/tensor/test_string.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/test/tensor/test_string.py b/test/tensor/test_string.py
index 662bf8d41b..e0aa81695b 100644
--- a/test/tensor/test_string.py
+++ b/test/tensor/test_string.py
@@ -385,6 +385,7 @@ def test_registers_as_a_pytorch_buffer() -> None:
     assert tuple(module.buffers()) == (tensor,)
 
 
+@pytest.mark.parametrize("native", [False, True])
 @pytest.mark.parametrize("descending", [False, True])
 @pytest.mark.parametrize(
     "values",
@@ -402,6 +403,7 @@ def test_registers_as_a_pytorch_buffer() -> None:
 def test_sort_stable(
     values: list[str | None],
     descending: bool,
+    native: bool,
     device: torch.device,
 ) -> None:
     # The sort indices are relative to the view, including a storage offset.
@@ -416,7 +418,12 @@ def test_sort_stable(
     )
     indices.extend(i for i, value in enumerate(values) if value is None)
 
-    out, perm = torch.sort(tensor, descending=descending, stable=True)
+    if native:
+        out, perm = _sort_native(
+            cast(StringTensor, tensor), descending=descending
+        )
+    else:
+        out, perm = torch.sort(tensor, descending=descending, stable=True)
     assert out.tolist() == [values[i] for i in indices]
     assert perm.tolist() == indices
     assert out.device == device

From e726995cd7c5c750f42b9581bef977148ea1e154 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Mon, 26 Oct 2026 16:58:26 -0700
Subject: [PATCH 3/4] Enable native fullgraph sorting on CPU while preserving
 eager Arrow

---
 sdm/tensor/string.py       | 3 ++-
 test/tensor/test_string.py | 4 +---
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/sdm/tensor/string.py b/sdm/tensor/string.py
index 8123e69691..a31889031c 100644
--- a/sdm/tensor/string.py
+++ b/sdm/tensor/string.py
@@ -526,7 +526,8 @@ def _sort(
         raise NotImplementedError("'sort' only supports one-dimensional input")
 
     # Bound radix passes; longer strings retain the external sort.
-    if inp.is_cuda and inp._max_length <= 56:
+    native = inp.is_cuda or (inp.is_cpu and torch.compiler.is_compiling())
+    if native and inp._max_length <= 56:
         return _sort_native(inp, descending=descending)
 
     backend: Literal["arrow", "cudf"] = "arrow"
diff --git a/test/tensor/test_string.py b/test/tensor/test_string.py
index e0aa81695b..9fbb3ec762 100644
--- a/test/tensor/test_string.py
+++ b/test/tensor/test_string.py
@@ -455,9 +455,7 @@ def test_sort_native_fullgraph(
     )[1::2]
 
     def sort(inp: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        if inp.is_cuda:
-            return torch.sort(inp, descending=descending, stable=True)
-        return _sort_native(cast(StringTensor, inp), descending=descending)
+        return torch.sort(inp, descending=descending, stable=True)
 
     out, perm = torch.compile(sort, fullgraph=True)(tensor)
     expected = [5, 0, 1, 3, 4, 2] if descending else [4, 1, 3, 0, 5, 2]

From 21923161fbb513a0853e5dcdf6882f68a158ce2d Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Mon, 26 Oct 2026 22:51:37 -0700
Subject: [PATCH 4/4] Extend CUDA native string sort to twelve words

---
 sdm/tensor/string.py       |  7 ++++---
 test/tensor/test_string.py | 24 ++++++++++++++++++++++++
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/sdm/tensor/string.py b/sdm/tensor/string.py
index a31889031c..9492ef7a8d 100644
--- a/sdm/tensor/string.py
+++ b/sdm/tensor/string.py
@@ -525,9 +525,10 @@ def _sort(
     if inp.dim() != 1:
         raise NotImplementedError("'sort' only supports one-dimensional input")
 
-    # Bound radix passes; longer strings retain the external sort.
-    native = inp.is_cuda or (inp.is_cpu and torch.compiler.is_compiling())
-    if native and inp._max_length <= 56:
+    # Bound radix passes to 12 words on CUDA and eight on CPU.
+    cuda = inp.is_cuda
+    native = cuda or (inp.is_cpu and torch.compiler.is_compiling())
+    if native and inp._max_length <= (84 if cuda else 56):
         return _sort_native(inp, descending=descending)
 
     backend: Literal["arrow", "cudf"] = "arrow"
diff --git a/test/tensor/test_string.py b/test/tensor/test_string.py
index 9fbb3ec762..31bcbea885 100644
--- a/test/tensor/test_string.py
+++ b/test/tensor/test_string.py
@@ -397,6 +397,10 @@ def test_registers_as_a_pytorch_buffer() -> None:
         ["a" * 7 + "b", "a" * 8, "a" * 7, "a" * 14, None],
         ["a" * 55 + "b", "a" * 56, "a" * 55, "a" * 56, None],
         ["a" * 57 + "z", "a" * 58, "a" * 57, "a" * 58, None, None],
+        ["é" * 37 + "b", "é" * 37 + "a", "é" * 37, None],
+        ["a" * 83 + "b", "a" * 84, "a" * 83, "a" * 84, None],
+        ["é" * 42, "é" * 41 + "e", "é" * 41, None],
+        ["a" * 84 + "b", "a" * 85, "a" * 84, "a" * 85, None],
     ],
 )
 @withCUDA
@@ -464,6 +468,26 @@ def sort(inp: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
     assert out.tolist() == [values[i] for i in expected]
 
 
+@pytest.mark.parametrize("width", [75, 84])
+@pytest.mark.parametrize("descending", [False, True])
+@onlyCUDA
+def test_sort_wide_fullgraph(
+    compilation: None,
+    width: int,
+    descending: bool,
+) -> None:
+    values = ["a" * (width - 1) + "b", "a" * width, "a" * width, None, ""]
+    tensor = StringTensor.from_list(values, device="cuda")
+
+    def sort(inp: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        return torch.sort(inp, descending=descending, stable=True)
+
+    out, perm = torch.compile(sort, fullgraph=True)(tensor)
+    expected = [0, 1, 2, 4, 3] if descending else [4, 1, 2, 0, 3]
+    assert perm.tolist() == expected
+    assert out.tolist() == [values[i] for i in expected]
+
+
 @withCUDA
 def test_sort_null_payload(device: torch.device) -> None:
     # Null slots may contain arbitrary bytes and must retain their input order.