chore: import upstream snapshot with attribution

tests/python/serve/test_radix_tree.py

@@ -0,0 +1,138 @@
+import pytest
+
+from mlc_llm.serve import PagedRadixTree
+
+# category "runtime_module"
+pytestmark = [pytest.mark.unittest]
+
+
+def test_add():
+    prt = PagedRadixTree()
+    prt.add(0)
+    assert list(prt.get(0)) == []
+    prt.add(1)
+    assert list(prt.get(1)) == []
+
+
+def test_remove():
+    prt = PagedRadixTree()
+    capacity = prt.free_capacity()
+    prt.add(0)
+    prt.remove(0)
+    prt.add(0)
+    prt.extend(0, [1 for _ in range(200)])
+    prt.remove(0)
+    assert prt.free_capacity() == capacity
+
+    prt.add(1)
+    prt.extend(1, [1 for _ in range(200)])
+    capacity = prt.free_capacity()
+    prt.add(2)
+    prt.extend(2, [1 for _ in range(100)] + [2 for _ in range(100)])
+    prt.remove(2)
+    assert prt.free_capacity() == capacity
+
+    prt.add(3)
+    prt.extend(3, [1 for _ in range(200)])
+    prt.remove(3)
+    assert prt.free_capacity() == capacity
+
+    prt.add(4)
+    prt.add(5)
+    prt.add(6)
+    assert prt.free_capacity() == capacity
+    prt.remove(4)
+    assert prt.free_capacity() == capacity
+    prt.remove(5)
+    assert prt.free_capacity() == capacity
+    prt.remove(6)
+    assert prt.free_capacity() == capacity
+
+
+def test_extend():
+    prt = PagedRadixTree()
+    L = prt.free_capacity() // 64
+    H = L // 2
+    Q = L // 4
+    seq_id = 0
+    for start_pos in [0, H, L, L + H]:
+        for length in [Q, L - H, L, 2 * L - H, 2 * L]:
+            prt.add(seq_id)
+            if start_pos:
+                tokens_1 = [seq_id for _ in range(start_pos)]
+                prt.extend(seq_id, tokens_1)
+                assert list(prt.get(seq_id)) == tokens_1
+            else:
+                tokens_1 = []
+            tokens_2 = [seq_id for _ in range(length)]
+            prt.extend(seq_id, tokens_2)
+            assert list(prt.get(seq_id)) == tokens_1 + tokens_2
+            seq_id += 1
+
+
+def test_fork():
+    prt = PagedRadixTree()
+    L = prt.free_capacity() // 64
+    H = L // 2
+    Q = L // 4
+    seq_id = 0
+    length_list = [Q, H, L, L + Q, L + H, L * 2]
+    for p_idx in range(1, len(length_list)):
+        for c_idx in range(0, p_idx + 1):
+            prt.add(seq_id)
+            tokens = [seq_id for _ in range(length_list[p_idx])]
+            prt.extend(seq_id, tokens)
+            prt.fork(seq_id + 1, seq_id, length_list[c_idx])
+            assert list(prt.get(seq_id + 1)) == tokens[: length_list[c_idx]]
+            seq_id += 2
+
+
+def test_fork_2():
+    prt = PagedRadixTree()
+    prt.add(0)
+    prt.extend(0, [0, 1, 2, 3])
+    prt.fork(1, 0, 3)
+    prt.extend(1, [4])
+    prt.fork(2, 0, 3)
+    prt.extend(2, [5])
+    assert prt.match([0, 1, 2, 4]) == (4, (1,))
+    assert prt.match([0, 1, 2, 5]) == (4, (2,))
+
+
+def test_rollback():
+    prt = PagedRadixTree()
+    L = prt.free_capacity() // 64
+    H = L // 2
+    Q = L // 4
+    seq_id = 0
+    for start_pos in [H, L, L + H, 2 * L, 3 * L + H]:
+        for length in [Q, H, L + Q, 2 * L, 2 * L + Q]:
+            if length > start_pos:
+                continue
+            prt.add(seq_id)
+            tokens = [seq_id for _ in range(start_pos)]
+            prt.extend(seq_id, tokens)
+            prt.rollback(seq_id, length)
+            assert list(prt.get(seq_id)) == tokens[:-length]
+            seq_id += 1
+
+    for start_pos in [H, L, L + H, 2 * L, 3 * L + H]:
+        for length in [Q, H, L + Q, 2 * L, 2 * L + Q]:
+            if length > start_pos:
+                continue
+            prt.add(seq_id)
+            tokens = [seq_id for _ in range(start_pos)]
+            prt.extend(seq_id, tokens)
+            prt.fork(seq_id + 1, seq_id, start_pos)
+            prt.rollback(seq_id + 1, length)
+            assert list(prt.get(seq_id + 1)) == tokens[:-length]
+            seq_id += 2
+
+
+if __name__ == "__main__":
+    test_add()
+    test_remove()
+    test_extend()
+    test_fork()
+    test_fork_2()
+    test_rollback()