diff --git a/benchmarks/test_benchmark_knn.py b/benchmarks/test_benchmark_knn.py
new file mode 100644
index 00000000..8ac1c9e7
--- /dev/null
+++ b/benchmarks/test_benchmark_knn.py
@@ -0,0 +1,495 @@
+import importlib.util
+from itertools import product
+
+import pytest
+import torch
+import torch_cluster as tc
+
+knn = tc.knn
+knn_graph = tc.knn_graph
+
+
+pytestmark = pytest.mark.skipif(
+    not (
+        torch.ops.torch_cluster.cuda_version() != -1
+        and importlib.util.find_spec('triton') is not None
+    ),
+    reason='CUDA and Triton are required for Triton benchmark tests.',
+)
+
+KNN_SIZES = [
+    (256, 128),
+    (512, 256),
+    (1024, 512),
+    (2048, 1024),
+    (4096, 2048),
+    (8192, 4096),
+    (8192, 8192),
+    (8201, 4103),
+    (32000, 32000),
+    (255, 127),
+    (256, 5),
+    (1024, 5),
+    (4096, 5),
+    (255, 5),
+]
+KNN_GROUPS = [1, 2, 4, 8, 16, 32]
+
+FEATURES = [3, 8, 64, 200]
+
+
+def to_set(edge_index):
+    return set([(i, j) for i, j in edge_index.t().tolist()])
+
+
+def _assert_knn_within_cuda(
+    out_cuda,
+    out_triton,
+    x,
+    y,
+    k,
+    cosine,
+    tol=None,
+):
+    if tol is None:
+        tol = 5 * torch.finfo(x.dtype).eps
+    m = y.size(0)
+    cuda_rows = out_cuda[0]
+    cuda_cols = out_cuda[1]
+    triton_rows = out_triton[0]
+    triton_cols = out_triton[1]
+    y_f = y.float()
+    x_f = x.float()
+    y_norm = torch.linalg.norm(y_f, dim=1)
+    if cosine:
+        x_cuda = x_f[cuda_cols]
+        y_cuda = y_f[cuda_rows]
+        cuda_dot = (x_cuda * y_cuda).sum(dim=1)
+        cuda_norm = torch.linalg.norm(x_cuda, dim=1)
+        cuda_dist = 1.0 - cuda_dot / (cuda_norm * y_norm[cuda_rows])
+        x_triton = x_f[triton_cols]
+        y_triton = y_f[triton_rows]
+        triton_dot = (x_triton * y_triton).sum(dim=1)
+        triton_norm = torch.linalg.norm(x_triton, dim=1)
+        triton_dist = 1.0 - triton_dot / (
+            triton_norm * y_norm[triton_rows]
+        )
+    else:
+        x_cuda = x_f[cuda_cols]
+        y_cuda = y_f[cuda_rows]
+        cuda_dist = ((x_cuda - y_cuda) ** 2).sum(dim=1)
+        x_triton = x_f[triton_cols]
+        y_triton = y_f[triton_rows]
+        triton_dist = ((x_triton - y_triton) ** 2).sum(dim=1)
+    cuda_max = torch.full(
+        (m,),
+        -float("inf"),
+        device=y.device,
+        dtype=torch.float32,
+    )
+    cuda_max.scatter_reduce_(
+        0,
+        cuda_rows,
+        cuda_dist,
+        reduce="amax",
+        include_self=True,
+    )
+    triton_thresh = cuda_max[triton_rows] + tol
+    margin = (triton_dist - triton_thresh).max().item()
+    if cosine:
+        x_ref = x_f[triton_cols]
+        y_ref = y_f[triton_rows]
+        ref_dot = (x_ref * y_ref).sum(dim=1)
+        ref_norm = torch.linalg.norm(x_ref, dim=1)
+        ref_dist = 1.0 - ref_dot / (ref_norm * y_norm[triton_rows])
+    else:
+        x_ref = x_f[triton_cols]
+        y_ref = y_f[triton_rows]
+        ref_dist = ((x_ref - y_ref) ** 2).sum(dim=1)
+    max_diff = torch.abs(triton_dist - ref_dist).max().item()
+    print(f"[knn][match] max_margin={margin:.6e} tol={tol:.1e}")
+    print(f"[knn][match] max_diff={max_diff:.6e} tol={tol:.1e}")
+    assert (triton_dist <= triton_thresh).all()
+    assert max_diff <= tol
+
+
+def _make_batch(
+    num_nodes: int,
+    num_groups: int,
+    device: torch.device,
+) -> torch.Tensor:
+    groups = max(1, min(num_groups, num_nodes))
+    counts = torch.full(
+        (groups,),
+        num_nodes // groups,
+        device=device,
+        dtype=torch.long,
+    )
+    remainder = num_nodes % groups
+    if remainder:
+        counts[:remainder] += 1
+    return torch.repeat_interleave(
+        torch.arange(groups, device=device),
+        counts,
+    )
+
+
+def _knn_param_grid():
+    return (
+        (*p[0], p[1], p[2])
+        for p in product(KNN_SIZES, KNN_GROUPS, FEATURES)
+        if p[1] <= min(p[0])
+    )
+
+
+@pytest.mark.parametrize(
+    'num_x,num_y,num_groups,num_features',
+    _knn_param_grid(),
+)
+@pytest.mark.benchmark(group="knn")
+def test_triton_knn_benchmark_cuda(
+    benchmark,
+    num_x,
+    num_y,
+    num_groups,
+    num_features,
+):
+    torch.manual_seed(99)
+    x = torch.randn(num_x, num_features, device='cuda')
+    y = torch.randn(num_y, num_features, device='cuda')
+    groups = min(num_groups, x.size(0), y.size(0))
+    batch_x = _make_batch(num_x, groups, x.device)
+    batch_y = _make_batch(num_y, groups, y.device)
+
+    def cuda_fn():
+        knn(
+            x,
+            y,
+            k=16,
+            batch_x=batch_x,
+            batch_y=batch_y,
+            cosine=False,
+            use_triton=False,
+        )
+
+    for _ in range(5):
+        cuda_fn()
+        torch.cuda.synchronize()
+
+    benchmark(cuda_fn)
+    print(
+        f"[knn][cuda] num_x={num_x} num_y={num_y} groups={groups}"
+    )
+
+
+@pytest.mark.parametrize(
+    'num_x,num_y,num_groups,num_features',
+    _knn_param_grid(),
+)
+@pytest.mark.benchmark(group="knn_cosine")
+def test_triton_knn_benchmark_cuda_cosine(
+    benchmark,
+    num_x,
+    num_y,
+    num_groups,
+    num_features,
+):
+    torch.manual_seed(99)
+    x = torch.randn(num_x, num_features, device='cuda')
+    y = torch.randn(num_y, num_features, device='cuda')
+    groups = min(num_groups, x.size(0), y.size(0))
+    batch_x = _make_batch(num_x, groups, x.device)
+    batch_y = _make_batch(num_y, groups, y.device)
+
+    def cuda_fn():
+        knn(
+            x,
+            y,
+            k=16,
+            batch_x=batch_x,
+            batch_y=batch_y,
+            cosine=True,
+            use_triton=False,
+        )
+
+    for _ in range(5):
+        cuda_fn()
+        torch.cuda.synchronize()
+
+    benchmark(cuda_fn)
+    print(
+        f"[knn][cuda] num_x={num_x} num_y={num_y} groups={groups}"
+    )
+
+
+@pytest.mark.parametrize(
+    'num_x,num_y,num_groups,num_features',
+    _knn_param_grid(),
+)
+@pytest.mark.benchmark(group="knn_cosine")
+def test_triton_knn_benchmark_triton_cosine(
+    benchmark,
+    num_x,
+    num_y,
+    num_groups,
+    num_features,
+):
+    torch.manual_seed(99)
+    x = torch.randn(num_x, num_features, device='cuda')
+    y = torch.randn(num_y, num_features, device='cuda')
+    groups = min(num_groups, x.size(0), y.size(0))
+    batch_x = _make_batch(num_x, groups, x.device)
+    batch_y = _make_batch(num_y, groups, y.device)
+
+    def cuda_fn():
+        return knn(
+            x,
+            y,
+            k=16,
+            batch_x=batch_x,
+            batch_y=batch_y,
+            cosine=True,
+            use_triton=False,
+        )
+
+    def triton_fn():
+        return knn(
+            x,
+            y,
+            k=16,
+            batch_x=batch_x,
+            batch_y=batch_y,
+            cosine=True,
+            use_triton=True,
+        )
+
+    for i in range(5):
+        if i == 0:
+            out_cuda = cuda_fn()
+            out_triton = triton_fn()
+            _assert_knn_within_cuda(
+                out_cuda,
+                out_triton,
+                x,
+                y,
+                k=16,
+                cosine=True,
+            )
+        else:
+            triton_fn()
+        torch.cuda.synchronize()
+
+    benchmark(triton_fn)
+    print(
+        f"[knn][triton] num_x={num_x} num_y={num_y} groups={groups}"
+    )
+
+
+@pytest.mark.parametrize(
+    'num_x,num_y,num_groups,num_features',
+    _knn_param_grid(),
+)
+@pytest.mark.benchmark(group="knn")
+def test_triton_knn_benchmark_triton(
+    benchmark,
+    num_x,
+    num_y,
+    num_groups,
+    num_features,
+):
+    torch.manual_seed(99)
+    x = torch.randn(num_x, num_features, device='cuda')
+    y = torch.randn(num_y, num_features, device='cuda')
+    groups = min(num_groups, x.size(0), y.size(0))
+    batch_x = _make_batch(num_x, groups, x.device)
+    batch_y = _make_batch(num_y, groups, y.device)
+
+    def cuda_fn():
+        return knn(
+            x,
+            y,
+            k=16,
+            batch_x=batch_x,
+            batch_y=batch_y,
+            cosine=False,
+            use_triton=False,
+        )
+
+    def triton_fn():
+        return knn(
+            x,
+            y,
+            k=16,
+            batch_x=batch_x,
+            batch_y=batch_y,
+            cosine=False,
+            use_triton=True,
+        )
+
+    for i in range(5):
+        if i == 0:
+            out_cuda = cuda_fn()
+            out_triton = triton_fn()
+            _assert_knn_within_cuda(
+                out_cuda,
+                out_triton,
+                x,
+                y,
+                k=16,
+                cosine=False,
+            )
+        else:
+            triton_fn()
+        torch.cuda.synchronize()
+
+    benchmark(triton_fn)
+    print(
+        f"[knn][triton] num_x={num_x} num_y={num_y} groups={groups}"
+    )
+
+
+@pytest.mark.parametrize('num_x', [256, 1024, 4096, 8192, 255])
+@pytest.mark.parametrize('num_groups', [1, 2, 4, 6, 8, 16, 24, 32])
+@pytest.mark.benchmark(group="knn_graph")
+def test_triton_knn_graph_benchmark_cuda(benchmark, num_x, num_groups):
+    torch.manual_seed(199)
+    x = torch.randn(num_x, 8, device='cuda')
+    groups = min(num_groups, x.size(0))
+    batch = _make_batch(num_x, groups, x.device)
+    k = min(16, max(1, num_x - 1))
+
+    def cuda_fn():
+        knn_graph(x, k=k, batch=batch, loop=False, use_triton=False)
+
+    for _ in range(5):
+        cuda_fn()
+        torch.cuda.synchronize()
+
+    benchmark(cuda_fn)
+    print(
+        f"[knn_graph][cuda] num_x={num_x} groups={groups} k={k}"
+    )
+
+
+@pytest.mark.parametrize('num_x', [256, 1024, 4096, 8192, 255])
+@pytest.mark.parametrize('num_groups', [1, 2, 4, 6, 8, 16, 24, 32])
+@pytest.mark.benchmark(group="knn_graph")
+def test_triton_knn_graph_benchmark_triton(benchmark, num_x, num_groups):
+    torch.manual_seed(199)
+    x = torch.randn(num_x, 8, device='cuda')
+    groups = min(num_groups, x.size(0))
+    batch = _make_batch(num_x, groups, x.device)
+    k = min(16, max(1, num_x - 1))
+
+    def cuda_fn():
+        return knn_graph(
+            x,
+            k=k,
+            batch=batch,
+            loop=False,
+            use_triton=False,
+        )
+
+    def triton_fn():
+        return knn_graph(
+            x,
+            k=k,
+            batch=batch,
+            loop=False,
+            use_triton=True,
+        )
+
+    for i in range(5):
+        if i == 0:
+            out_cuda = cuda_fn()
+            out_triton = triton_fn()
+            for a, b in zip(
+                sorted(list(to_set(out_cuda))),
+                sorted(list(to_set(out_triton))),
+            ):
+                assert a == b
+        else:
+            triton_fn()
+        torch.cuda.synchronize()
+
+    benchmark(triton_fn)
+    print(f"[knn_graph][triton] num_x={num_x} groups={groups} k={k}")
+
+
+@pytest.mark.parametrize('num_x', [256, 1024, 4096, 8192, 255])
+@pytest.mark.parametrize('num_groups', [1, 2, 4, 6, 8, 16, 24, 32])
+@pytest.mark.benchmark(group="knn_graph_cosine")
+def test_triton_knn_graph_benchmark_cuda_cosine(benchmark, num_x, num_groups):
+    torch.manual_seed(199)
+    x = torch.randn(num_x, 8, device='cuda')
+    groups = min(num_groups, x.size(0))
+    batch = _make_batch(num_x, groups, x.device)
+    k = min(16, max(1, num_x - 1))
+
+    def cuda_fn():
+        knn_graph(
+            x,
+            k=k,
+            batch=batch,
+            loop=False,
+            use_triton=False,
+            cosine=True,
+        )
+
+    for _ in range(5):
+        cuda_fn()
+        torch.cuda.synchronize()
+
+    benchmark(cuda_fn)
+    print(f"[knn_graph][cuda] num_x={num_x} groups={groups} k={k}")
+
+
+@pytest.mark.parametrize('num_x', [256, 1024, 4096, 8192, 255])
+@pytest.mark.parametrize('num_groups', [1, 2, 4, 6, 8, 16, 24, 32])
+@pytest.mark.benchmark(group="knn_graph_cosine")
+def test_triton_knn_graph_benchmark_triton_cosine(
+    benchmark,
+    num_x,
+    num_groups,
+):
+    torch.manual_seed(199)
+    x = torch.randn(num_x, 8, device='cuda')
+    groups = min(num_groups, x.size(0))
+    batch = _make_batch(num_x, groups, x.device)
+    k = min(16, max(1, num_x - 1))
+
+    def cuda_fn():
+        return knn_graph(
+            x,
+            k=k,
+            batch=batch,
+            loop=False,
+            use_triton=False,
+            cosine=True,
+        )
+
+    def triton_fn():
+        return knn_graph(
+            x,
+            k=k,
+            batch=batch,
+            loop=False,
+            use_triton=True,
+            cosine=True,
+        )
+
+    for i in range(5):
+        if i == 0:
+            out_cuda = cuda_fn()
+            out_triton = triton_fn()
+            for a, b in zip(
+                sorted(list(to_set(out_cuda))),
+                sorted(list(to_set(out_triton))),
+            ):
+                assert a == b
+        else:
+            triton_fn()
+        torch.cuda.synchronize()
+
+    benchmark(triton_fn)
+    print(f"[knn_graph][triton] num_x={num_x} groups={groups} k={k}")
diff --git a/benchmarks/test_benchmark_nearest.py b/benchmarks/test_benchmark_nearest.py
new file mode 100644
index 00000000..0d1c151d
--- /dev/null
+++ b/benchmarks/test_benchmark_nearest.py
@@ -0,0 +1,157 @@
+import importlib.util
+from itertools import product
+from typing import Optional
+
+import pytest
+import torch
+import torch_cluster as tc
+
+nearest = tc.nearest
+
+pytestmark = pytest.mark.skipif(
+    not (
+        torch.ops.torch_cluster.cuda_version() != -1
+        and importlib.util.find_spec('triton') is not None
+    ),
+    reason='CUDA and Triton are required for Triton benchmark tests.',
+)
+
+NEAREST_SIZES = [
+    (256, 128),
+    (1024, 512),
+    (4096, 2048),
+    (8192, 4096),
+    (8192, 8192),
+    (8201, 4103),
+    (32000, 32000),
+    (255, 127),
+    (256, 5),
+    (1024, 5),
+    (4096, 5),
+    (255, 5),
+]
+NEAREST_GROUPS = [1, 2, 4, 8, 16, 32]
+FEATURES = [8, 64, 200]
+
+
+def _assert_nearest_within_cuda(
+    out_cuda: torch.Tensor,
+    out_triton: torch.Tensor,
+    x: torch.Tensor,
+    y: torch.Tensor,
+    tol: Optional[float] = None,
+) -> None:
+    if tol is None:
+        tol = 5 * torch.finfo(x.dtype).eps
+    x_f = x.float()
+    y_f = y.float()
+    cuda_dist = ((x_f - y_f[out_cuda]) ** 2).sum(dim=1)
+    triton_dist = ((x_f - y_f[out_triton]) ** 2).sum(dim=1)
+    thresh = cuda_dist + tol
+    margin = (triton_dist - thresh).max().item()
+    max_diff = torch.abs(triton_dist - cuda_dist).max().item()
+    print(f"[nearest][match] max_margin={margin:.6e} tol={tol:.1e}")
+    print(f"[nearest][match] max_diff={max_diff:.6e} tol={tol:.1e}")
+    assert (triton_dist <= thresh).all()
+    assert max_diff <= tol
+
+
+def _make_batch(
+    num_nodes: int,
+    num_groups: int,
+    device: torch.device,
+) -> torch.Tensor:
+    groups = max(1, min(num_groups, num_nodes))
+    counts = torch.full(
+        (groups,),
+        num_nodes // groups,
+        device=device,
+        dtype=torch.long,
+    )
+    remainder = num_nodes % groups
+    if remainder:
+        counts[:remainder] += 1
+    return torch.repeat_interleave(
+        torch.arange(groups, device=device),
+        counts,
+    )
+
+
+def _nearest_param_grid():
+    return (
+        (*p[0], p[1], p[2])
+        for p in product(NEAREST_SIZES, NEAREST_GROUPS, FEATURES)
+        if p[1] <= min(p[0])
+    )
+
+
+@pytest.mark.parametrize(
+    'num_x,num_y,num_groups,num_features',
+    _nearest_param_grid(),
+)
+@pytest.mark.benchmark(group="nearest")
+def test_triton_nearest_benchmark_cuda(
+    benchmark,
+    num_x,
+    num_y,
+    num_groups,
+    num_features,
+):
+    torch.manual_seed(123)
+    x = torch.randn(num_x, num_features, device='cuda')
+    y = torch.randn(num_y, num_features, device='cuda')
+    groups = min(num_groups, x.size(0), y.size(0))
+    batch_x = _make_batch(num_x, groups, x.device)
+    batch_y = _make_batch(num_y, groups, y.device)
+
+    def cuda_fn():
+        nearest(x, y, batch_x, batch_y, use_triton=False)
+
+    for _ in range(5):
+        cuda_fn()
+        torch.cuda.synchronize()
+
+    benchmark(cuda_fn)
+    print(
+        f"[nearest][cuda] num_x={num_x} num_y={num_y} groups={groups}"
+    )
+
+
+@pytest.mark.parametrize(
+    'num_x,num_y,num_groups,num_features',
+    _nearest_param_grid(),
+)
+@pytest.mark.benchmark(group="nearest")
+def test_triton_nearest_benchmark_triton(
+    benchmark,
+    num_x,
+    num_y,
+    num_groups,
+    num_features,
+):
+    torch.manual_seed(123)
+    x = torch.randn(num_x, num_features, device='cuda')
+    y = torch.randn(num_y, num_features, device='cuda')
+    groups = min(num_groups, x.size(0), y.size(0))
+    batch_x = _make_batch(num_x, groups, x.device)
+    batch_y = _make_batch(num_y, groups, y.device)
+
+    def cuda_fn():
+        return nearest(x, y, batch_x, batch_y, use_triton=False)
+
+    def triton_fn():
+        return nearest(x, y, batch_x, batch_y, use_triton=True)
+
+    for i in range(5):
+        if i == 0:
+            out_cuda = cuda_fn()
+            out_triton = triton_fn()
+            _assert_nearest_within_cuda(out_cuda, out_triton, x, y)
+        else:
+            triton_fn()
+        torch.cuda.synchronize()
+
+    benchmark(triton_fn)
+    print(
+        f"[nearest][triton] num_x={num_x} num_y={num_y} groups={groups}"
+    )
diff --git a/benchmarks/test_benchmark_radius.py b/benchmarks/test_benchmark_radius.py
new file mode 100644
index 00000000..52a9532d
--- /dev/null
+++ b/benchmarks/test_benchmark_radius.py
@@ -0,0 +1,279 @@
+import importlib.util
+from itertools import product
+
+import pytest
+import torch
+import torch_cluster as tc
+
+radius = tc.radius
+radius_graph = tc.radius_graph
+
+pytestmark = pytest.mark.skipif(
+    not (
+        torch.ops.torch_cluster.cuda_version() != -1
+        and importlib.util.find_spec('triton') is not None
+    ),
+    reason='CUDA and Triton are required for Triton benchmark tests.',
+)
+
+RADIUS_SIZES = [
+    (256, 128),
+    (512, 256),
+    (1024, 512),
+    (2048, 1024),
+    (4096, 2048),
+    (8192, 4096),
+    (8192, 8192),
+    (8201, 4103),
+    (32000, 32000),
+    (255, 127),
+    (256, 5),
+    (1024, 5),
+    (4096, 5),
+    (255, 5),
+]
+RADIUS_GROUPS = [1, 2, 4, 8, 16, 32]
+FEATURES = [3, 8, 64, 200]
+
+
+def to_set(edge_index):
+    return set([(i, j) for i, j in edge_index.t().tolist()])
+
+
+def _assert_radius_within_cuda(
+    edge_index,
+    x,
+    y,
+    r,
+    max_num_neighbors,
+    ignore_same_index,
+    check_max_neighbors=True,
+    tol=None,
+):
+    if edge_index.numel() == 0:
+        return
+    if tol is None:
+        tol = 5 * torch.finfo(x.dtype).eps
+    row, col = edge_index
+    x_f = x.float()
+    y_f = y.float()
+    diffs = x_f[col] - y_f[row]
+    dist = (diffs * diffs).sum(dim=1)
+    r2 = float(r) * float(r)
+    assert (dist <= (r2 + tol)).all()
+    if ignore_same_index:
+        assert (row != col).all()
+    if check_max_neighbors:
+        counts = torch.bincount(row, minlength=y.size(0))
+        assert (counts <= max_num_neighbors).all()
+
+
+def _make_batch(
+    num_nodes: int,
+    num_groups: int,
+    device: torch.device,
+) -> torch.Tensor:
+    groups = max(1, min(num_groups, num_nodes))
+    counts = torch.full(
+        (groups,),
+        num_nodes // groups,
+        device=device,
+        dtype=torch.long,
+    )
+    remainder = num_nodes % groups
+    if remainder:
+        counts[:remainder] += 1
+    return torch.repeat_interleave(
+        torch.arange(groups, device=device),
+        counts,
+    )
+
+
+def _radius_param_grid():
+    return (
+        (*p[0], p[1], p[2])
+        for p in product(RADIUS_SIZES, RADIUS_GROUPS, FEATURES)
+        if p[1] <= min(p[0])
+    )
+
+
+@pytest.mark.parametrize(
+    'num_x,num_y,num_groups,num_features',
+    _radius_param_grid(),
+)
+@pytest.mark.benchmark(group="radius")
+def test_triton_radius_benchmark_cuda(
+    benchmark,
+    num_x,
+    num_y,
+    num_groups,
+    num_features,
+):
+    torch.manual_seed(99)
+    x = torch.randn(num_x, num_features, device='cuda')
+    y = torch.randn(num_y, num_features, device='cuda')
+    groups = min(num_groups, x.size(0), y.size(0))
+    batch_x = _make_batch(num_x, groups, x.device)
+    batch_y = _make_batch(num_y, groups, y.device)
+
+    def cuda_fn():
+        radius(
+            x,
+            y,
+            r=1.5,
+            batch_x=batch_x,
+            batch_y=batch_y,
+            max_num_neighbors=32,
+            use_triton=False,
+        )
+
+    for _ in range(5):
+        cuda_fn()
+        torch.cuda.synchronize()
+
+    benchmark(cuda_fn)
+    print(
+        f"[radius][cuda] num_x={num_x} num_y={num_y} groups={groups}"
+    )
+
+
+@pytest.mark.parametrize(
+    'num_x,num_y,num_groups,num_features',
+    _radius_param_grid(),
+)
+@pytest.mark.benchmark(group="radius")
+def test_triton_radius_benchmark_triton(
+    benchmark,
+    num_x,
+    num_y,
+    num_groups,
+    num_features,
+):
+    torch.manual_seed(99)
+    x = torch.randn(num_x, num_features, device='cuda')
+    y = torch.randn(num_y, num_features, device='cuda')
+    groups = min(num_groups, x.size(0), y.size(0))
+    batch_x = _make_batch(num_x, groups, x.device)
+    batch_y = _make_batch(num_y, groups, y.device)
+
+    def cuda_fn():
+        return radius(
+            x,
+            y,
+            r=1.5,
+            batch_x=batch_x,
+            batch_y=batch_y,
+            max_num_neighbors=32,
+            use_triton=False,
+        )
+
+    def triton_fn():
+        return radius(
+            x,
+            y,
+            r=1.5,
+            batch_x=batch_x,
+            batch_y=batch_y,
+            max_num_neighbors=32,
+            use_triton=True,
+        )
+
+    for i in range(5):
+        if i == 0:
+            out_triton = triton_fn()
+            _assert_radius_within_cuda(
+                out_triton,
+                x,
+                y,
+                r=1.5,
+                max_num_neighbors=32,
+                ignore_same_index=False,
+            )
+        else:
+            triton_fn()
+        torch.cuda.synchronize()
+
+    benchmark(triton_fn)
+    print(
+        f"[radius][triton] num_x={num_x} num_y={num_y} groups={groups}"
+    )
+
+
+@pytest.mark.parametrize('num_x', [256, 1024, 4096, 8192, 255])
+@pytest.mark.parametrize('num_groups', [1, 2, 4, 6, 8, 16, 24, 32])
+@pytest.mark.benchmark(group="radius_graph")
+def test_triton_radius_graph_benchmark_cuda(benchmark, num_x, num_groups):
+    torch.manual_seed(199)
+    x = torch.randn(num_x, 8, device='cuda')
+    groups = min(num_groups, x.size(0))
+    batch = _make_batch(num_x, groups, x.device)
+
+    def cuda_fn():
+        radius_graph(
+            x,
+            r=1.5,
+            batch=batch,
+            loop=False,
+            max_num_neighbors=32,
+            use_triton=False,
+        )
+
+    for _ in range(5):
+        cuda_fn()
+        torch.cuda.synchronize()
+
+    benchmark(cuda_fn)
+    print(
+        f"[radius_graph][cuda] num_x={num_x} groups={groups}"
+    )
+
+
+@pytest.mark.parametrize('num_x', [256, 1024, 4096, 8192, 255])
+@pytest.mark.parametrize('num_groups', [1, 2, 4, 6, 8, 16, 24, 32])
+@pytest.mark.benchmark(group="radius_graph")
+def test_triton_radius_graph_benchmark_triton(benchmark, num_x, num_groups):
+    torch.manual_seed(199)
+    x = torch.randn(num_x, 8, device='cuda')
+    groups = min(num_groups, x.size(0))
+    batch = _make_batch(num_x, groups, x.device)
+
+    def cuda_fn():
+        return radius_graph(
+            x,
+            r=1.5,
+            batch=batch,
+            loop=False,
+            max_num_neighbors=32,
+            use_triton=False,
+        )
+
+    def triton_fn():
+        return radius_graph(
+            x,
+            r=1.5,
+            batch=batch,
+            loop=False,
+            max_num_neighbors=32,
+            use_triton=True,
+        )
+
+    for i in range(5):
+        if i == 0:
+            out_triton = triton_fn()
+            _assert_radius_within_cuda(
+                out_triton,
+                x,
+                x,
+                r=1.5,
+                max_num_neighbors=32,
+                ignore_same_index=True,
+                check_max_neighbors=False,
+            )
+        else:
+            triton_fn()
+        torch.cuda.synchronize()
+
+    benchmark(triton_fn)
+    print(
+        f"[radius_graph][triton] num_x={num_x} groups={groups}"
+    )
diff --git a/csrc/cluster.h b/csrc/cluster.h
index bd0afdf5..6f2e4f64 100644
--- a/csrc/cluster.h
+++ b/csrc/cluster.h
@@ -1,4 +1,5 @@
 #pragma once
+#include <torch/library.h>
 
 #include "extensions.h"
 
@@ -37,3 +38,15 @@ random_walk(torch::Tensor rowptr, torch::Tensor col, torch::Tensor start,
 
 CLUSTER_API torch::Tensor neighbor_sampler(torch::Tensor start, torch::Tensor rowptr,
                                int64_t count, double factor);
+
+TORCH_LIBRARY(torch_cluster, m) {
+    m.def("fps(Tensor src, Tensor ptr, Tensor ratio, bool random_start = False) -> Tensor");
+    m.def("graclus(Tensor rowptr, Tensor col, Tensor? weight = None) -> Tensor");
+    m.def("grid(Tensor pos, Tensor size, Tensor? start = None, Tensor? end = None) -> Tensor");
+    m.def("knn(Tensor a, Tensor b, Tensor? ptr_x, Tensor? ptr_y, int k, bool cosine = False, int num_workers = 1) -> Tensor");
+    m.def("nearest(Tensor x, Tensor y, Tensor ptr_x, Tensor ptr_y) -> Tensor");
+    m.def("radius(Tensor x, Tensor y, Tensor? ptr_x, Tensor? ptr_y, float r, int max_num_neighbors, int num_workers = 1, bool ignore_same_index = False) -> Tensor");
+    m.def("random_walk(Tensor rowptr, Tensor col, Tensor start, int walk_length, float p = 1, float q = 1) -> (Tensor, Tensor)");
+    m.def("neighbor_sampler(Tensor start, Tensor rowptr, int count, float factor) -> Tensor");
+    m.def("cuda_version() -> int");
+}
diff --git a/csrc/fps.cpp b/csrc/fps.cpp
index db395336..be137422 100644
--- a/csrc/fps.cpp
+++ b/csrc/fps.cpp
@@ -1,7 +1,8 @@
 #ifdef WITH_PYTHON
 #include <Python.h>
 #endif
-#include <torch/script.h>
+#include <torch/torch.h>
+#include <torch/library.h>
 
 #include "cpu/fps_cpu.h"
 
@@ -32,5 +33,15 @@ CLUSTER_API torch::Tensor fps(torch::Tensor src, torch::Tensor ptr, torch::Tenso
   }
 }
 
-static auto registry =
-    torch::RegisterOperators().op("torch_cluster::fps", &fps);
+TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
+  m.impl("fps", &fps_cpu);
+}
+
+#ifdef WITH_CUDA
+TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
+  m.impl("fps", &fps_cuda);
+}
+TORCH_LIBRARY_IMPL(torch_cluster, HIP, m) {
+  m.impl("fps", &fps_cuda);
+}
+#endif
diff --git a/csrc/graclus.cpp b/csrc/graclus.cpp
index 73338553..7b9b7b19 100644
--- a/csrc/graclus.cpp
+++ b/csrc/graclus.cpp
@@ -1,7 +1,8 @@
 #ifdef WITH_PYTHON
 #include <Python.h>
 #endif
-#include <torch/script.h>
+#include <torch/torch.h>
+#include <torch/library.h>
 
 #include "cpu/graclus_cpu.h"
 
@@ -32,5 +33,15 @@ CLUSTER_API torch::Tensor graclus(torch::Tensor rowptr, torch::Tensor col,
   }
 }
 
-static auto registry =
-    torch::RegisterOperators().op("torch_cluster::graclus", &graclus);
+TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
+  m.impl("graclus", &graclus_cpu);
+}
+
+#ifdef WITH_CUDA
+TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
+  m.impl("graclus", &graclus_cuda);
+}
+TORCH_LIBRARY_IMPL(torch_cluster, HIP, m) {
+  m.impl("graclus", &graclus_cuda);
+}
+#endif
diff --git a/csrc/grid.cpp b/csrc/grid.cpp
index f77512d5..ba1cb47b 100644
--- a/csrc/grid.cpp
+++ b/csrc/grid.cpp
@@ -1,7 +1,8 @@
 #ifdef WITH_PYTHON
 #include <Python.h>
 #endif
-#include <torch/script.h>
+#include <torch/torch.h>
+#include <torch/library.h>
 
 #include "cpu/grid_cpu.h"
 
@@ -33,5 +34,15 @@ CLUSTER_API torch::Tensor grid(torch::Tensor pos, torch::Tensor size,
   }
 }
 
-static auto registry =
-    torch::RegisterOperators().op("torch_cluster::grid", &grid);
+TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
+    m.impl("grid", &grid_cpu);
+}
+
+#ifdef WITH_CUDA
+TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
+    m.impl("grid", &grid_cuda);
+}
+TORCH_LIBRARY_IMPL(torch_cluster, HIP, m) {
+    m.impl("grid", &grid_cuda);
+}
+#endif
diff --git a/csrc/knn.cpp b/csrc/knn.cpp
index 149bf5f9..dd69dcd0 100644
--- a/csrc/knn.cpp
+++ b/csrc/knn.cpp
@@ -1,7 +1,8 @@
 #ifdef WITH_PYTHON
 #include <Python.h>
 #endif
-#include <torch/script.h>
+#include <torch/torch.h>
+#include <torch/library.h>
 
 #include "cpu/knn_cpu.h"
 
@@ -36,5 +37,29 @@ CLUSTER_API torch::Tensor knn(torch::Tensor x, torch::Tensor y,
   }
 }
 
-static auto registry =
-    torch::RegisterOperators().op("torch_cluster::knn", &knn);
+inline torch::Tensor knn_cpu_wrap(torch::Tensor x, torch::Tensor y,
+                  std::optional<torch::Tensor> ptr_x,
+                  std::optional<torch::Tensor> ptr_y, int64_t k, bool cosine,
+                  int64_t num_workers = 1) {
+    TORCH_CHECK(!cosine, "`cosine` argument not supported on CPU");
+    return knn_cpu(x, y, ptr_x, ptr_y, k, num_workers);
+}
+
+TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
+  m.impl("knn", &knn_cpu_wrap);
+}
+
+#ifdef WITH_CUDA
+inline torch::Tensor knn_cuda_wrap(torch::Tensor x, torch::Tensor y,
+                  std::optional<torch::Tensor> ptr_x,
+                  std::optional<torch::Tensor> ptr_y, int64_t k, bool cosine,
+                  int64_t num_workers) {
+    return knn_cuda(x, y, ptr_x, ptr_y, k, cosine);
+}
+TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
+    m.impl("knn", &knn_cuda_wrap);
+}
+TORCH_LIBRARY_IMPL(torch_cluster, HIP, m) {
+    m.impl("knn", &knn_cuda_wrap);
+}
+#endif
diff --git a/csrc/nearest.cpp b/csrc/nearest.cpp
index 83eb41e9..e752d1a8 100644
--- a/csrc/nearest.cpp
+++ b/csrc/nearest.cpp
@@ -1,7 +1,8 @@
 #ifdef WITH_PYTHON
 #include <Python.h>
 #endif
-#include <torch/script.h>
+#include <torch/torch.h>
+#include <torch/library.h>
 
 #include "extensions.h"
 
@@ -32,5 +33,11 @@ CLUSTER_API torch::Tensor nearest(torch::Tensor x, torch::Tensor y, torch::Tenso
   }
 }
 
-static auto registry =
-    torch::RegisterOperators().op("torch_cluster::nearest", &nearest);
+#ifdef WITH_CUDA
+TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
+    m.impl("nearest", &nearest_cuda);
+}
+TORCH_LIBRARY_IMPL(torch_cluster, HIP, m) {
+  m.impl("nearest", &nearest_cuda);
+}
+#endif
diff --git a/csrc/radius.cpp b/csrc/radius.cpp
index 1ab88844..d8136716 100644
--- a/csrc/radius.cpp
+++ b/csrc/radius.cpp
@@ -1,7 +1,8 @@
 #ifdef WITH_PYTHON
 #include <Python.h>
 #endif
-#include <torch/script.h>
+#include <torch/torch.h>
+#include <torch/library.h>
 
 #include "cpu/radius_cpu.h"
 
@@ -35,5 +36,23 @@ CLUSTER_API torch::Tensor radius(torch::Tensor x, torch::Tensor y,
   }
 }
 
-static auto registry =
-    torch::RegisterOperators().op("torch_cluster::radius", &radius);
+TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
+  m.impl("radius", &radius_cpu);
+}
+
+#ifdef WITH_CUDA
+inline torch::Tensor radius_cuda_wrap(torch::Tensor x, torch::Tensor y,
+                     std::optional<torch::Tensor> ptr_x,
+                     std::optional<torch::Tensor> ptr_y, double r,
+                     int64_t max_num_neighbors, int64_t num_workers,
+                     bool ignore_same_index) {
+    return radius_cuda(x, y, ptr_x, ptr_y, r, max_num_neighbors, ignore_same_index);
+}
+
+TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
+  m.impl("radius", &radius_cuda_wrap);
+}
+TORCH_LIBRARY_IMPL(torch_cluster, HIP, m) {
+  m.impl("radius", &radius_cuda_wrap);
+}
+#endif
diff --git a/csrc/rw.cpp b/csrc/rw.cpp
index 0f8de62d..184a40f1 100644
--- a/csrc/rw.cpp
+++ b/csrc/rw.cpp
@@ -1,7 +1,8 @@
 #ifdef WITH_PYTHON
 #include <Python.h>
 #endif
-#include <torch/script.h>
+#include <torch/torch.h>
+#include <torch/library.h>
 
 #include "cpu/rw_cpu.h"
 
@@ -33,5 +34,16 @@ random_walk(torch::Tensor rowptr, torch::Tensor col, torch::Tensor start,
   }
 }
 
-static auto registry =
-    torch::RegisterOperators().op("torch_cluster::random_walk", &random_walk);
+TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
+  m.impl("random_walk", &random_walk_cpu);
+}
+
+#ifdef WITH_CUDA
+TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
+  m.impl("random_walk", &random_walk_cuda);
+}
+
+TORCH_LIBRARY_IMPL(torch_cluster, HIP, m) {
+  m.impl("random_walk", &random_walk_cuda);
+}
+#endif
diff --git a/csrc/sampler.cpp b/csrc/sampler.cpp
index d0b3be8d..df25144e 100644
--- a/csrc/sampler.cpp
+++ b/csrc/sampler.cpp
@@ -1,7 +1,8 @@
 #ifdef WITH_PYTHON
 #include <Python.h>
 #endif
-#include <torch/script.h>
+#include <torch/torch.h>
+#include <torch/library.h>
 
 #include "cpu/sampler_cpu.h"
 
@@ -28,5 +29,6 @@ CLUSTER_API torch::Tensor neighbor_sampler(torch::Tensor start, torch::Tensor ro
   }
 }
 
-static auto registry = torch::RegisterOperators().op(
-    "torch_cluster::neighbor_sampler", &neighbor_sampler);
+TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
+  m.impl("neighbor_sampler", &neighbor_sampler_cpu);
+}
diff --git a/csrc/version.cpp b/csrc/version.cpp
index 84395d82..cfca9b23 100644
--- a/csrc/version.cpp
+++ b/csrc/version.cpp
@@ -1,9 +1,11 @@
 #ifdef WITH_PYTHON
 #include <Python.h>
 #endif
+
+#include <torch/library.h>
+
 #include "cluster.h"
 #include "macros.h"
-#include <torch/script.h>
 
 #ifdef WITH_CUDA
 #ifdef USE_ROCM
@@ -37,5 +39,6 @@ CLUSTER_API int64_t cuda_version() noexcept {
 }
 } // namespace cluster
 
-static auto registry = torch::RegisterOperators().op(
-    "torch_cluster::cuda_version", [] { return cluster::cuda_version(); });
+TORCH_LIBRARY_IMPL(torch_cluster, CompositeImplicitAutograd, m) {
+  m.impl("cuda_version", [] { return cluster::cuda_version(); });
+}
diff --git a/setup.py b/setup.py
index 1ef90ab9..9375ce57 100644
--- a/setup.py
+++ b/setup.py
@@ -11,7 +11,7 @@
 from torch.utils.cpp_extension import (CUDA_HOME, BuildExtension, CppExtension,
                                        CUDAExtension)
 
-__version__ = '1.6.3'
+__version__ = '2.0.0'
 URL = 'https://github.com/rusty1s/pytorch_cluster'
 
 WITH_CUDA = False
@@ -25,7 +25,6 @@
     suffices = ['cuda']
 if os.getenv('FORCE_ONLY_CPU', '0') == '1':
     suffices = ['cpu']
-
 BUILD_DOCS = os.getenv('BUILD_DOCS', '0') == '1'
 
 
diff --git a/test/test_fps.py b/test/test_fps.py
index 0f10e51a..1c8e997a 100644
--- a/test/test_fps.py
+++ b/test/test_fps.py
@@ -4,14 +4,24 @@
 import torch
 from torch import Tensor
 from torch_cluster import fps
-from torch_cluster.testing import devices, grad_dtypes, tensor
+from torch_cluster.testing import (
+    devices,
+    grad_dtypes,
+    has_compiler,
+    tensor,
+)
 
 
-@torch.jit.script
-def fps2(x: Tensor, ratio: Tensor) -> Tensor:
+def fps2_impl(x: Tensor, ratio: Tensor) -> Tensor:
     return fps(x, None, ratio, False)
 
 
+if has_compiler():
+    fps2 = torch.compile(fps2_impl)
+else:
+    fps2 = fps2_impl
+
+
 @pytest.mark.parametrize('dtype,device', product(grad_dtypes, devices))
 def test_fps(dtype, device):
     x = tensor([
diff --git a/test/test_graclus.py b/test/test_graclus.py
index c8e1f39f..08704842 100644
--- a/test/test_graclus.py
+++ b/test/test_graclus.py
@@ -51,6 +51,4 @@ def test_graclus_cluster(test, dtype, device):
     cluster = graclus_cluster(row, col, weight)
     assert_correct(row, col, cluster)
 
-    jit = torch.jit.script(graclus_cluster)
-    cluster = jit(row, col, weight)
-    assert_correct(row, col, cluster)
+    # graclus is not torch.compile compatible.
diff --git a/test/test_grid.py b/test/test_grid.py
index 2d53220f..d7b09615 100644
--- a/test/test_grid.py
+++ b/test/test_grid.py
@@ -3,7 +3,7 @@
 import pytest
 import torch
 from torch_cluster import grid_cluster
-from torch_cluster.testing import devices, dtypes, tensor
+from torch_cluster.testing import devices, dtypes, has_compiler, tensor
 
 tests = [{
     'pos': [2, 6],
@@ -39,5 +39,6 @@ def test_grid_cluster(test, dtype, device):
     cluster = grid_cluster(pos, size, start, end)
     assert cluster.tolist() == test['cluster']
 
-    jit = torch.jit.script(grid_cluster)
-    assert torch.equal(jit(pos, size, start, end), cluster)
+    if has_compiler():
+        jit = torch.compile(grid_cluster)
+        assert torch.equal(jit(pos, size, start, end), cluster)
diff --git a/test/test_knn.py b/test/test_knn.py
index 32852fe0..855536a5 100644
--- a/test/test_knn.py
+++ b/test/test_knn.py
@@ -1,18 +1,33 @@
+import importlib.util
 from itertools import product
 
 import pytest
 import scipy.spatial
 import torch
 from torch_cluster import knn, knn_graph
-from torch_cluster.testing import devices, grad_dtypes, tensor
+from torch_cluster.testing import (
+    devices,
+    grad_dtypes,
+    has_compiler,
+    tensor,
+    triton_wrap,
+)
+
+HAS_CUDA = torch.cuda.is_available()
+HAS_TRITON = importlib.util.find_spec('triton') is not None
+
+torch._dynamo.config.capture_scalar_outputs = True
 
 
 def to_set(edge_index):
     return set([(i, j) for i, j in edge_index.t().tolist()])
 
 
-@pytest.mark.parametrize('dtype,device', product(grad_dtypes, devices))
-def test_knn(dtype, device):
+@pytest.mark.parametrize(
+    'dtype,device,use_triton',
+    triton_wrap(product(grad_dtypes, devices)),
+)
+def test_knn(dtype, device, use_triton):
     x = tensor([
         [-1, -1],
         [-1, +1],
@@ -31,29 +46,41 @@ def test_knn(dtype, device):
     batch_x = tensor([0, 0, 0, 0, 1, 1, 1, 1], torch.long, device)
     batch_y = tensor([0, 1], torch.long, device)
 
-    edge_index = knn(x, y, 2)
-    assert to_set(edge_index) == set([(0, 2), (0, 3), (1, 0), (1, 1)])
+    edge_index = knn(x, y, 2, use_triton=use_triton)
+    assert to_set(edge_index) == {(0, 2), (0, 3), (1, 0), (1, 1)}
 
-    jit = torch.jit.script(knn)
-    edge_index = jit(x, y, 2)
-    assert to_set(edge_index) == set([(0, 2), (0, 3), (1, 0), (1, 1)])
+    if has_compiler():
+        jit = torch.compile(knn)
+        edge_index = jit(x, y, 2, use_triton=use_triton)
+        assert to_set(edge_index) == {(0, 2), (0, 3), (1, 0), (1, 1)}
 
-    edge_index = knn(x, y, 2, batch_x, batch_y)
-    assert to_set(edge_index) == set([(0, 2), (0, 3), (1, 4), (1, 5)])
+    edge_index = knn(x, y, 2, batch_x, batch_y, use_triton=use_triton)
+    assert to_set(edge_index) == {(0, 2), (0, 3), (1, 4), (1, 5)}
 
     if x.is_cuda:
-        edge_index = knn(x, y, 2, batch_x, batch_y, cosine=True)
-        assert to_set(edge_index) == set([(0, 2), (0, 3), (1, 4), (1, 5)])
+        edge_index = knn(
+            x,
+            y,
+            2,
+            batch_x,
+            batch_y,
+            cosine=True,
+            use_triton=use_triton,
+        )
+        assert to_set(edge_index) == {(0, 2), (0, 3), (1, 4), (1, 5)}
 
     # Skipping a batch
     batch_x = tensor([0, 0, 0, 0, 2, 2, 2, 2], torch.long, device)
     batch_y = tensor([0, 2], torch.long, device)
-    edge_index = knn(x, y, 2, batch_x, batch_y)
-    assert to_set(edge_index) == set([(0, 2), (0, 3), (1, 4), (1, 5)])
+    edge_index = knn(x, y, 2, batch_x, batch_y, use_triton=use_triton)
+    assert to_set(edge_index) == {(0, 2), (0, 3), (1, 4), (1, 5)}
 
 
-@pytest.mark.parametrize('dtype,device', product(grad_dtypes, devices))
-def test_knn_graph(dtype, device):
+@pytest.mark.parametrize(
+    'dtype,device,use_triton',
+    triton_wrap(product(grad_dtypes, devices)),
+)
+def test_knn_graph(dtype, device, use_triton):
     x = tensor([
         [-1, -1],
         [-1, +1],
@@ -61,28 +88,137 @@ def test_knn_graph(dtype, device):
         [+1, -1],
     ], dtype, device)
 
-    edge_index = knn_graph(x, k=2, flow='target_to_source')
-    assert to_set(edge_index) == set([(0, 1), (0, 3), (1, 0), (1, 2), (2, 1),
-                                      (2, 3), (3, 0), (3, 2)])
-
-    edge_index = knn_graph(x, k=2, flow='source_to_target')
-    assert to_set(edge_index) == set([(1, 0), (3, 0), (0, 1), (2, 1), (1, 2),
-                                      (3, 2), (0, 3), (2, 3)])
-
-    jit = torch.jit.script(knn_graph)
-    edge_index = jit(x, k=2, flow='source_to_target')
-    assert to_set(edge_index) == set([(1, 0), (3, 0), (0, 1), (2, 1), (1, 2),
-                                      (3, 2), (0, 3), (2, 3)])
-
-
-@pytest.mark.parametrize('dtype,device', product([torch.float], devices))
-def test_knn_graph_large(dtype, device):
+    edge_index = knn_graph(
+        x,
+        k=2,
+        flow='target_to_source',
+        use_triton=use_triton,
+    )
+    assert to_set(edge_index) == {
+        (0, 1),
+        (0, 3),
+        (1, 0),
+        (1, 2),
+        (2, 1),
+        (2, 3),
+        (3, 0),
+        (3, 2),
+    }
+
+    edge_index = knn_graph(
+        x,
+        k=2,
+        flow='source_to_target',
+        use_triton=use_triton,
+    )
+    assert to_set(edge_index) == {
+        (1, 0),
+        (3, 0),
+        (0, 1),
+        (2, 1),
+        (1, 2),
+        (3, 2),
+        (0, 3),
+        (2, 3),
+    }
+
+    if has_compiler():
+        jit = torch.compile(knn_graph)
+        edge_index = jit(x, k=2, flow='source_to_target')
+        assert to_set(edge_index) == {
+            (1, 0),
+            (3, 0),
+            (0, 1),
+            (2, 1),
+            (1, 2),
+            (3, 2),
+            (0, 3),
+            (2, 3),
+        }
+
+
+@pytest.mark.parametrize(
+    'dtype,device,use_triton',
+    triton_wrap(product([torch.float], devices)),
+)
+def test_knn_graph_large(dtype, device, use_triton):
+    torch.manual_seed(29)
     x = torch.randn(1000, 3, dtype=dtype, device=device)
 
-    edge_index = knn_graph(x, k=5, flow='target_to_source', loop=True)
+    edge_index = knn_graph(
+        x,
+        k=5,
+        flow='target_to_source',
+        loop=True,
+        use_triton=use_triton,
+    )
 
     tree = scipy.spatial.cKDTree(x.cpu().numpy())
     _, col = tree.query(x.cpu(), k=5)
     truth = set([(i, j) for i, ns in enumerate(col) for j in ns])
 
     assert to_set(edge_index.cpu()) == truth
+
+
+@pytest.mark.skipif(
+    not (HAS_CUDA and HAS_TRITON),
+    reason='CUDA and Triton are required for Triton parity tests.',
+)
+def test_knn_triton_matches_cuda():
+    torch.manual_seed(42)
+    x = torch.randn(128, 16, device='cuda')
+    y = torch.randn(64, 16, device='cuda')
+    batch_x = torch.zeros(x.size(0), dtype=torch.long, device='cuda')
+    batch_y = torch.zeros(y.size(0), dtype=torch.long, device='cuda')
+
+    out_cuda = knn(
+        x,
+        y,
+        k=8,
+        batch_x=batch_x,
+        batch_y=batch_y,
+        use_triton=False,
+    )
+    out_triton = knn(
+        x,
+        y,
+        k=8,
+        batch_x=batch_x,
+        batch_y=batch_y,
+        use_triton=True,
+    )
+    assert to_set(out_cuda) == to_set(out_triton)
+
+    out_cuda = knn(
+        x,
+        y,
+        k=8,
+        batch_x=batch_x,
+        batch_y=batch_y,
+        cosine=True,
+        use_triton=False,
+    )
+    out_triton = knn(
+        x,
+        y,
+        k=8,
+        batch_x=batch_x,
+        batch_y=batch_y,
+        cosine=True,
+        use_triton=True,
+    )
+    assert to_set(out_cuda) == to_set(out_triton)
+
+
+@pytest.mark.skipif(
+    not (HAS_CUDA and HAS_TRITON),
+    reason='CUDA and Triton are required for Triton parity tests.',
+)
+def test_knn_graph_triton_matches_cuda():
+    torch.manual_seed(1)
+    x = torch.randn(64, 8, device='cuda')
+    batch = torch.zeros(x.size(0), dtype=torch.long, device='cuda')
+
+    out_cuda = knn_graph(x, k=4, batch=batch, loop=False, use_triton=False)
+    out_triton = knn_graph(x, k=4, batch=batch, loop=False, use_triton=True)
+    assert to_set(out_cuda) == to_set(out_triton)
diff --git a/test/test_nearest.py b/test/test_nearest.py
index 582818e1..b1393113 100644
--- a/test/test_nearest.py
+++ b/test/test_nearest.py
@@ -1,13 +1,45 @@
+import importlib.util
 from itertools import product
 
 import pytest
 import torch
 from torch_cluster import nearest
-from torch_cluster.testing import devices, grad_dtypes, tensor
+from torch_cluster.testing import devices, grad_dtypes, tensor, triton_wrap
 
 
-@pytest.mark.parametrize('dtype,device', product(grad_dtypes, devices))
-def test_nearest(dtype, device):
+@pytest.mark.skipif(
+    not (
+        torch.cuda.is_available()
+        and importlib.util.find_spec('triton') is not None
+    ),
+    reason='CUDA and Triton are required for Triton parity tests.',
+)
+def test_nearest_triton_matches_cuda():
+    torch.manual_seed(123)
+    x = torch.randn(128, 8, device='cuda')
+    y = torch.randn(32, 8, device='cuda')
+    batch_x = torch.zeros(x.size(0), dtype=torch.long, device='cuda')
+    batch_y = torch.zeros(y.size(0), dtype=torch.long, device='cuda')
+
+    out_cuda = nearest(x, y, batch_x, batch_y, use_triton=False)
+    out_triton = nearest(x, y, batch_x, batch_y, use_triton=True)
+    assert torch.equal(out_cuda, out_triton)
+
+    batch_x = torch.zeros(x.size(0), dtype=torch.long, device='cuda')
+    batch_y = torch.zeros(y.size(0), dtype=torch.long, device='cuda')
+    batch_x[x.size(0) // 2:] = 1
+    batch_y[y.size(0) // 2:] = 1
+
+    out_cuda = nearest(x, y, batch_x, batch_y, use_triton=False)
+    out_triton = nearest(x, y, batch_x, batch_y, use_triton=True)
+    assert torch.equal(out_cuda, out_triton)
+
+
+@pytest.mark.parametrize(
+    'dtype,device,use_triton',
+    triton_wrap(product(grad_dtypes, devices)),
+)
+def test_nearest(dtype, device, use_triton):
     x = tensor([
         [-1, -1],
         [-1, +1],
@@ -28,7 +60,7 @@ def test_nearest(dtype, device):
     batch_x = tensor([0, 0, 0, 0, 1, 1, 1, 1], torch.long, device)
     batch_y = tensor([0, 0, 1, 1], torch.long, device)
 
-    out = nearest(x, y, batch_x, batch_y)
+    out = nearest(x, y, batch_x, batch_y, use_triton=use_triton)
     assert out.tolist() == [0, 0, 1, 1, 2, 2, 3, 3]
 
     out = nearest(x, y)
@@ -38,27 +70,27 @@ def test_nearest(dtype, device):
     batch_x = tensor([0, 0, 0, 0, 1, 1, 1, 1], torch.long, device)
     batch_y = tensor([0, 0, 0, 0], torch.long, device)
     with pytest.raises(ValueError):
-        nearest(x, y, batch_x, batch_y)
+        nearest(x, y, batch_x, batch_y, use_triton=use_triton)
 
     # Invalid input: instance 1 only in batch_x (implicitly as batch_y=None)
     with pytest.raises(ValueError):
-        nearest(x, y, batch_x, batch_y=None)
+        nearest(x, y, batch_x, batch_y=None, use_triton=use_triton)
 
     # Invalid input: instance 2 only in batch_x
     # (i.e.instance in the middle missing)
     batch_x = tensor([0, 0, 1, 1, 2, 2, 3, 3], torch.long, device)
     batch_y = tensor([0, 1, 3, 3], torch.long, device)
     with pytest.raises(ValueError):
-        nearest(x, y, batch_x, batch_y)
+        nearest(x, y, batch_x, batch_y, use_triton=use_triton)
 
     # Invalid input: batch_x unsorted
     batch_x = tensor([0, 0, 1, 0, 0, 0, 0], torch.long, device)
     batch_y = tensor([0, 0, 1, 1], torch.long, device)
     with pytest.raises(ValueError):
-        nearest(x, y, batch_x, batch_y)
+        nearest(x, y, batch_x, batch_y, use_triton=use_triton)
 
     # Invalid input: batch_y unsorted
     batch_x = tensor([0, 0, 0, 0, 1, 1, 1, 1], torch.long, device)
     batch_y = tensor([0, 0, 1, 0], torch.long, device)
     with pytest.raises(ValueError):
-        nearest(x, y, batch_x, batch_y)
+        nearest(x, y, batch_x, batch_y, use_triton=use_triton)
diff --git a/test/test_radius.py b/test/test_radius.py
index 4289bfc6..e4f2a7e3 100644
--- a/test/test_radius.py
+++ b/test/test_radius.py
@@ -4,7 +4,13 @@
 import scipy.spatial
 import torch
 from torch_cluster import radius, radius_graph
-from torch_cluster.testing import devices, floating_dtypes, tensor
+from torch_cluster.testing import (
+    devices,
+    floating_dtypes,
+    has_compiler,
+    tensor,
+    triton_wrap,
+)
 
 
 def to_set(edge_index):
@@ -20,8 +26,11 @@ def to_batch(nodes):
     return [int(i / 4) for i in nodes]
 
 
-@pytest.mark.parametrize('dtype,device', product(floating_dtypes, devices))
-def test_radius(dtype, device):
+@pytest.mark.parametrize(
+    'dtype,device,use_triton',
+    triton_wrap(product(floating_dtypes, devices)),
+)
+def test_radius(dtype, device, use_triton):
     x = tensor([
         [-1, -1],
         [-1, +1],
@@ -40,29 +49,57 @@ def test_radius(dtype, device):
     batch_x = tensor([0, 0, 0, 0, 1, 1, 1, 1], torch.long, device)
     batch_y = tensor([0, 1], torch.long, device)
 
-    edge_index = radius(x, y, 2, max_num_neighbors=4)
+    edge_index = radius(x, y, 2, max_num_neighbors=4, use_triton=use_triton)
     assert to_set(edge_index) == set([(0, 0), (0, 1), (0, 2), (0, 3), (1, 1),
                                       (1, 2), (1, 5), (1, 6)])
 
-    jit = torch.jit.script(radius)
-    edge_index = jit(x, y, 2, max_num_neighbors=4)
-    assert to_set(edge_index) == set([(0, 0), (0, 1), (0, 2), (0, 3), (1, 1),
-                                      (1, 2), (1, 5), (1, 6)])
-
-    edge_index = radius(x, y, 2, batch_x, batch_y, max_num_neighbors=4)
+    if has_compiler():
+        jit = torch.compile(radius)
+        edge_index = jit(x, y, 2, max_num_neighbors=4, use_triton=use_triton)
+        assert to_set(edge_index) == set([
+            (0, 0),
+            (0, 1),
+            (0, 2),
+            (0, 3),
+            (1, 1),
+            (1, 2),
+            (1, 5),
+            (1, 6),
+        ])
+
+    edge_index = radius(
+        x,
+        y,
+        2,
+        batch_x,
+        batch_y,
+        max_num_neighbors=4,
+        use_triton=use_triton,
+    )
     assert to_set(edge_index) == set([(0, 0), (0, 1), (0, 2), (0, 3), (1, 5),
                                       (1, 6)])
 
     # Skipping a batch
     batch_x = tensor([0, 0, 0, 0, 2, 2, 2, 2], torch.long, device)
     batch_y = tensor([0, 2], torch.long, device)
-    edge_index = radius(x, y, 2, batch_x, batch_y, max_num_neighbors=4)
+    edge_index = radius(
+        x,
+        y,
+        2,
+        batch_x,
+        batch_y,
+        max_num_neighbors=4,
+        use_triton=use_triton,
+    )
     assert to_set(edge_index) == set([(0, 0), (0, 1), (0, 2), (0, 3), (1, 5),
                                       (1, 6)])
 
 
-@pytest.mark.parametrize('dtype,device', product(floating_dtypes, devices))
-def test_radius_graph(dtype, device):
+@pytest.mark.parametrize(
+    'dtype,device,use_triton',
+    triton_wrap(product(floating_dtypes, devices)),
+)
+def test_radius_graph(dtype, device, use_triton):
     x = tensor([
         [-1, -1],
         [-1, +1],
@@ -70,21 +107,50 @@ def test_radius_graph(dtype, device):
         [+1, -1],
     ], dtype, device)
 
-    edge_index = radius_graph(x, r=2.5, flow='target_to_source')
+    edge_index = radius_graph(
+        x,
+        r=2.5,
+        flow='target_to_source',
+        use_triton=use_triton,
+    )
     assert to_set(edge_index) == set([(0, 1), (0, 3), (1, 0), (1, 2), (2, 1),
                                       (2, 3), (3, 0), (3, 2)])
 
-    edge_index = radius_graph(x, r=2.5, flow='source_to_target')
-    assert to_set(edge_index) == set([(1, 0), (3, 0), (0, 1), (2, 1), (1, 2),
-                                      (3, 2), (0, 3), (2, 3)])
-
-    jit = torch.jit.script(radius_graph)
-    edge_index = jit(x, r=2.5, flow='source_to_target')
+    edge_index = radius_graph(
+        x,
+        r=2.5,
+        flow='source_to_target',
+        use_triton=use_triton,
+    )
     assert to_set(edge_index) == set([(1, 0), (3, 0), (0, 1), (2, 1), (1, 2),
                                       (3, 2), (0, 3), (2, 3)])
 
-    edge_index = radius_graph(x, r=100, flow='source_to_target',
-                              max_num_neighbors=1)
+    if has_compiler():
+        jit = torch.compile(radius_graph)
+        edge_index = jit(
+            x,
+            r=2.5,
+            flow='source_to_target',
+            use_triton=use_triton,
+        )
+        assert to_set(edge_index) == set([
+            (1, 0),
+            (3, 0),
+            (0, 1),
+            (2, 1),
+            (1, 2),
+            (3, 2),
+            (0, 3),
+            (2, 3),
+        ])
+
+    edge_index = radius_graph(
+        x,
+        r=100,
+        flow='source_to_target',
+        max_num_neighbors=1,
+        use_triton=use_triton,
+    )
     assert set(to_degree(edge_index)) == set([1])
 
     x = tensor([
@@ -94,8 +160,13 @@ def test_radius_graph(dtype, device):
         [-1, -1],
     ], dtype, device)
 
-    edge_index = radius_graph(x, r=100, flow='source_to_target',
-                              max_num_neighbors=1)
+    edge_index = radius_graph(
+        x,
+        r=100,
+        flow='source_to_target',
+        max_num_neighbors=1,
+        use_triton=use_triton,
+    )
     assert set(to_degree(edge_index)) == set([1])
 
     x = tensor([
@@ -110,21 +181,33 @@ def test_radius_graph(dtype, device):
     ], dtype, device)
     batch_x = tensor([0, 0, 0, 0, 1, 1, 1, 1], torch.long, device)
 
-    edge_index = radius_graph(x, r=100, batch=batch_x, flow='source_to_target',
-                              max_num_neighbors=1)
+    edge_index = radius_graph(
+        x,
+        r=100,
+        batch=batch_x,
+        flow='source_to_target',
+        max_num_neighbors=1,
+        use_triton=use_triton,
+    )
     assert set(to_degree(edge_index)) == set([1])
     assert to_batch(edge_index[0]) == batch_x.tolist()
 
 
-@pytest.mark.parametrize('dtype,device', product([torch.float], devices))
-def test_radius_graph_large(dtype, device):
+@pytest.mark.parametrize(
+    'dtype,device,use_triton',
+    triton_wrap(product([torch.float], devices)),
+)
+def test_radius_graph_large(dtype, device, use_triton):
     x = torch.randn(1000, 3, dtype=dtype, device=device)
 
-    edge_index = radius_graph(x,
-                              r=0.5,
-                              flow='target_to_source',
-                              loop=True,
-                              max_num_neighbors=2000)
+    edge_index = radius_graph(
+        x,
+        r=0.5,
+        flow='target_to_source',
+        loop=True,
+        max_num_neighbors=2000,
+        use_triton=use_triton,
+    )
 
     tree = scipy.spatial.cKDTree(x.cpu().numpy())
     col = tree.query_ball_point(x.cpu(), r=0.5)
diff --git a/test/test_rw.py b/test/test_rw.py
index 82a8b77d..80ad2cef 100644
--- a/test/test_rw.py
+++ b/test/test_rw.py
@@ -1,7 +1,7 @@
 import pytest
 import torch
 from torch_cluster import random_walk
-from torch_cluster.testing import devices, tensor
+from torch_cluster.testing import devices, has_compiler, tensor
 
 
 @pytest.mark.parametrize('device', devices)
@@ -31,8 +31,12 @@ def test_rw_small(device):
     out = random_walk(row, col, start, walk_length, num_nodes=3)
     assert out.tolist() == [[0, 1, 0, 1, 0], [1, 0, 1, 0, 1], [2, 2, 2, 2, 2]]
 
-    jit = torch.jit.script(random_walk)
-    assert torch.equal(jit(row, col, start, walk_length, num_nodes=3), out)
+    if has_compiler():
+        jit = torch.compile(random_walk)
+        assert torch.equal(
+            jit(row, col, start, walk_length, num_nodes=3),
+            out,
+        )
 
 
 @pytest.mark.parametrize('device', devices)
diff --git a/torch_cluster/fps.py b/torch_cluster/fps.py
index 7baf9813..a18261c0 100644
--- a/torch_cluster/fps.py
+++ b/torch_cluster/fps.py
@@ -6,28 +6,14 @@
 import torch_cluster.typing
 
 
-@torch.jit._overload  # noqa
-def fps(src, batch, ratio, random_start, batch_size, ptr):  # noqa
-    # type: (Tensor, Optional[Tensor], Optional[float], bool, Optional[int], Optional[Tensor]) -> Tensor  # noqa
-    pass  # pragma: no cover
-
-
-@torch.jit._overload  # noqa
-def fps(src, batch, ratio, random_start, batch_size, ptr):  # noqa
-    # type: (Tensor, Optional[Tensor], Optional[Tensor], bool, Optional[int], Optional[Tensor]) -> Tensor  # noqa
-    pass  # pragma: no cover
-
-
-@torch.jit._overload  # noqa
-def fps(src, batch, ratio, random_start, batch_size, ptr):  # noqa
-    # type: (Tensor, Optional[Tensor], Optional[float], bool, Optional[int], Optional[List[int]]) -> Tensor  # noqa
-    pass  # pragma: no cover
-
-
-@torch.jit._overload  # noqa
-def fps(src, batch, ratio, random_start, batch_size, ptr):  # noqa
-    # type: (Tensor, Optional[Tensor], Optional[Tensor], bool, Optional[int], Optional[List[int]]) -> Tensor  # noqa
-    pass  # pragma: no cover
+@torch.library.register_fake("torch_cluster::fps")
+def _(src, ptr, ratio, random_start=True):
+    torch._check(src.device == ptr.device)
+    torch._check(ptr.ndim == 1)
+
+    ctx = torch.library.get_ctx()
+    nnz = ctx.new_dynamic_size()
+    return ptr.new_empty((nnz,))
 
 
 def fps(  # noqa
diff --git a/torch_cluster/graclus.py b/torch_cluster/graclus.py
index 7fa834d6..90317568 100644
--- a/torch_cluster/graclus.py
+++ b/torch_cluster/graclus.py
@@ -3,6 +3,17 @@
 import torch
 
 
+@torch.library.register_fake("torch_cluster::graclus")
+def _(rowptr, col, weight=None):
+    torch._check(rowptr.device == col.device)
+    torch._check(rowptr.ndim == 1)
+    torch._check(col.ndim == 1)
+    if weight is not None:
+        torch._check(weight.device == col.device)
+        torch._check(weight.ndim == 1)
+    return rowptr.new_empty((rowptr.numel() - 1,), dtype=torch.long)
+
+
 def graclus_cluster(
     row: torch.Tensor,
     col: torch.Tensor,
diff --git a/torch_cluster/grid.py b/torch_cluster/grid.py
index da59d51e..c4ebdaf4 100644
--- a/torch_cluster/grid.py
+++ b/torch_cluster/grid.py
@@ -3,6 +3,16 @@
 import torch
 
 
+@torch.library.register_fake("torch_cluster::grid")
+def _(pos, size, start=None, end=None):
+    torch._check(pos.device == size.device)
+    if start is not None:
+        torch._check(start.device == pos.device)
+    if end is not None:
+        torch._check(end.device == pos.device)
+    return pos.new_empty((pos.size(0),), dtype=torch.long)
+
+
 def grid_cluster(
     pos: torch.Tensor,
     size: torch.Tensor,
diff --git a/torch_cluster/knn.py b/torch_cluster/knn.py
index cf8f0875..472fc8ed 100644
--- a/torch_cluster/knn.py
+++ b/torch_cluster/knn.py
@@ -1,8 +1,23 @@
+import importlib.util
 from typing import Optional
 
 import torch
 
 
+@torch.library.register_fake("torch_cluster::knn")
+def _(x, y, batch_x, batch_y, k, cosine=False, num_workers=1):
+    torch._check(x.device == y.device)
+    if batch_x is not None:
+        torch._check(x.device == batch_x.device)
+        torch._check(batch_x.ndim == 1)
+    if batch_y is not None:
+        torch._check(y.device == batch_y.device)
+        torch._check(batch_y.ndim == 1)
+    ctx = torch.library.get_ctx()
+    nnz = ctx.new_dynamic_size()
+    return x.new_empty((2, nnz), dtype=torch.long)
+
+
 def knn(
     x: torch.Tensor,
     y: torch.Tensor,
@@ -12,6 +27,7 @@ def knn(
     cosine: bool = False,
     num_workers: int = 1,
     batch_size: Optional[int] = None,
+    use_triton: bool = False,
 ) -> torch.Tensor:
     r"""Finds for each element in :obj:`y` the :obj:`k` nearest points in
     :obj:`x`.
@@ -44,6 +60,7 @@ def knn(
     .. code-block:: python
 
         import torch
+
         from torch_cluster import knn
 
         x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
@@ -69,6 +86,25 @@ def knn(
             batch_size = max(batch_size, int(batch_y.max()) + 1)
     assert batch_size > 0
 
+    if (use_triton and x.is_cuda and y.is_cuda
+            and x.dtype is not torch.float64 and y.dtype is not torch.float64):
+        if importlib.util.find_spec("triton") is None:
+            print(
+                "Triton is not available. Falling back to general "
+                "implementation."
+            )
+        else:
+            from .triton.knn import knn as triton_knn
+            return triton_knn(
+                x,
+                y,
+                k,
+                batch_x,
+                batch_y,
+                cosine,
+                batch_size,
+            )
+
     ptr_x: Optional[torch.Tensor] = None
     ptr_y: Optional[torch.Tensor] = None
     if batch_size > 1:
@@ -78,8 +114,15 @@ def knn(
         ptr_x = torch.bucketize(arange, batch_x)
         ptr_y = torch.bucketize(arange, batch_y)
 
-    return torch.ops.torch_cluster.knn(x, y, ptr_x, ptr_y, k, cosine,
-                                       num_workers)
+    return torch.ops.torch_cluster.knn(
+        x,
+        y,
+        ptr_x,
+        ptr_y,
+        k,
+        cosine,
+        num_workers,
+    )
 
 
 def knn_graph(
@@ -91,6 +134,7 @@ def knn_graph(
     cosine: bool = False,
     num_workers: int = 1,
     batch_size: Optional[int] = None,
+    use_triton: bool = False,
 ) -> torch.Tensor:
     r"""Computes graph edges to the nearest :obj:`k` points.
 
@@ -121,24 +165,31 @@ def knn_graph(
     .. code-block:: python
 
         import torch
+
         from torch_cluster import knn_graph
 
         x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
         batch = torch.tensor([0, 0, 0, 0])
         edge_index = knn_graph(x, k=2, batch=batch, loop=False)
     """
-
     assert flow in ['source_to_target', 'target_to_source']
-    edge_index = knn(x, x, k if loop else k + 1, batch, batch, cosine,
-                     num_workers, batch_size)
+    edge_index = knn(
+        x,
+        x,
+        k if loop else k + 1,
+        batch,
+        batch,
+        cosine,
+        num_workers,
+        batch_size,
+        use_triton=use_triton,
+    )
 
     if flow == 'source_to_target':
-        row, col = edge_index[1], edge_index[0]
-    else:
-        row, col = edge_index[0], edge_index[1]
+        edge_index = edge_index.flip(0)
 
     if not loop:
-        mask = row != col
-        row, col = row[mask], col[mask]
+        mask = edge_index[0] != edge_index[1]
+        edge_index = edge_index[:, mask]
 
-    return torch.stack([row, col], dim=0)
+    return edge_index.contiguous()
diff --git a/torch_cluster/nearest.py b/torch_cluster/nearest.py
index 1ba4db6e..1e2cfed2 100644
--- a/torch_cluster/nearest.py
+++ b/torch_cluster/nearest.py
@@ -1,14 +1,26 @@
+import importlib.util
 from typing import Optional
 
 import scipy.cluster
 import torch
 
 
+@torch.library.register_fake("torch_cluster::nearest")
+def _(x, y, ptr_x, ptr_y):
+    torch._check(x.device == y.device)
+    torch._check(ptr_x.device == x.device)
+    torch._check(ptr_y.device == y.device)
+    torch._check(ptr_x.ndim == 1)
+    torch._check(ptr_y.ndim == 1)
+    return x.new_empty((y.size(0),), dtype=torch.long)
+
+
 def nearest(
     x: torch.Tensor,
     y: torch.Tensor,
     batch_x: Optional[torch.Tensor] = None,
     batch_y: Optional[torch.Tensor] = None,
+    use_triton: bool = False,
 ) -> torch.Tensor:
     r"""Clusters points in :obj:`x` together which are nearest to a given query
     point in :obj:`y`.
@@ -40,7 +52,6 @@ def nearest(
         batch_y = torch.tensor([0, 0])
         cluster = nearest(x, y, batch_x, batch_y)
     """
-
     x = x.view(-1, 1) if x.dim() == 1 else x
     y = y.view(-1, 1) if y.dim() == 1 else y
     assert x.size(1) == y.size(1)
@@ -51,6 +62,17 @@ def nearest(
         raise ValueError("'batch_y' is not sorted")
 
     if x.is_cuda:
+        if (use_triton and x.dtype is not torch.float64
+                and y.dtype is not torch.float64):
+            if importlib.util.find_spec("triton") is None:
+                print(
+                    "Triton is not available. Falling back to general "
+                    "implementation."
+                )
+            else:
+                from .triton.nearest import nearest as triton_nearest
+                return triton_nearest(x, y, batch_x, batch_y)
+
         if batch_x is not None:
             assert x.size(0) == batch_x.numel()
             batch_size = int(batch_x.max()) + 1
@@ -119,6 +141,10 @@ def nearest(
             x = torch.cat([x, 2 * D * batch_x.view(-1, 1).to(x.dtype)], -1)
             y = torch.cat([y, 2 * D * batch_y.view(-1, 1).to(y.dtype)], -1)
 
+        if x.dtype in (torch.float16, torch.bfloat16):
+            x = x.float()
+        if y.dtype in (torch.float16, torch.bfloat16):
+            y = y.float()
         return torch.from_numpy(
             scipy.cluster.vq.vq(x.detach().cpu(),
                                 y.detach().cpu())[0]).to(torch.long)
diff --git a/torch_cluster/radius.py b/torch_cluster/radius.py
index 92187eb6..ffea2b97 100644
--- a/torch_cluster/radius.py
+++ b/torch_cluster/radius.py
@@ -1,8 +1,57 @@
+import importlib.util
 from typing import Optional
 
 import torch
 
 
+def _cap_neighbors_by_col(
+    edge_index: torch.Tensor,
+    max_num_neighbors: int,
+) -> torch.Tensor:
+    if max_num_neighbors <= 0 or edge_index.numel() == 0:
+        return edge_index
+    row, col = edge_index
+    max_row = int(row.max().item()) if row.numel() > 0 else -1
+    key = row * (max_row + 2) + col
+    order = torch.argsort(key)
+    row = row[order]
+    col = col[order]
+    idx = torch.arange(row.numel(), device=row.device)
+    row_diff = torch.ones_like(row, dtype=torch.bool)
+    if row.numel() > 1:
+        row_diff[1:] = row[1:] != row[:-1]
+    start_idx = torch.where(row_diff, idx, torch.zeros_like(idx))
+    start_idx = torch.cummax(start_idx, 0).values
+    pos = idx - start_idx
+    mask = pos < max_num_neighbors
+    row = row[mask]
+    col = col[mask]
+    return torch.stack([row, col], dim=0)
+
+
+@torch.library.register_fake("torch_cluster::radius")
+def _(
+    x,
+    y,
+    ptr_x,
+    ptr_y,
+    r,
+    max_num_neighbors=32,
+    num_workers=1,
+    ignore_same_index=False,
+):
+    torch._check(x.device == y.device)
+    if ptr_x is not None:
+        torch._check(ptr_x.device == x.device)
+        torch._check(ptr_x.ndim == 1)
+    if ptr_y is not None:
+        torch._check(ptr_y.device == y.device)
+        torch._check(ptr_y.ndim == 1)
+    ctx = torch.library.get_ctx()
+    nnz = ctx.new_dynamic_size()
+    return x.new_empty((2, nnz), dtype=torch.long)
+
+
 def radius(
     x: torch.Tensor,
     y: torch.Tensor,
@@ -12,7 +61,8 @@ def radius(
     max_num_neighbors: int = 32,
     num_workers: int = 1,
     batch_size: Optional[int] = None,
-    ignore_same_index: bool = False
+    ignore_same_index: bool = False,
+    use_triton: bool = False,
 ) -> torch.Tensor:
     r"""Finds for each element in :obj:`y` all points in :obj:`x` within
     distance :obj:`r`.
@@ -44,6 +94,8 @@ def radius(
         ignore_same_index (bool, optional): If :obj:`True`, each element in
             :obj:`y` ignores the point in :obj:`x` with the same index.
             (default: :obj:`False`)
+        use_triton (bool, optional): If :obj:`True`, use Triton kernels when
+            available. (default: :obj:`False`)
 
     .. code-block:: python
 
@@ -73,6 +125,26 @@ def radius(
             batch_size = max(batch_size, int(batch_y.max()) + 1)
     assert batch_size > 0
 
+    if (use_triton and x.is_cuda and y.is_cuda
+            and x.dtype is not torch.float64 and y.dtype is not torch.float64):
+        if importlib.util.find_spec("triton") is None:
+            print(
+                "Triton is not available. Falling back to general "
+                "implementation."
+            )
+        else:
+            from .triton.radius import radius as triton_radius
+            return triton_radius(
+                x,
+                y,
+                r,
+                batch_x,
+                batch_y,
+                max_num_neighbors,
+                batch_size,
+                ignore_same_index,
+            )
+
     ptr_x: Optional[torch.Tensor] = None
     ptr_y: Optional[torch.Tensor] = None
 
@@ -97,6 +169,7 @@ def radius_graph(
     flow: str = 'source_to_target',
     num_workers: int = 1,
     batch_size: Optional[int] = None,
+    use_triton: bool = False,
 ) -> torch.Tensor:
     r"""Computes graph edges to all points within a given distance.
 
@@ -123,6 +196,8 @@ def radius_graph(
             on the GPU. (default: :obj:`1`)
         batch_size (int, optional): The number of examples :math:`B`.
             Automatically calculated if not given. (default: :obj:`None`)
+        use_triton (bool, optional): If :obj:`True`, use Triton kernels when
+            available. (default: :obj:`False`)
 
     :rtype: :class:`LongTensor`
 
@@ -137,12 +212,18 @@ def radius_graph(
     """
 
     assert flow in ['source_to_target', 'target_to_source']
-    edge_index = radius(x, x, r, batch, batch,
-                        max_num_neighbors,
-                        num_workers, batch_size, not loop)
+    edge_index = radius(
+        x,
+        x,
+        r,
+        batch,
+        batch,
+        max_num_neighbors,
+        num_workers,
+        batch_size,
+        not loop,
+        use_triton=use_triton,
+    )
     if flow == 'source_to_target':
-        row, col = edge_index[1], edge_index[0]
-    else:
-        row, col = edge_index[0], edge_index[1]
-
-    return torch.stack([row, col], dim=0)
+        return edge_index.flip(0).contiguous()
+    return edge_index
diff --git a/torch_cluster/rw.py b/torch_cluster/rw.py
index ecd07e92..83b3ea52 100644
--- a/torch_cluster/rw.py
+++ b/torch_cluster/rw.py
@@ -4,6 +4,19 @@
 from torch import Tensor
 
 
+@torch.library.register_fake("torch_cluster::random_walk")
+def _(rowptr, col, start, walk_length, p=1.0, q=1.0):
+    torch._check(rowptr.device == col.device)
+    torch._check(start.device == rowptr.device)
+    torch._check(rowptr.ndim == 1)
+    torch._check(col.ndim == 1)
+    torch._check(start.ndim == 1)
+    num_walks = start.numel()
+    node_seq = start.new_empty((num_walks, walk_length + 1), dtype=torch.long)
+    edge_seq = start.new_empty((num_walks, walk_length), dtype=torch.long)
+    return node_seq, edge_seq
+
+
 def random_walk(
     row: Tensor,
     col: Tensor,
diff --git a/torch_cluster/sampler.py b/torch_cluster/sampler.py
index 1b68de0a..8cf98d1c 100644
--- a/torch_cluster/sampler.py
+++ b/torch_cluster/sampler.py
@@ -1,6 +1,16 @@
 import torch
 
 
+@torch.library.register_fake("torch_cluster::neighbor_sampler")
+def _(start, rowptr, count, factor):
+    torch._check(start.device == rowptr.device)
+    torch._check(start.ndim == 1)
+    torch._check(rowptr.ndim == 1)
+    ctx = torch.library.get_ctx()
+    out_len = ctx.new_dynamic_size()
+    return start.new_empty((out_len,), dtype=torch.long)
+
+
 def neighbor_sampler(start: torch.Tensor, rowptr: torch.Tensor, size: float):
     assert not start.is_cuda
 
diff --git a/torch_cluster/testing.py b/torch_cluster/testing.py
index 68949fa5..0d2e7db0 100644
--- a/torch_cluster/testing.py
+++ b/torch_cluster/testing.py
@@ -1,4 +1,6 @@
 from typing import Any
+import platform
+import shutil
 
 import torch
 
@@ -19,3 +21,17 @@
 
 def tensor(x: Any, dtype: torch.dtype, device: torch.device):
     return None if x is None else torch.tensor(x, dtype=dtype, device=device)
+
+
+def triton_wrap(dt_device_seq):
+    return [
+        (dt, device, use_triton)
+        for dt, device in dt_device_seq
+        for use_triton in ([False, True] if device.type == 'cuda' else [False])
+    ]
+
+
+def has_compiler() -> bool:
+    if platform.system().lower() != 'windows':
+        return True
+    return shutil.which('cl') is not None
diff --git a/torch_cluster/triton/__init__.py b/torch_cluster/triton/__init__.py
new file mode 100644
index 00000000..15dc9335
--- /dev/null
+++ b/torch_cluster/triton/__init__.py
@@ -0,0 +1,9 @@
+from .knn import knn  # noqa
+from .nearest import nearest  # noqa
+from .radius import radius  # noqa
+
+__all__ = [
+    'knn',
+    'nearest',
+    'radius',
+]
diff --git a/torch_cluster/triton/_kernels.py b/torch_cluster/triton/_kernels.py
new file mode 100644
index 00000000..af463c16
--- /dev/null
+++ b/torch_cluster/triton/_kernels.py
@@ -0,0 +1,723 @@
+from __future__ import annotations
+
+
+import triton
+import triton.language as tl
+from triton.language.extra import libdevice
+
+from triton.language import topk as tl_topk
+import inspect
+sig = inspect.signature(tl_topk)
+if 'descending' not in sig.parameters:
+    from triton.language.standard import sort_impl as tl_topk
+
+
+def _get_block_d(D, BLOCK_D: int = 16):
+    return min(triton.next_power_of_2(D), BLOCK_D)
+
+
+@triton.jit
+def _neumaier_add(acc, acc_c, term):
+    t = acc + term
+    cond = tl.abs(acc) >= tl.abs(term)
+    acc_c += tl.where(
+        cond,
+        (acc - t) + term,
+        (term - t) + acc,
+    )
+    return t, acc_c
+
+
+@triton.jit
+def _load_tile(ptr, even: tl.constexpr):
+    if even:
+        return tl.load(ptr)
+    return tl.load(
+        ptr,
+        boundary_check=(0, 1),
+        padding_option="zero",
+    )
+
+
+@triton.jit
+def _compute_mask_x(
+    x_block_start,
+    offs_n,
+    x_end,
+    N,
+    mask_y,
+    BLOCK_N: tl.constexpr,
+):
+    x_idx = x_block_start + offs_n
+    full_block = (
+        (x_block_start + BLOCK_N <= x_end)
+        & (x_block_start + BLOCK_N <= N)
+        & mask_y
+    )
+    mask_x = tl.where(
+        full_block,
+        mask_y,
+        (x_idx < x_end) & (x_idx < N) & mask_y,
+    )
+    return x_idx, full_block, mask_x
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({'BLOCK_M': 32, 'BLOCK_N': 32},
+                      num_warps=2,
+                      num_stages=2),
+        triton.Config({'BLOCK_M': 32, 'BLOCK_N': 64},
+                      num_warps=2,
+                      num_stages=2),
+        triton.Config({'BLOCK_M': 64, 'BLOCK_N': 32},
+                      num_warps=2,
+                      num_stages=2),
+        triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64},
+                      num_warps=4,
+                      num_stages=3),
+        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64},
+                      num_warps=4,
+                      num_stages=3),
+        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 32},
+                      num_warps=4,
+                      num_stages=3),
+        triton.Config({'BLOCK_M': 32, 'BLOCK_N': 128},
+                      num_warps=4,
+                      num_stages=3),
+        triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128},
+                      num_warps=4,
+                      num_stages=3),
+        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128},
+                      num_warps=8,
+                      num_stages=4),
+        triton.Config({'BLOCK_M': 256, 'BLOCK_N': 64},
+                      num_warps=8,
+                      num_stages=2),
+        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256},
+                      num_warps=8,
+                      num_stages=2),
+    ],
+    key=['M', 'N', 'D'],
+)
+@triton.heuristics({
+    'BLOCK_D': lambda args: _get_block_d(args['D']),
+    'EVEN_D': lambda args: args['D'] % _get_block_d(args['D']) == 0,
+    'EVEN_N': lambda args: args['N'] % args['BLOCK_N'] == 0,
+    'EVEN_K': lambda args: args['D'] % args['BLOCK_K'] == 0,
+})
+@triton.jit
+def _nearest_kernel(
+    x_ptr,
+    y_ptr,
+    out_ptr,
+    M,
+    N,
+    D,
+    stride_xm,
+    stride_xd,
+    stride_ym,
+    stride_yd,
+    ptr_y_ptr,
+    batch_x_ptr,
+    USE_BATCH: tl.constexpr,
+    COSINE: tl.constexpr,
+    EPS: tl.constexpr,
+    INPUT_PRECISION: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    EVEN_N: tl.constexpr,
+    EVEN_D: tl.constexpr,
+):
+    r"""Compute nearest neighbor indices for blocks of x.
+
+    Args:
+        x_ptr: Pointer to x matrix.
+        y_ptr: Pointer to y matrix.
+        out_ptr: Pointer to output indices.
+        M (int): Number of y rows.
+        N (int): Number of x rows.
+        D (int): Feature dimension.
+        stride_xm/stride_xd: Strides for x.
+        stride_ym/stride_yd: Strides for y.
+        ptr_y_ptr: Pointer to y batch prefix sums.
+        batch_x_ptr: Pointer to x batch ids.
+        USE_BATCH (constexpr bool): Whether to apply batch masking.
+        COSINE (constexpr bool): Use cosine distance if True.
+        EPS (constexpr float): Numerical epsilon.
+        BLOCK_M/BLOCK_N/BLOCK_K (constexpr int): Tile sizes.
+        EVEN_N (constexpr bool): Alignment hints.
+    """
+    # X block id.
+    pid = tl.program_id(0)
+    # X indices.
+    offs_n = pid * BLOCK_N + tl.arange(0, BLOCK_N)
+    # Valid x rows.
+    mask_x = offs_n < N
+    if EVEN_N:
+        # Hint contiguous access.
+        tl.multiple_of(offs_n, 8)
+        # Hint vectorization.
+        tl.max_contiguous(offs_n, 8)
+
+    # Best dist.
+    best_dist = tl.full((BLOCK_N,), float('inf'), tl.float32)
+    # Best y index.
+    best_idx = tl.zeros((BLOCK_N,), dtype=tl.int32)
+
+    if USE_BATCH:
+        # Batch id per x.
+        batch_id = tl.load(
+            batch_x_ptr + offs_n,
+            mask=mask_x,
+            other=0,
+        )
+        left = tl.load(
+            ptr_y_ptr + batch_id,
+            mask=mask_x,
+            other=0,
+        )
+        # y range start.
+        right = tl.load(
+            ptr_y_ptr + batch_id + 1,
+            mask=mask_x,
+            other=0,
+        )
+        # y range end.
+    else:
+        # Full y range.
+        left = 0
+        right = M
+
+    # ||x||^2.
+    x_sq = tl.zeros((BLOCK_N,), dtype=tl.float32)
+    x_sq_c = tl.zeros((BLOCK_N,), dtype=tl.float32)
+    x_block_ptr_sq = tl.make_block_ptr(
+        base=x_ptr,
+        shape=(D, N),
+        strides=(stride_xd, stride_xm),
+        offsets=(0, pid * BLOCK_N),
+        block_shape=(BLOCK_D, BLOCK_N),
+        order=(0, 1),
+    )
+    for _ in range(0, D, BLOCK_D):
+        x = _load_tile(x_block_ptr_sq, EVEN_N and EVEN_D)
+        x_term = tl.sum(x * x, axis=0)
+        x_sq, x_sq_c = _neumaier_add(x_sq, x_sq_c, x_term)
+        x_block_ptr_sq = tl.advance(x_block_ptr_sq, (BLOCK_D, 0))
+    x_sq += x_sq_c
+    if COSINE:
+        # 1/||x||.
+        inv_x = tl.rsqrt(x_sq + EPS)
+
+    x_block_ptr_base = tl.make_block_ptr(
+        base=x_ptr,
+        shape=(D, N),
+        strides=(stride_xd, stride_xm),
+        offsets=(0, pid * BLOCK_N),
+        block_shape=(BLOCK_D, BLOCK_N),
+        order=(0, 1),
+    )
+
+    for y_start in range(0, M, BLOCK_M):
+        # y indices.
+        offs_m = y_start + tl.arange(0, BLOCK_M)
+        # Valid y rows.
+        mask_y = offs_m < M
+        # Full tile.
+        full_y = y_start + BLOCK_M <= M
+
+        # Dot acc.
+        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        acc_c = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        # ||y||^2.
+        y_sq = tl.zeros((BLOCK_M,), dtype=tl.float32)
+        y_sq_c = tl.zeros((BLOCK_M,), dtype=tl.float32)
+
+        x_block_ptr = x_block_ptr_base
+        y_block_ptr = tl.make_block_ptr(
+            base=y_ptr,
+            shape=(M, D),
+            strides=(stride_ym, stride_yd),
+            offsets=(y_start, 0),
+            block_shape=(BLOCK_M, BLOCK_D),
+            order=(1, 0),
+        )
+
+        for _ in range(0, D, BLOCK_D):
+            full_tile = full_y and EVEN_D
+            y = _load_tile(y_block_ptr, full_tile)
+
+            full_x = full_tile & EVEN_N
+            x = _load_tile(x_block_ptr, full_x)
+            dot_term = tl.dot(y, x, input_precision=INPUT_PRECISION)
+            acc, acc_c = _neumaier_add(acc, acc_c, dot_term)
+            y_term = tl.sum(y * y, axis=1)
+            y_sq, y_sq_c = _neumaier_add(y_sq, y_sq_c, y_term)
+            y_block_ptr = tl.advance(y_block_ptr, (0, BLOCK_D))
+            x_block_ptr = tl.advance(x_block_ptr, (BLOCK_D, 0))
+
+        acc += acc_c
+        y_sq += y_sq_c
+        if COSINE:
+            # 1/||y||.
+            inv_y = tl.rsqrt(y_sq + EPS)
+            dist = 1.0 - acc * (inv_y[:, None] * inv_x[None, :])
+        else:
+            dist = (
+                (y_sq)[:, None]
+                + (x_sq)[None, :]
+                - 2.0 * acc
+            )
+
+        if full_y:
+            valid = tl.broadcast_to(mask_x[None, :], (BLOCK_M, BLOCK_N))
+        else:
+            valid = mask_y[:, None] & mask_x[None, :]
+        if USE_BATCH:
+            valid &= (
+                (offs_m[:, None] >= left[None, :])
+                & (offs_m[:, None] < right[None, :])
+            )
+        dist = tl.where(valid, dist, float('inf'))
+
+        block_min, block_arg = tl.min(
+            dist,
+            axis=0,
+            return_indices=True,
+            return_indices_tie_break_left=True,
+        )
+        block_idx = y_start + block_arg
+        better = block_min < best_dist
+        best_dist = tl.where(better, block_min, best_dist)
+        best_idx = tl.where(better, block_idx, best_idx)
+
+    # Write output.
+    tl.store(
+        out_ptr + offs_n,
+        best_idx.to(tl.int64),
+        mask=mask_x,
+    )
+
+
+@triton.autotune(
+    configs=[
+        triton.Config(
+            {'BLOCK_N': 32},
+            num_warps=2,
+            num_stages=2,
+        ),
+        triton.Config(
+            {'BLOCK_N': 64},
+            num_warps=2,
+            num_stages=2,
+        ),
+        triton.Config(
+            {'BLOCK_N': 64},
+            num_warps=4,
+            num_stages=3,
+        ),
+        triton.Config(
+            {'BLOCK_N': 128},
+            num_warps=8,
+            num_stages=4,
+        ),
+        triton.Config(
+            {'BLOCK_N': 256},
+            num_warps=8,
+            num_stages=4,
+        ),
+    ],
+    key=['D', 'MAX_CAND'],
+)
+@triton.heuristics({
+    'BLOCK_D': lambda args: _get_block_d(args['D']),
+    'EVEN_D': lambda args: args['D'] % _get_block_d(args['D']) == 0,
+    'N_BLOCKS_PER_K': (
+        lambda args: triton.cdiv(args['K'], args['BLOCK_N'])
+    ),
+})
+@triton.jit
+def _knn_segmented_kernel(
+    x_ptr,
+    y_ptr,
+    ptr_x_ptr,
+    batch_y_ptr,
+    grid_ptr,
+    M,
+    N,
+    D,
+    MAX_CAND,
+    stride_xm,
+    stride_xd,
+    stride_ym,
+    stride_yd,
+    K: tl.constexpr,
+    USE_BATCH: tl.constexpr,
+    INPUT_PRECISION: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    EVEN_D: tl.constexpr,
+    COSINE: tl.constexpr,
+    EPS: tl.constexpr,
+    IGNORE_SAME_INDEX: tl.constexpr,
+    N_BLOCKS_PER_K: tl.constexpr,
+):
+    # Program id over y rows.
+    pid = tl.program_id(0)
+    # Current y index.
+    n_y = pid
+    # Mask for valid y.
+    mask_y = n_y < M
+
+    # Offsets for padded top-k buffer.
+    k_offsets = tl.arange(0, N_BLOCKS_PER_K * BLOCK_N)
+    INF_KEY = 0x7f800000ffffffff
+
+    best_dist_key = tl.full((N_BLOCKS_PER_K * BLOCK_N,), INF_KEY, tl.int64)
+
+    if N_BLOCKS_PER_K > 1:
+        acc_key = tl.full((N_BLOCKS_PER_K, BLOCK_N), INF_KEY, tl.int64)
+        k_rows = tl.arange(0, N_BLOCKS_PER_K)[:, None].broadcast_to(
+            (N_BLOCKS_PER_K, BLOCK_N)
+        )
+
+    if USE_BATCH:
+        # Segment id.
+        example_idx = tl.load(batch_y_ptr + n_y, mask=mask_y, other=0)
+    else:
+        example_idx = 0
+
+    # Segment start.
+    x_start = tl.load(ptr_x_ptr + example_idx, mask=mask_y, other=0)
+    # Segment end.
+    x_end = tl.load(ptr_x_ptr + example_idx + 1, mask=mask_y, other=0)
+    if COSINE:
+        # ||y||^2 accumulator.
+        y_sq = tl.zeros((1,), tl.float32)
+        y_sq_c = tl.zeros((1,), tl.float32)
+        y_norm_ptr = tl.make_block_ptr(
+            base=y_ptr,
+            shape=(D, M),
+            strides=(stride_yd, stride_ym),
+            offsets=(0, n_y),
+            block_shape=(BLOCK_D, 1),
+            order=(0, 1),
+        )
+        for yi in range(0, D, BLOCK_D):
+            y = _load_tile(y_norm_ptr, EVEN_D)
+            y_term = tl.sum(y * y, axis=0)
+            y_sq, y_sq_c = _neumaier_add(y_sq, y_sq_c, y_term)
+            y_norm_ptr = tl.advance(y_norm_ptr, (BLOCK_D, 0))
+        # 1/||y|| for cosine.
+        y_rnorm = tl.rsqrt(y_sq + EPS)
+
+    for xb in range(0, MAX_CAND, BLOCK_N):
+        # Start of x block.
+        x_block_start = x_start + xb
+        # Offsets within block.
+        offs_n = tl.arange(0, BLOCK_N)
+        x_idx, full_block, mask_x = _compute_mask_x(
+            x_block_start,
+            offs_n,
+            x_end,
+            N,
+            mask_y,
+            BLOCK_N,
+        )
+        if MAX_CAND > BLOCK_N:
+            # Hint alignment for block start.
+            tl.multiple_of(x_block_start, 8)
+        # Hint vectorization.
+        tl.multiple_of(offs_n, 8)
+        tl.max_contiguous(offs_n, 8)
+
+        if COSINE:
+            # Dot accumulator.
+            acc_dot = tl.zeros((BLOCK_N,), tl.float32)
+            acc_dot_c = tl.zeros((BLOCK_N,), tl.float32)
+            # ||x||^2 accumulator.
+            acc_x_sq = tl.zeros((BLOCK_N,), tl.float32)
+            acc_x_sq_c = tl.zeros((BLOCK_N,), tl.float32)
+        else:
+            acc_dist = tl.zeros((BLOCK_N,), tl.float32)
+            acc_dist_c = tl.zeros((BLOCK_N,), tl.float32)
+        # Block ptr needs int32 offsets.
+        x_block_start_i32 = x_block_start.to(tl.int32)
+        n_y_i32 = n_y.to(tl.int32)
+        x_block_ptr = tl.make_block_ptr(
+            base=x_ptr,
+            shape=(N, D),
+            strides=(stride_xm, stride_xd),
+            offsets=(x_block_start_i32, 0),
+            block_shape=(BLOCK_N, BLOCK_D),
+            order=(1, 0),
+        )
+        if COSINE:
+            y_block_ptr = tl.make_block_ptr(
+                base=y_ptr,
+                shape=(D, M),
+                strides=(stride_yd, stride_ym),
+                offsets=(0, n_y_i32),
+                block_shape=(BLOCK_D, 1),
+                order=(0, 1),
+            )
+        else:
+            y_block_ptr = tl.make_block_ptr(
+                base=y_ptr,
+                shape=(M, D),
+                strides=(stride_ym, stride_yd),
+                offsets=(n_y_i32, 0),
+                block_shape=(1, BLOCK_D),
+                order=(0, 1),
+            )
+
+        for nd in range(0, D, BLOCK_D):
+            x = _load_tile(x_block_ptr, EVEN_D)
+            y = _load_tile(y_block_ptr, EVEN_D)
+            if COSINE:
+                # MxV dot.
+                prod = tl.dot(x, y, input_precision=INPUT_PRECISION)
+                dot_term = tl.sum(prod, axis=1)
+                acc_dot, acc_dot_c = _neumaier_add(
+                    acc_dot,
+                    acc_dot_c,
+                    dot_term,
+                )
+                x_term = tl.sum(x * x, axis=1)
+                acc_x_sq, acc_x_sq_c = _neumaier_add(
+                    acc_x_sq,
+                    acc_x_sq_c,
+                    x_term,
+                )
+                y_block_ptr = tl.advance(y_block_ptr, (BLOCK_D, 0))
+            else:
+                diff = x - y
+                term = tl.sum(diff * diff, axis=1)
+                acc_dist, acc_dist_c = _neumaier_add(
+                    acc_dist,
+                    acc_dist_c,
+                    term,
+                )
+                y_block_ptr = tl.advance(y_block_ptr, (0, BLOCK_D))
+            x_block_ptr = tl.advance(x_block_ptr, (0, BLOCK_D))
+        if COSINE:
+            # 1/||x||.
+            x_rnorm = tl.rsqrt(acc_x_sq + acc_x_sq_c + EPS)
+            # Cosine distance.
+            dist = 1.0 - acc_dot * (x_rnorm * y_rnorm)
+        else:
+            # L2^2 distance.
+            dist = acc_dist + acc_dist_c
+
+        if not full_block:
+            # Mask invalid x.
+            dist = tl.where(mask_x, dist, float("inf"))
+
+        if IGNORE_SAME_INDEX:
+            same_idx = x_idx == n_y
+            if not full_block:
+                same_idx = same_idx & mask_x
+            dist = tl.where(same_idx, float("inf"), dist)
+
+        dist_bits = dist.to(tl.int32, bitcast=True).to(tl.int64)
+        # Pack idx.
+        idx_bits = x_idx & 0xFFFFFFFF
+        # Lexicographic key.
+        key = (dist_bits << 32) | idx_bits
+        key = tl.where(libdevice.isinf(dist), INF_KEY, key)
+
+        absorb = False
+        if N_BLOCKS_PER_K > 1:
+            row = (xb // BLOCK_N) % N_BLOCKS_PER_K
+            acc_key = tl.where(k_rows == row, key[None, :], acc_key)
+            if row == (N_BLOCKS_PER_K - 1) or xb + BLOCK_N > MAX_CAND:
+                absorb = True
+                final_key = acc_key.ravel()
+                acc_key = tl.full(
+                    (N_BLOCKS_PER_K, BLOCK_N),
+                    INF_KEY,
+                    tl.int64,
+                )
+        else:
+            absorb = True
+            final_key = key
+
+        if absorb:
+            # Merge buffers.
+            combo_key = tl.cat(best_dist_key, final_key, can_reorder=True)
+            # Get the best.
+            best_dist_key = tl_topk(
+                combo_key,
+                N_BLOCKS_PER_K * BLOCK_N,
+                descending=False,
+            )
+
+    # Mask for real k.
+    k_mask = k_offsets < K
+    best_dist_key = tl.where(k_mask, best_dist_key, INF_KEY)
+    best_idx = (best_dist_key & 0xFFFFFFFF).to(tl.int32)
+
+    # Output offsets.
+    out_offsets = n_y * K + k_offsets
+    # Valid output mask.
+    out_mask = mask_y & k_mask
+
+    tl.store(grid_ptr + out_offsets, n_y, mask=out_mask)
+    # Write col idx.
+    tl.store(
+        grid_ptr + M * K + out_offsets,
+        best_idx.to(tl.int64),
+        mask=out_mask,
+    )
+
+
+@triton.autotune(
+    configs=[
+        triton.Config(
+            {'BLOCK_N': 32},
+            num_warps=2,
+            num_stages=2,
+        ),
+        triton.Config(
+            {'BLOCK_N': 64},
+            num_warps=4,
+            num_stages=2,
+        ),
+        triton.Config(
+            {'BLOCK_N': 64},
+            num_warps=4,
+            num_stages=3,
+        ),
+        triton.Config(
+            {'BLOCK_N': 128},
+            num_warps=4,
+            num_stages=3,
+        ),
+        triton.Config(
+            {'BLOCK_N': 256},
+            num_warps=4,
+            num_stages=3,
+        ),
+    ],
+    key=['N', 'D', 'MAX_CAND'],
+)
+@triton.heuristics({
+    'BLOCK_D': lambda args: _get_block_d(args['D']),
+    'EVEN_D': lambda args: args['D'] % _get_block_d(args['D']) == 0,
+    'MAX_X_BLOCKS': lambda args: triton.cdiv(
+        args['MAX_CAND'],
+        args['BLOCK_N'],
+    ),
+})
+@triton.jit
+def _radius_segmented_kernel(
+    x_ptr,
+    y_ptr,
+    ptr_x_ptr,
+    example_idx_ptr,
+    grid_ptr,
+    M,
+    N,
+    D,
+    MAX_CAND,
+    stride_xm,
+    stride_xd,
+    stride_ym,
+    stride_yd,
+    R2,
+    MAX_NEIGHBORS,
+    USE_BATCH: tl.constexpr,
+    IGNORE_SAME_INDEX: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    MAX_X_BLOCKS: tl.constexpr,
+    EVEN_D: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    n_y = pid
+    mask_y = n_y < M
+
+    if USE_BATCH:
+        example_idx = tl.load(example_idx_ptr + n_y, mask=mask_y, other=0)
+    else:
+        example_idx = 0
+
+    x_start = tl.load(ptr_x_ptr + example_idx, mask=mask_y, other=0)
+    x_end = tl.load(ptr_x_ptr + example_idx + 1, mask=mask_y, other=0)
+
+    if stride_xm % 8 == 0:
+        tl.multiple_of(stride_xm, 8)
+    if stride_xd % 8 == 0:
+        tl.multiple_of(stride_xd, 8)
+    if stride_ym % 8 == 0:
+        tl.multiple_of(stride_ym, 8)
+    if stride_yd % 8 == 0:
+        tl.multiple_of(stride_yd, 8)
+
+    count = tl.zeros((), dtype=tl.int32)
+    max_neighbors = MAX_NEIGHBORS.to(tl.int32)
+    n_y_i32 = n_y.to(tl.int32)
+    xb = 0
+    offs_n = tl.arange(0, BLOCK_N)
+
+    while (xb < MAX_X_BLOCKS) and (count < max_neighbors):
+        x_block_start = x_start + xb * BLOCK_N
+        x_idx, full_block, mask_x = _compute_mask_x(
+            x_block_start,
+            offs_n,
+            x_end,
+            N,
+            mask_y,
+            BLOCK_N,
+        )
+
+        if MAX_X_BLOCKS > 1:
+            tl.multiple_of(x_block_start, 8)
+
+        acc_dist = tl.zeros((BLOCK_N,), tl.float32)
+        acc_dist_c = tl.zeros((BLOCK_N,), tl.float32)
+        x_block_start_i32 = x_block_start.to(tl.int32)
+        x_block_ptr = tl.make_block_ptr(
+            base=x_ptr,
+            shape=(N, D),
+            strides=(stride_xm, stride_xd),
+            offsets=(x_block_start_i32, 0),
+            block_shape=(BLOCK_N, BLOCK_D),
+            order=(1, 0),
+        )
+        y_row_ptr = tl.make_block_ptr(
+            base=y_ptr,
+            shape=(M, D),
+            strides=(stride_ym, stride_yd),
+            offsets=(n_y_i32, 0),
+            block_shape=(1, BLOCK_D),
+            order=(0, 1),
+        )
+
+        for nd in range(0, D, BLOCK_D):
+            x = _load_tile(x_block_ptr, EVEN_D)
+            y = _load_tile(y_row_ptr, EVEN_D)
+
+            diff = x - y
+            term = tl.sum(diff * diff, axis=1)
+            acc_dist, acc_dist_c = _neumaier_add(acc_dist, acc_dist_c, term)
+            x_block_ptr = tl.advance(x_block_ptr, (0, BLOCK_D))
+            y_row_ptr = tl.advance(y_row_ptr, (0, BLOCK_D))
+
+        dist = acc_dist + acc_dist_c
+        active = count < max_neighbors
+        mask = mask_x & (dist < R2) & active
+        if IGNORE_SAME_INDEX:
+            mask &= x_idx != n_y.to(tl.int64)
+
+        prefix = tl.cumsum(mask, axis=0).to(tl.int32)
+        pos = count + prefix - 1
+        write = mask & (pos < max_neighbors)
+        pos_safe = tl.where(write, pos, 0)
+        out_offset = n_y * MAX_NEIGHBORS + pos_safe
+        tl.store(grid_ptr + out_offset, n_y, mask=write)
+        tl.store(grid_ptr + M * MAX_NEIGHBORS + out_offset, x_idx, mask=write)
+        count += tl.sum(write, axis=0).to(tl.int32)
+        xb += 1
diff --git a/torch_cluster/triton/knn.py b/torch_cluster/triton/knn.py
new file mode 100644
index 00000000..448c7059
--- /dev/null
+++ b/torch_cluster/triton/knn.py
@@ -0,0 +1,43 @@
+from __future__ import annotations
+from typing import Optional
+
+from torch import Tensor
+
+from .segmented import segmented_topk_search
+
+
+def knn(
+    x: Tensor,
+    y: Tensor,
+    k: int,
+    batch_x: Optional[Tensor] = None,
+    batch_y: Optional[Tensor] = None,
+    cosine: bool = False,
+    batch_size: Optional[int] = None,
+) -> Tensor:
+    r"""Compute k-NN indices using Triton kernels.
+
+    Args:
+        x (Tensor): Input features of shape [N, D].
+        y (Tensor): Query features of shape [M, D].
+        k (int): Number of neighbors.
+        batch_x (Tensor, optional): Batch vector for x.
+        batch_y (Tensor, optional): Batch vector for y.
+        cosine (bool, optional): Use cosine distance if True.
+        num_workers (int, optional): Unused for Triton path.
+        batch_size (int, optional): Number of batches.
+        ignore_same_index (bool, optional): Exclude exact self matches.
+
+    Returns:
+        Tensor: Edge index with shape [2, M * k].
+    """
+    grid = segmented_topk_search(
+        x,
+        y,
+        k,
+        batch_x,
+        batch_y,
+        cosine,
+        batch_size,
+    )
+    return grid[:, grid[1] >= 0]
diff --git a/torch_cluster/triton/nearest.py b/torch_cluster/triton/nearest.py
new file mode 100644
index 00000000..83a684d3
--- /dev/null
+++ b/torch_cluster/triton/nearest.py
@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import triton
+from torch import Tensor
+
+from ._kernels import _nearest_kernel
+
+
+def nearest(
+    x: Tensor,
+    y: Tensor,
+    batch_x: Optional[Tensor] = None,
+    batch_y: Optional[Tensor] = None,
+) -> Tensor:
+    r"""Find nearest neighbors using Triton pairwise distances.
+
+    Args:
+        x (Tensor): Input features of shape [N, D].
+        y (Tensor): Query features of shape [M, D].
+        batch_x (Tensor, optional): Batch vector for x.
+        batch_y (Tensor, optional): Batch vector for y.
+
+    Returns:
+        Tensor: Index of the nearest y for each x.
+    """
+    if batch_x is None and batch_y is None:
+        out = torch.empty((x.size(0),), dtype=torch.long, device=x.device)
+
+        def grid(meta):
+            return (triton.cdiv(x.size(0), meta['BLOCK_N']),)
+
+        eps = torch.finfo(torch.float32).eps
+        _nearest_kernel[grid](
+            x,
+            y,
+            out,
+            y.size(0),
+            x.size(0),
+            x.size(1),
+            x.stride(0),
+            x.stride(1),
+            y.stride(0),
+            y.stride(1),
+            y,
+            x,
+            USE_BATCH=False,
+            COSINE=False,
+            EPS=eps,
+            INPUT_PRECISION="ieee",
+        )
+        return out
+
+    if batch_x is None:
+        batch_x = x.new_zeros(x.size(0), dtype=torch.long)
+    if batch_y is None:
+        batch_y = y.new_zeros(y.size(0), dtype=torch.long)
+
+    batch_x_unique = batch_x.unique_consecutive()
+    batch_y_unique = batch_y.unique_consecutive()
+    if not torch.equal(batch_x_unique, batch_y_unique):
+        raise ValueError("Some batch indices occur in 'batch_x' "
+                         "that do not occur in 'batch_y'")
+
+    batch_size = int(max(batch_x.max(), batch_y.max())) + 1
+    use_batch = batch_size > 1
+    if use_batch:
+        arange = torch.arange(batch_size + 1, device=x.device)
+        ptr_y = torch.bucketize(arange, batch_y)
+    else:
+        ptr_y = y
+        batch_x = x
+    out = torch.empty((x.size(0),), dtype=torch.long, device=x.device)
+
+    def grid(meta):
+        return (triton.cdiv(x.size(0), meta['BLOCK_N']),)
+
+    eps = torch.finfo(torch.float32).eps
+    _nearest_kernel[grid](
+        x,
+        y,
+        out,
+        y.size(0),
+        x.size(0),
+        x.size(1),
+        x.stride(0),
+        x.stride(1),
+        y.stride(0),
+        y.stride(1),
+        ptr_y,
+        batch_x,
+        USE_BATCH=use_batch,
+        COSINE=False,
+        EPS=eps,
+        INPUT_PRECISION="ieee",
+    )
+    return out
diff --git a/torch_cluster/triton/radius.py b/torch_cluster/triton/radius.py
new file mode 100644
index 00000000..3d62750b
--- /dev/null
+++ b/torch_cluster/triton/radius.py
@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+from ._kernels import _radius_segmented_kernel
+
+
+def radius(
+    x: Tensor,
+    y: Tensor,
+    r: float,
+    batch_x: Optional[Tensor] = None,
+    batch_y: Optional[Tensor] = None,
+    max_num_neighbors: int = 32,
+    batch_size: Optional[int] = None,
+    ignore_same_index: bool = False,
+) -> Tensor:
+    r"""Find all neighbors within a radius using Triton kernels.
+
+    Args:
+        x (Tensor): Input features of shape [N, D].
+        y (Tensor): Query features of shape [M, D].
+        r (float): Search radius.
+        batch_x (Tensor, optional): Batch vector for x.
+        batch_y (Tensor, optional): Batch vector for y.
+        max_num_neighbors (int, optional): Maximum neighbors per y.
+        batch_size (int, optional): Number of batches.
+        ignore_same_index (bool, optional): Exclude exact self matches.
+
+    Returns:
+        Tensor: Edge index with shape [2, num_edges].
+    """
+    use_batch = (batch_size or 1) > 1
+    if use_batch:
+        assert batch_x is not None
+        assert batch_y is not None
+        arange = torch.arange((batch_size or 1) + 1, device=x.device)
+        ptr_x = torch.bucketize(arange, batch_x)
+        ptr_y = torch.bucketize(arange, batch_y)
+    else:
+        ptr_x = torch.tensor(
+            [0, x.size(0)],
+            device=x.device,
+            dtype=torch.int64,
+        )
+        ptr_y = torch.tensor(
+            [0, y.size(0)],
+            device=y.device,
+            dtype=torch.int64,
+        )
+
+    if ptr_x.numel() != ptr_y.numel():
+        raise ValueError(
+            "ptr_x and ptr_y must have the same number of segments."
+        )
+
+    M, N = y.size(0), x.size(0)
+    D = x.size(1)
+    if use_batch:
+        example_idx = batch_y.to(torch.int64).contiguous()
+    else:
+        example_idx = ptr_x  # Dummy; kernel ignores when USE_BATCH=False.
+
+    seg_sizes = torch.diff(ptr_x).to(torch.int64)
+    if seg_sizes.numel() > 0:
+        max_candidates = int(seg_sizes.max().item())
+    else:
+        max_candidates = 0
+
+    if max_candidates == 0 or max_num_neighbors <= 0:
+        return torch.empty(2, 0, device=x.device, dtype=torch.int64)
+
+    grid_out = torch.full(
+        (2, M * max_num_neighbors),
+        -1,
+        device=x.device,
+        dtype=torch.int64,
+    )
+    grid = (M,)
+    _radius_segmented_kernel[grid](
+        x,
+        y,
+        ptr_x,
+        example_idx,
+        grid_out,
+        M,
+        N,
+        D,
+        max_candidates,
+        x.stride(0),
+        x.stride(1),
+        y.stride(0),
+        y.stride(1),
+        float(r) * float(r),
+        max_num_neighbors,
+        USE_BATCH=use_batch,
+        IGNORE_SAME_INDEX=ignore_same_index,
+    )
+
+    return grid_out[:, grid_out[0] != -1]
diff --git a/torch_cluster/triton/segmented.py b/torch_cluster/triton/segmented.py
new file mode 100644
index 00000000..e5c07ccc
--- /dev/null
+++ b/torch_cluster/triton/segmented.py
@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+from ._kernels import _knn_segmented_kernel
+
+
+def segmented_topk_search(
+    x: Tensor,
+    y: Tensor,
+    k: int,
+    batch_x: Optional[Tensor] = None,
+    batch_y: Optional[Tensor] = None,
+    cosine: bool = False,
+    batch_size: Optional[int] = None,
+    ignore_same_index: bool = False,
+) -> Tensor:
+    r"""Compute top-k neighbor indices."""
+    use_batch = (batch_size or 1) > 1
+    if use_batch:
+        assert batch_x is not None
+        assert batch_y is not None
+        arange = torch.arange((batch_size or 1) + 1, device=x.device)
+        ptr_x = torch.bucketize(arange, batch_x)
+        ptr_y = torch.bucketize(arange, batch_y)
+    else:
+        ptr_x = torch.tensor(
+            [0, x.size(0)],
+            device=x.device,
+            dtype=torch.int64,
+        )
+        ptr_y = torch.tensor(
+            [0, y.size(0)],
+            device=y.device,
+            dtype=torch.int64,
+        )
+
+    if ptr_x.numel() != ptr_y.numel():
+        raise ValueError(
+            "ptr_x and ptr_y must have the same number of segments."
+        )
+
+    M, N = y.size(0), x.size(0)
+    D = x.size(1)
+    if use_batch:
+        batch_y = batch_y.to(torch.int64).contiguous()
+
+    seg_sizes = torch.diff(ptr_x).to(torch.int64)
+    if seg_sizes.numel() > 0:
+        max_candidates = int(seg_sizes.max().item())
+    else:
+        max_candidates = 0
+    eps = 0.0 if cosine else torch.finfo(torch.float32).eps
+
+    if max_candidates == 0:
+        return torch.full((2, M * k), -1, device=y.device, dtype=torch.int64)
+
+    grid_out = torch.empty(2, M * k, device=y.device, dtype=torch.int64)
+    grid = (M,)
+    _knn_segmented_kernel[grid](
+        x,
+        y,
+        ptr_x,
+        batch_y if use_batch else ptr_x,  # dummy ptr when unused
+        grid_out,
+        M,
+        N,
+        D,
+        max_candidates,
+        x.stride(0),
+        x.stride(1),
+        y.stride(0),
+        y.stride(1),
+        K=k,
+        USE_BATCH=use_batch,
+        INPUT_PRECISION="ieee",
+        COSINE=cosine,
+        EPS=eps,
+        IGNORE_SAME_INDEX=ignore_same_index,
+    )
+    return grid_out