chore: import upstream snapshot with attribution

tests/python/disco/test_callback.py

@@ -0,0 +1,175 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Test sharded loader"""
+# pylint: disable=missing-docstring
+
+import pathlib
+import tempfile
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import relax as R
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(tvm.runtime.disco is None, reason="disco runtime is not available")
+@pytest.mark.skipif(not env.has_nccl(), reason="need nccl")
+@pytest.mark.skipif(not env.has_multi_gpu(), reason="need multiple gpus")
+def test_callback():
+    """Simulate lazy loading of parameters in a callback
+
+    The output of a lazy parameter loading, which would accept a
+    callback to load the parameters.
+    """
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(private=True, s_tir=True)
+        def slice_A(
+            A: T.Buffer((4, 4), "int32"),
+            rank: T.int64,
+            A_sharded: T.Buffer((2, 4), "int32"),
+        ):
+            for i, j in T.grid(2, 4):
+                with T.sblock("slice_A"):
+                    vi, vj = T.axis.remap("SS", [i, j])
+                    A_sharded[vi, vj] = A[rank * 2 + vi, vj]
+
+        @T.prim_func(private=True, s_tir=True)
+        def slice_B(
+            B: T.Buffer((2, 2), "float32"),
+            rank: T.int64,
+            B_sharded: T.Buffer((2, 1), "float32"),
+        ):
+            for i in range(2):
+                with T.sblock("slice_B"):
+                    vi = T.axis.spatial(2, i)
+                    B_sharded[vi, 0] = B[vi, rank]
+
+        @R.function
+        def transform_params(
+            rank_arg: R.Prim("int64"),
+            fget_item: R.Callable([R.Any, R.Prim("int64")], R.Any),
+        ):
+            cls = Module
+
+            A = fget_item(R.str("A"), R.prim_value(0))
+            A = R.match_cast(A, R.Tensor([4, 4], "int32"))
+            A = R.call_tir(
+                cls.slice_A,
+                (A, rank_arg),
+                out_ty=R.Tensor([2, 4], "int32"),
+            )
+
+            B = fget_item(R.str("B"), R.prim_value(1))
+            B = R.match_cast(B, R.Tensor([2, 2], "float32"))
+            B = R.call_tir(
+                cls.slice_B,
+                (B, rank_arg),
+                out_ty=R.Tensor([2, 1], "float32"),
+            )
+
+            return (A, B)
+
+    pipeline = tvm.ir.transform.Sequential(
+        [
+            tvm.relax.transform.LegalizeOps(),
+            tvm.s_tir.dlight.ApplyDefaultSchedule(tvm.s_tir.dlight.gpu.Fallback()),
+        ],
+        name="pipeline",
+    )
+
+    with tvm.target.Target("cuda"):
+        mod = Module
+        mod = pipeline(mod)
+        built = tvm.compile(mod, "cuda")
+
+    num_shards = 2
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        temp_dir = pathlib.Path(temp_dir)
+
+        # TODO(Lunderberg): Update `disco.Session.load_vm_module` to
+        # allow a `tvm.runtime.Module` argument.  This would avoid the
+        # need for a temporary file.
+        shlib_path = temp_dir.joinpath("libtemp.so")
+        built.export_library(shlib_path)
+
+        def run_and_check():
+            session = tvm.runtime.disco.ProcessSession(num_workers=num_shards)
+            try:
+                session.import_python_module("tvm.exec.disco_worker")
+                session.init_ccl("nccl", *range(num_shards))
+
+                worker_device = session.get_global_func("runtime.disco.device")()
+                worker_id = session.get_global_func("runtime.disco.worker_rank")()
+                callback_maker = session.get_global_func("tests.disco.test_callback")
+                fget_item = callback_maker(worker_device)
+                vm = session.load_vm_module(shlib_path.as_posix())
+                transform_params = vm["transform_params"]
+
+                params = transform_params(worker_id, fget_item)
+
+                # Worker 0 is the same PID as the controlling scope, so
+                # `debug_get_from_remote(0)` returns the Tensor containing
+                # the output.
+                params_gpu0 = params.debug_get_from_remote(0)
+                assert params_gpu0[0].device == tvm.cuda(0)
+                assert params_gpu0[1].device == tvm.cuda(0)
+                np.testing.assert_array_equal(
+                    params_gpu0[0].numpy(),
+                    [
+                        [0, 1, 2, 3],
+                        [4, 5, 6, 7],
+                    ],
+                )
+                np.testing.assert_array_equal(
+                    params_gpu0[1].numpy(),
+                    [[0], [2]],
+                )
+
+                # Worker 1 is a different PID altogether, so
+                # `debug_get_from_remote(1)` returns a new Tensor within the
+                # calling scope's PID.
+                params_gpu1 = params.debug_get_from_remote(1)
+                assert params_gpu1[0].device == tvm.cpu()
+                assert params_gpu1[1].device == tvm.cpu()
+                np.testing.assert_array_equal(
+                    params_gpu1[0].numpy(),
+                    [
+                        [8, 9, 10, 11],
+                        [12, 13, 14, 15],
+                    ],
+                )
+                np.testing.assert_array_equal(
+                    params_gpu1[1].numpy(),
+                    [[1], [3]],
+                )
+            finally:
+                session.shutdown()
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/disco/test_ccl.py

@@ -0,0 +1,708 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-docstring
+"""Tests for NCCL/RCCL"""
+
+import tempfile
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import get_global_func
+from tvm import relax as rx
+from tvm.runtime import disco as di
+from tvm.runtime.vm import VirtualMachine
+from tvm.s_tir import dlight as dl
+from tvm.script import relax as R
+
+if di is None:
+    pytest.skip("disco runtime is not available", allow_module_level=True)
+
+_all_session_kinds = [di.ThreadedSession, di.ProcessSession]
+_compiled_ccl = get_global_func("runtime.disco.compiled_ccl", allow_missing=True)
+if _compiled_ccl is None:
+    pytest.skip("Disco CCL is not enabled in this TVM build", allow_module_level=True)
+_ccl = [_compiled_ccl()]
+
+
+def create_device_target(ccl):
+    if ccl == "nccl":
+        dev = tvm.cuda(0)
+    else:
+        dev = tvm.rocm(0)
+    target = tvm.target.Target.from_device(dev)
+    return (dev, target)
+
+
+def _run_with_ccl_session(session_kind, ccl, devices, func, *, num_groups=1):
+    def run_and_check():
+        sess = session_kind(num_workers=len(devices), num_groups=num_groups)
+        try:
+            sess.init_ccl(ccl, *devices)
+            return func(sess)
+        finally:
+            sess.shutdown()
+
+    return tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+@pytest.mark.parametrize("ccl", _ccl)
+def test_init(session_kind, ccl):
+    devices = [0, 1]
+
+    def run_test(_sess):
+        pass
+
+    _run_with_ccl_session(session_kind, ccl, devices, run_test)
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+@pytest.mark.parametrize("ccl", _ccl)
+def test_allreduce(session_kind, ccl):
+    devices = [0, 1]
+    array_1 = np.arange(12, dtype="float32").reshape(3, 4)
+    array_2 = np.arange(start=1, stop=-11, step=-1, dtype="float32").reshape(3, 4)
+
+    def run_test(sess):
+        d_array = sess.empty((3, 4), "float32")
+        d_array.debug_copy_from(0, array_1)
+        d_array.debug_copy_from(1, array_2)
+        for op, np_op in [  # pylint: disable=invalid-name
+            ("sum", np.add),
+            ("prod", np.multiply),
+            ("min", np.minimum),
+            ("max", np.maximum),
+            ("avg", lambda a, b: (a + b) * 0.5),
+        ]:
+            dst_array = sess.empty((3, 4), "float32")
+            sess.allreduce(d_array, dst_array, op=op)
+            result = dst_array.debug_get_from_remote(0).numpy()
+            expected = np_op(array_1, array_2)
+            np.testing.assert_equal(result, expected)
+
+    _run_with_ccl_session(session_kind, ccl, devices, run_test)
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+@pytest.mark.parametrize("ccl", _ccl)
+def test_group_allreduce(session_kind, ccl):
+    devices = [0, 1, 2, 3]
+    array_1 = np.arange(12, dtype="float32").reshape(3, 4)
+    array_2 = np.arange(start=1, stop=-11, step=-1, dtype="float32").reshape(3, 4)
+    array_3 = np.arange(30, dtype="float32").reshape(5, 6)
+    array_4 = np.arange(start=1, stop=-29, step=-1, dtype="float32").reshape(5, 6)
+
+    def run_test(sess):
+        d_array_1 = sess.empty((3, 4), "float32")
+        d_array_2 = sess.empty((5, 6), "float32")
+        d_array_1.debug_copy_from(0, array_1)
+        d_array_1.debug_copy_from(1, array_2)
+        d_array_2.debug_copy_from(2, array_3)
+        d_array_2.debug_copy_from(3, array_4)
+        for op, np_op in [  # pylint: disable=invalid-name
+            ("sum", np.add),
+            ("prod", np.multiply),
+            ("min", np.minimum),
+            ("max", np.maximum),
+            ("avg", lambda a, b: (a + b) * 0.5),
+        ]:
+            dst_array_1 = sess.empty((3, 4), "float32")
+            dst_array_2 = sess.empty((5, 6), "float32")
+            sess.allreduce(d_array_1, dst_array_1, op=op, in_group=True)
+            sess.allreduce(d_array_2, dst_array_2, op=op, in_group=True)
+            result_1 = dst_array_1.debug_get_from_remote(0).numpy()
+            result_2 = dst_array_2.debug_get_from_remote(2).numpy()
+            expected_1 = np_op(array_1, array_2)
+            expected_2 = np_op(array_3, array_4)
+            np.testing.assert_equal(result_1, expected_1)
+            np.testing.assert_equal(result_2, expected_2)
+
+    _run_with_ccl_session(session_kind, ccl, devices, run_test, num_groups=2)
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+@pytest.mark.parametrize("ccl", _ccl)
+def test_allgather(session_kind, ccl):
+    devices = [0, 1]
+    array = np.arange(36, dtype="float32")
+
+    def run_test(sess):
+        d_src = sess.empty((3, 3, 2), "float32")
+        d_dst = sess.empty((3, 4, 3), "float32")
+        d_src.debug_copy_from(0, array[:18])
+        d_src.debug_copy_from(1, array[18:])
+        sess.allgather(d_src, d_dst)
+        np.testing.assert_equal(
+            d_dst.debug_get_from_remote(0).numpy(),
+            array.reshape(3, 4, 3),
+        )
+        np.testing.assert_equal(
+            d_dst.debug_get_from_remote(1).numpy(),
+            array.reshape(3, 4, 3),
+        )
+
+    _run_with_ccl_session(session_kind, ccl, devices, run_test)
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+@pytest.mark.parametrize("ccl", _ccl)
+def test_group_allgather(session_kind, ccl):
+    devices = [0, 1, 2, 3]
+    array_1 = np.arange(36, dtype="float32")
+    array_2 = np.arange(48, dtype="float32")
+
+    def run_test(sess):
+        d_src_1 = sess.empty((3, 3, 2), "float32")
+        d_dst_1 = sess.empty((3, 4, 3), "float32")
+        d_src_2 = sess.empty((2, 4, 3), "float32")
+        d_dst_2 = sess.empty((2, 6, 4), "float32")
+        d_src_1.debug_copy_from(0, array_1[:18])
+        d_src_1.debug_copy_from(1, array_1[18:])
+        d_src_2.debug_copy_from(2, array_2[:24])
+        d_src_2.debug_copy_from(3, array_2[24:])
+        sess.allgather(d_src_1, d_dst_1, in_group=True)
+        sess.allgather(d_src_2, d_dst_2, in_group=True)
+        np.testing.assert_equal(
+            d_dst_1.debug_get_from_remote(0).numpy(),
+            array_1.reshape(3, 4, 3),
+        )
+        np.testing.assert_equal(
+            d_dst_1.debug_get_from_remote(1).numpy(),
+            array_1.reshape(3, 4, 3),
+        )
+        np.testing.assert_equal(
+            d_dst_2.debug_get_from_remote(2).numpy(),
+            array_2.reshape(2, 6, 4),
+        )
+        np.testing.assert_equal(
+            d_dst_2.debug_get_from_remote(3).numpy(),
+            array_2.reshape(2, 6, 4),
+        )
+
+    _run_with_ccl_session(session_kind, ccl, devices, run_test, num_groups=2)
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+@pytest.mark.parametrize("ccl", _ccl)
+@pytest.mark.parametrize("use_explicit_output", [True, False])
+def test_broadcast(session_kind, ccl, use_explicit_output):
+    devices = [0, 1]
+    array = np.arange(12, dtype="float32").reshape(3, 4)
+
+    def run_test(sess):
+        if use_explicit_output:
+            src_array = sess.empty((3, 4), "float32", worker0_only=True)
+            src_array.debug_copy_from(0, array)
+            dst_array = sess.empty((3, 4), "float32")
+            sess.broadcast_from_worker0(src_array, dst_array)
+        else:
+            dst_array = sess.broadcast(array)
+
+        result = dst_array.debug_get_from_remote(1).numpy()
+        np.testing.assert_equal(result, array)
+
+    _run_with_ccl_session(session_kind, ccl, devices, run_test)
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+@pytest.mark.parametrize("ccl", _ccl)
+def test_group_broadcast(session_kind, ccl):
+    devices = [0, 1, 2, 3]
+    array_1 = np.arange(12, dtype="float32").reshape(3, 4)
+    array_2 = np.multiply(array_1, -1)
+
+    def run_test(sess):
+        src_array = sess.empty((3, 4), "float32", worker0_only=True, in_group=True)
+        src_array.debug_copy_from(0, array_1)
+        src_array.debug_copy_from(2, array_2)
+        dst_array = sess.empty((3, 4), "float32")
+        sess.broadcast_from_worker0(src_array, dst_array)
+
+        result_1 = dst_array.debug_get_from_remote(1).numpy()
+        np.testing.assert_equal(result_1, array_1)
+
+        result_3 = dst_array.debug_get_from_remote(3).numpy()
+        np.testing.assert_equal(result_3, array_2)
+
+    _run_with_ccl_session(session_kind, ccl, devices, run_test, num_groups=2)
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+@pytest.mark.parametrize("ccl", _ccl)
+@pytest.mark.parametrize("use_explicit_output", [True, False])
+def test_scatter(session_kind, ccl, use_explicit_output, capfd):
+    devices = [0, 1]
+    array = np.arange(36, dtype="float32").reshape(2, 6, 3)
+
+    def run_test(sess):
+        if use_explicit_output:
+            d_src = sess.empty((2, 6, 3), "float32", worker0_only=True)
+            d_dst = sess.empty((6, 3), "float32")
+            d_src.debug_copy_from(0, array)
+            sess.scatter_from_worker0(d_src, d_dst)
+        else:
+            d_dst = sess.scatter(array)
+
+        np.testing.assert_equal(
+            d_dst.debug_get_from_remote(0).numpy(),
+            array[0, :, :],
+        )
+        np.testing.assert_equal(
+            d_dst.debug_get_from_remote(1).numpy(),
+            array[1, :, :],
+        )
+
+        captured = capfd.readouterr()
+        assert not captured.err, (
+            "No warning messages should be generated from disco.Session.scatter_from_worker0"
+        )
+
+    _run_with_ccl_session(session_kind, ccl, devices, run_test)
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+@pytest.mark.parametrize("ccl", _ccl)
+def test_group_scatter(session_kind, ccl, capfd):
+    devices = [0, 1, 2, 3]
+    array_1 = np.arange(36, dtype="float32").reshape(2, 6, 3)
+    array_2 = np.multiply(array_1, -1)
+
+    def run_test(sess):
+        d_src = sess.empty((2, 6, 3), "float32", worker0_only=True, in_group=True)
+        d_src.debug_copy_from(0, array_1)
+        d_src.debug_copy_from(2, array_2)
+        d_dst = sess.empty((6, 3), "float32")
+        sess.scatter_from_worker0(d_src, d_dst)
+
+        np.testing.assert_equal(
+            d_dst.debug_get_from_remote(0).numpy(),
+            array_1[0, :, :],
+        )
+        np.testing.assert_equal(
+            d_dst.debug_get_from_remote(1).numpy(),
+            array_1[1, :, :],
+        )
+        np.testing.assert_equal(
+            d_dst.debug_get_from_remote(2).numpy(),
+            array_2[0, :, :],
+        )
+        np.testing.assert_equal(
+            d_dst.debug_get_from_remote(3).numpy(),
+            array_2[1, :, :],
+        )
+
+        captured = capfd.readouterr()
+        assert not captured.err, (
+            "No warning messages should be generated from disco.Session.scatter_from_worker0"
+        )
+
+    _run_with_ccl_session(session_kind, ccl, devices, run_test, num_groups=2)
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+@pytest.mark.parametrize("ccl", _ccl)
+def test_scatter_with_implicit_reshape(session_kind, ccl, capfd):
+    """Scatter may perform an implicit reshape
+
+    Scattering elements to the workers requires the total number of
+    elements to be divisible by the number of workers.  It does not
+    necessarily correspond to scattering across the outermost
+    dimension.  Here, the number of workers (2) and the outermost
+    dimension (3) are not divisible, but the scatter may still be
+    performed.
+
+    This is only allowed when the caller explicitly uses the
+    `sess.scatter_from_worker0` method, and is not allowed in
+    `sess.scatter` method.  Because the `sess.scatter` method may
+    perform an allocation on the disco workers, it requires that the
+    scatter occur across the outermost dimension.
+
+    """
+    devices = [0, 1]
+    array = np.arange(36, dtype="float32").reshape(3, 4, 3)
+
+    def run_test(sess):
+        d_src = sess.empty((3, 4, 3), "float32", worker0_only=True)
+        d_dst = sess.empty((3, 3, 2), "float32")
+        d_src.debug_copy_from(0, array)
+        sess.scatter_from_worker0(d_src, d_dst)
+
+        np.testing.assert_equal(
+            d_dst.debug_get_from_remote(0).numpy(),
+            array.flat[:18].reshape(3, 3, 2),
+        )
+        np.testing.assert_equal(
+            d_dst.debug_get_from_remote(1).numpy(),
+            array.flat[18:].reshape(3, 3, 2),
+        )
+
+        captured = capfd.readouterr()
+        assert not captured.err, (
+            "No warning messages should be generated from disco.Session.scatter_from_worker0"
+        )
+
+    _run_with_ccl_session(session_kind, ccl, devices, run_test)
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+@pytest.mark.parametrize("ccl", _ccl)
+def test_gather(session_kind, ccl, capfd):
+    devices = [0, 1]
+    array = np.arange(36, dtype="float32")
+
+    def run_test(sess):
+        d_src = sess.empty((3, 3, 2), "float32")
+        d_dst = sess.empty((3, 4, 3), "float32", worker0_only=True)
+        d_src.debug_copy_from(0, array[:18])
+        d_src.debug_copy_from(1, array[18:])
+        sess.gather_to_worker0(d_src, d_dst)
+        np.testing.assert_equal(
+            d_dst.debug_get_from_remote(0).numpy(),
+            array.reshape(3, 4, 3),
+        )
+
+        captured = capfd.readouterr()
+        assert not captured.err, (
+            "No warning messages should be generated from disco.Session.gather_to_worker0"
+        )
+
+    _run_with_ccl_session(session_kind, ccl, devices, run_test)
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+@pytest.mark.parametrize("ccl", _ccl)
+def test_group_gather(session_kind, ccl, capfd):
+    devices = [0, 1, 2, 3]
+    array_1 = np.arange(36, dtype="float32")
+    array_2 = np.multiply(array_1, -1)
+
+    def run_test(sess):
+        d_src = sess.empty((3, 3, 2), "float32")
+        d_dst = sess.empty((3, 4, 3), "float32", worker0_only=True, in_group=True)
+        d_src.debug_copy_from(0, array_1[:18])
+        d_src.debug_copy_from(1, array_1[18:])
+        d_src.debug_copy_from(2, array_2[:18])
+        d_src.debug_copy_from(3, array_2[18:])
+        sess.gather_to_worker0(d_src, d_dst)
+        np.testing.assert_equal(
+            d_dst.debug_get_from_remote(0).numpy(),
+            array_1.reshape(3, 4, 3),
+        )
+        np.testing.assert_equal(
+            d_dst.debug_get_from_remote(2).numpy(),
+            array_2.reshape(3, 4, 3),
+        )
+
+        captured = capfd.readouterr()
+        assert not captured.err, (
+            "No warning messages should be generated from disco.Session.gather_to_worker0"
+        )
+
+    _run_with_ccl_session(session_kind, ccl, devices, run_test, num_groups=2)
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+@pytest.mark.parametrize("ccl", _ccl)
+def test_send_to_next_group_receive_from_prev_group(session_kind, ccl):
+    devices = [0, 1, 2, 3]
+    array_1 = np.arange(12, dtype="float32").reshape(3, 4)
+    array_2 = np.arange(start=1, stop=-11, step=-1, dtype="float32").reshape(3, 4)
+
+    def run_test(sess):
+        d_array = sess.empty((3, 4), "float32")
+        d_array.debug_copy_from(0, array_1)
+        d_array.debug_copy_from(1, array_2)
+        sess.get_global_func(
+            "runtime.disco." + ccl + ".test_send_to_next_group_recv_from_prev_group"
+        )(d_array)
+
+        result_1 = d_array.debug_get_from_remote(2).numpy()
+        result_2 = d_array.debug_get_from_remote(3).numpy()
+        np.testing.assert_equal(result_1, array_1)
+        np.testing.assert_equal(result_2, array_2)
+
+    _run_with_ccl_session(session_kind, ccl, devices, run_test, num_groups=2)
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+@pytest.mark.parametrize("ccl", _ccl)
+def test_worker2_send_to_worker0(session_kind, ccl):
+    devices = [0, 1, 2, 3]
+    array = np.arange(start=1, stop=-11, step=-1, dtype="float32").reshape(3, 4)
+
+    def run_test(sess):
+        d_array = sess.empty((3, 4), "float32")
+        d_array.debug_copy_from(2, array)
+        sess.get_global_func("runtime.disco." + ccl + ".test_worker2_sends_to_worker0")(d_array)
+
+        result = d_array.debug_get_from_remote(0).numpy()
+        np.testing.assert_equal(result, array)
+
+    _run_with_ccl_session(session_kind, ccl, devices, run_test, num_groups=2)
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+@pytest.mark.parametrize("ccl", _ccl)
+def test_mlp(session_kind, ccl):  # pylint: disable=too-many-locals
+    devices = [0, 1]
+
+    # pylint: disable=invalid-name
+    @tvm.script.ir_module
+    class MLP:  # pylint: disable=too-few-public-methods
+        @R.function
+        def main(
+            x: R.Tensor((128, 128), "float32"),
+            W1: R.Tensor((128, 128), "float32"),
+            W2: R.Tensor((128, 128), "float32"),
+        ) -> R.Tensor((128, 128), "float32"):
+            R.func_attr({"global_symbol": "main"})
+            with R.dataflow():
+                lv0: R.Tensor((128, 128), "float32") = R.matmul(x, W1)
+                lv1: R.Tensor((128, 128), "float32") = R.nn.gelu(lv0)
+                lv2: R.Tensor((128, 128), "float32") = R.matmul(lv1, W2)
+                R.output(lv2)
+            return lv2
+
+    @tvm.script.ir_module
+    class ShardedMLP:  # pylint: disable=too-few-public-methods
+        @R.function
+        def main(
+            x: R.Tensor((128, 128), "float32"),
+            W1: R.Tensor((128, 64), "float32"),  # shard along axis 1
+            W2: R.Tensor((64, 128), "float32"),  # shard along axis 0
+        ) -> R.Tensor((128, 128), "float32"):
+            R.func_attr({"global_symbol": "main"})
+            with R.dataflow():
+                broadcast_x: R.Tensor((128, 128), "float32") = R.ccl.broadcast_from_worker0(x)
+                lv0: R.Tensor((128, 64), "float32") = R.matmul(broadcast_x, W1)
+                lv1: R.Tensor((128, 64), "float32") = R.nn.gelu(lv0)
+                lv2: R.Tensor((128, 128), "float32") = R.matmul(lv1, W2)
+                lv3: R.Tensor((128, 128), "float32") = R.ccl.allreduce(lv2, "sum")
+                R.output(lv3)
+            return lv3
+
+    # pylint: enable=invalid-name
+    dev, target = create_device_target(ccl)
+
+    def relax_build(mod, target):
+        with target:
+            mod = rx.get_pipeline("zero")(mod)  # pylint: disable=no-value-for-parameter
+            mod = dl.ApplyDefaultSchedule(  # pylint: disable=not-callable
+                dl.gpu.Matmul(),
+                dl.gpu.GEMV(),
+                dl.gpu.Reduction(),
+                dl.gpu.GeneralReduction(),
+                dl.gpu.Fallback(),
+            )(mod)
+            return tvm.compile(mod, target=target)
+
+    # pylint: disable=invalid-name
+    X = np.random.randn(128, 128).astype("float32")
+    W1 = np.random.randn(128, 128).astype("float32")
+    W2 = np.random.randn(128, 128).astype("float32")
+    expected_ex = relax_build(MLP, target)
+    sharded_ex = relax_build(ShardedMLP, target)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        path = tmpdir + "/test.so"
+        sharded_ex.export_library(path)
+
+        def run_test(sess):
+            Y_expected = VirtualMachine(expected_ex, device=dev)["main"](
+                tvm.runtime.tensor(X, device=dev),
+                tvm.runtime.tensor(W1, device=dev),
+                tvm.runtime.tensor(W2, device=dev),
+            ).numpy()
+
+            mod = sess.load_vm_module(path)
+            d_X = sess.empty((128, 128), "float32")
+            d_W1 = sess.empty((128, 64), "float32")
+            d_W2 = sess.empty((64, 128), "float32")
+
+            d_X.debug_copy_from(0, X)
+            d_W1.debug_copy_from(0, W1[:, :64])
+            d_W1.debug_copy_from(1, W1[:, 64:])
+            d_W2.debug_copy_from(0, W2[:64, :])
+            d_W2.debug_copy_from(1, W2[64:, :])
+            d_Y = mod["main"](d_X, d_W1, d_W2)
+            Y_result = tvm.runtime.empty((128, 128), "float32", device=dev)
+            sess.copy_from_worker_0(Y_result, d_Y)
+            sess.sync_worker_0()
+            Y_result = Y_result.numpy()
+            tvm.testing.assert_allclose(Y_result, Y_expected, rtol=1e-4, atol=1e-4)
+
+        _run_with_ccl_session(session_kind, ccl, devices, run_test)
+    # pylint: enable=invalid-name
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+@pytest.mark.parametrize("ccl", _ccl)
+def test_attention(session_kind, ccl):  # pylint: disable=too-many-locals,too-many-statements
+    devices = [0, 1]
+
+    # pylint: disable=invalid-name
+    @tvm.script.ir_module
+    class Attention:  # pylint: disable=too-few-public-methods
+        @R.function
+        def main(  # pylint: disable=too-many-locals
+            x: R.Tensor((1, 10, 128), "float32"),
+            Wq: R.Tensor((128, 512), "float32"),
+            Wk: R.Tensor((128, 512), "float32"),
+            Wv: R.Tensor((128, 512), "float32"),
+            Wo: R.Tensor((512, 128), "float32"),
+        ) -> R.Tensor((128, 128), "float32"):
+            R.func_attr({"global_symbol": "main"})
+            with R.dataflow():
+                # q
+                lv0: R.Tensor((1, 10, 512), "float32") = R.matmul(x, Wq)
+                lv1: R.Tensor((1, 10, 8, 64), "float32") = R.reshape(lv0, [1, 10, 8, 64])
+                lv2: R.Tensor((1, 8, 10, 64), "float32") = R.permute_dims(lv1, [0, 2, 1, 3])
+                # k
+                lv3: R.Tensor((1, 10, 512), "float32") = R.matmul(x, Wk)
+                lv4: R.Tensor((1, 10, 8, 64), "float32") = R.reshape(lv3, [1, 10, 8, 64])
+                lv5: R.Tensor((1, 8, 10, 64), "float32") = R.permute_dims(lv4, [0, 2, 1, 3])
+                # v
+                lv6: R.Tensor((1, 10, 512), "float32") = R.matmul(x, Wv)
+                lv7: R.Tensor((1, 10, 8, 64), "float32") = R.reshape(lv6, [1, 10, 8, 64])
+                lv8: R.Tensor((1, 8, 10, 64), "float32") = R.permute_dims(lv7, [0, 2, 1, 3])
+                # softmax(q @ k / sqrt(dk))
+                lv9: R.Tensor((1, 8, 64, 10), "float32") = R.permute_dims(lv5, [0, 1, 3, 2])
+                lv10: R.Tensor((1, 8, 10, 10), "float32") = R.matmul(lv2, lv9)
+                lv11: R.Tensor((1, 8, 10, 10), "float32") = R.multiply(
+                    lv10, R.const(1 / 8, "float32")
+                )
+                lv12: R.Tensor((1, 8, 10, 10), "float32") = R.nn.softmax(lv11, axis=-1)
+                # attn_weight @ v
+                lv13: R.Tensor((1, 8, 10, 64), "float32") = R.matmul(lv12, lv8)
+                lv14: R.Tensor((1, 10, 8, 64), "float32") = R.permute_dims(lv13, [0, 2, 1, 3])
+                lv15: R.Tensor((1, 10, 512), "float32") = R.reshape(lv14, [1, 10, 512])
+                # attn_output @ o
+                lv16: R.Tensor((1, 10, 128), "float32") = R.matmul(lv15, Wo)
+                R.output(lv16)
+            return lv16
+
+    @tvm.script.ir_module
+    class ShardedAttention:  # pylint: disable=too-few-public-methods
+        @R.function
+        def main(  # pylint: disable=too-many-locals
+            x: R.Tensor((1, 10, 128), "float32"),
+            Wq: R.Tensor((128, 256), "float32"),  # shard along axis 1
+            Wk: R.Tensor((128, 256), "float32"),  # shard along axis 1
+            Wv: R.Tensor((128, 256), "float32"),  # shard along axis 1
+            Wo: R.Tensor((256, 128), "float32"),  # shard along axis 0
+        ) -> R.Tensor((128, 128), "float32"):
+            R.func_attr({"global_symbol": "main"})
+            with R.dataflow():
+                broadcast_x: R.Tensor((1, 10, 128), "float32") = R.ccl.broadcast_from_worker0(x)
+                # q
+                lv0: R.Tensor((1, 10, 256), "float32") = R.matmul(broadcast_x, Wq)
+                lv1: R.Tensor((1, 10, 4, 64), "float32") = R.reshape(lv0, [1, 10, 4, 64])
+                lv2: R.Tensor((1, 4, 10, 64), "float32") = R.permute_dims(lv1, [0, 2, 1, 3])
+                # k
+                lv3: R.Tensor((1, 10, 256), "float32") = R.matmul(broadcast_x, Wk)
+                lv4: R.Tensor((1, 10, 4, 64), "float32") = R.reshape(lv3, [1, 10, 4, 64])
+                lv5: R.Tensor((1, 4, 10, 64), "float32") = R.permute_dims(lv4, [0, 2, 1, 3])
+                # v
+                lv6: R.Tensor((1, 10, 256), "float32") = R.matmul(broadcast_x, Wv)
+                lv7: R.Tensor((1, 10, 4, 64), "float32") = R.reshape(lv6, [1, 10, 4, 64])
+                lv8: R.Tensor((1, 4, 10, 64), "float32") = R.permute_dims(lv7, [0, 2, 1, 3])
+                # softmax(q @ k / sqrt(dk))
+                lv9: R.Tensor((1, 4, 64, 10), "float32") = R.permute_dims(lv5, [0, 1, 3, 2])
+                lv10: R.Tensor((1, 4, 10, 10), "float32") = R.matmul(lv2, lv9)
+                lv11: R.Tensor((1, 4, 10, 10), "float32") = R.multiply(
+                    lv10, R.const(1 / 8, "float32")
+                )
+                lv12: R.Tensor((1, 4, 10, 10), "float32") = R.nn.softmax(lv11, axis=-1)
+                # attn_weight @ v
+                lv13: R.Tensor((1, 4, 10, 64), "float32") = R.matmul(lv12, lv8)
+                lv14: R.Tensor((1, 10, 4, 64), "float32") = R.permute_dims(lv13, [0, 2, 1, 3])
+                lv15: R.Tensor((1, 10, 256), "float32") = R.reshape(lv14, [1, 10, 256])
+                # attn_output @ o
+                lv16: R.Tensor((1, 10, 128), "float32") = R.matmul(lv15, Wo)
+                lv17: R.Tensor((1, 10, 128), "float32") = R.ccl.allreduce(lv16, "sum")
+                R.output(lv17)
+            return lv17
+
+    # pylint: enable=invalid-name
+    dev, target = create_device_target(ccl)
+
+    def relax_build(mod, target):
+        with target:
+            mod = rx.get_pipeline("zero")(mod)  # pylint: disable=no-value-for-parameter
+            mod = dl.ApplyDefaultSchedule(  # pylint: disable=not-callable
+                dl.gpu.Matmul(),
+                dl.gpu.GEMV(),
+                dl.gpu.Reduction(),
+                dl.gpu.GeneralReduction(),
+                dl.gpu.Fallback(),
+            )(mod)
+            return tvm.compile(mod, target=target)
+
+    # pylint: disable=invalid-name
+    X = np.random.randn(1, 10, 128).astype("float32")
+    Wq = np.random.randn(128, 512).astype("float32")
+    Wk = np.random.randn(128, 512).astype("float32")
+    Wv = np.random.randn(128, 512).astype("float32")
+    Wo = np.random.randn(512, 128).astype("float32")
+    expected_ex = relax_build(Attention, target)
+    sharded_ex = relax_build(ShardedAttention, target)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        path = tmpdir + "/test.so"
+        sharded_ex.export_library(path)
+
+        def run_test(sess):
+            Y_expected = VirtualMachine(expected_ex, device=dev)["main"](
+                tvm.runtime.tensor(X, device=dev),
+                tvm.runtime.tensor(Wq, device=dev),
+                tvm.runtime.tensor(Wk, device=dev),
+                tvm.runtime.tensor(Wv, device=dev),
+                tvm.runtime.tensor(Wo, device=dev),
+            ).numpy()
+
+            mod = sess.load_vm_module(path)
+            d_X = sess.empty((1, 10, 128), "float32")
+            d_Wq = sess.empty((128, 256), "float32")
+            d_Wk = sess.empty((128, 256), "float32")
+            d_Wv = sess.empty((128, 256), "float32")
+            d_Wo = sess.empty((256, 128), "float32")
+
+            d_X.debug_copy_from(0, X)
+            d_Wq.debug_copy_from(0, Wq[:, :256])
+            d_Wq.debug_copy_from(1, Wq[:, 256:])
+            d_Wk.debug_copy_from(0, Wk[:, :256])
+            d_Wk.debug_copy_from(1, Wk[:, 256:])
+            d_Wv.debug_copy_from(0, Wv[:, :256])
+            d_Wv.debug_copy_from(1, Wv[:, 256:])
+            d_Wo.debug_copy_from(0, Wo[:256, :])
+            d_Wo.debug_copy_from(1, Wo[256:, :])
+            d_Y = mod["main"](d_X, d_Wq, d_Wk, d_Wv, d_Wo)
+            Y_result = tvm.runtime.empty((1, 10, 128), "float32", device=dev)
+            sess.copy_from_worker_0(Y_result, d_Y)
+            sess.sync_worker_0()
+            Y_result = Y_result.numpy()
+            tvm.testing.assert_allclose(Y_result, Y_expected, rtol=1e-3, atol=1e-3)
+
+        _run_with_ccl_session(session_kind, ccl, devices, run_test)
+    # pylint: enable=invalid-name
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/disco/test_custom_allreduce.py

@@ -0,0 +1,88 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import enum
+from functools import reduce
+from itertools import product
+
+import numpy as np
+import pytest
+from tvm_ffi import Shape
+
+import tvm
+import tvm.testing
+from tvm.runtime import DataType, disco
+
+if disco is None:
+    pytest.skip("disco runtime is not available", allow_module_level=True)
+
+
+class AllReduceStrategyType(enum.IntEnum):
+    RING = 0
+    ONESHOT = 1
+    TWOSHOT = 2
+    AUTO = 3
+
+
+_shapes = [(2, 3), (3, 4), (128, 128)]
+
+_strategies = [
+    AllReduceStrategyType.RING,
+    AllReduceStrategyType.ONESHOT,
+    AllReduceStrategyType.TWOSHOT,
+    AllReduceStrategyType.AUTO,
+]
+
+_compiled_ccl = tvm.get_global_func("runtime.disco.compiled_ccl", allow_missing=True)
+if _compiled_ccl is None:
+    pytest.skip("Disco CCL is not enabled in this TVM build", allow_module_level=True)
+_ccl = [ccl for ccl in _compiled_ccl() if ccl == "nccl"]
+
+
+@pytest.mark.parametrize("shape", _shapes)
+@pytest.mark.parametrize("ccl", _ccl)
+@pytest.mark.parametrize("strategy", _strategies)
+def test_allreduce(shape, ccl, strategy):
+    devices = [0, 1]
+    sess = disco.ProcessSession(num_workers=len(devices))
+    sess.init_ccl(ccl, *devices)
+
+    num_elements = reduce(lambda x, y: x * y, shape)
+    dtype = "float32"
+    falloc_ipc_storage = sess.get_global_func("runtime.disco.cuda_ipc.alloc_storage")
+    falloc_tensor = sess.get_global_func("vm.builtin.alloc_tensor")
+    fallreduce = sess.get_global_func("runtime.disco.cuda_ipc.custom_allreduce")
+    d_storage = sess.call_packed(falloc_ipc_storage, Shape(shape), DataType(dtype))
+    d_input = sess.call_packed(falloc_tensor, d_storage, 0, Shape(shape), DataType(dtype))
+
+    array_1 = np.arange(num_elements, dtype="float32").reshape(*shape)
+    array_2 = np.arange(start=1, stop=-(num_elements - 1), step=-1, dtype="float32").reshape(*shape)
+    d_input.debug_copy_from(0, array_1)
+    d_input.debug_copy_from(1, array_2)
+    d_output = sess.empty(shape, "float32")
+
+    sess.call_packed(fallreduce, d_input, strategy, d_output)
+    result_1 = d_output.debug_get_from_remote(0).numpy()
+    result_2 = d_output.debug_get_from_remote(1).numpy()
+    expected = np.add(array_1, array_2)
+    np.testing.assert_equal(result_1, expected)
+    np.testing.assert_equal(result_2, expected)
+
+
+if __name__ == "__main__":
+    for shape, strategy in product(_shapes, _strategies):
+        test_allreduce(shape, "nccl", strategy)

tests/python/disco/test_loader.py

@@ -0,0 +1,496 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F401, F841
+"""Test sharded loader"""
+
+# pylint: disable=missing-docstring
+import json
+import tempfile
+
+import numpy as np
+import pytest
+from tvm_ffi import Shape, register_global_func
+
+import tvm
+import tvm.testing
+from tvm import relax as rx
+from tvm.contrib import tvmjs
+from tvm.runtime import disco as di
+from tvm.s_tir import dlight as dl
+from tvm.script import ir as I
+from tvm.script import relax as R
+from tvm.target import Target
+from tvm.testing import env
+
+# `runtime.disco.compiled_ccl` is registered together with the CCL runtime
+# functions, so its absence means the disco CCL runtime is not in this build.
+_compiled_ccl = tvm.get_global_func("runtime.disco.compiled_ccl", allow_missing=True)
+if _compiled_ccl is None or _compiled_ccl() != "nccl":
+    pytest.skip("Disco NCCL support is not available", allow_module_level=True)
+
+# All tests in this file shard across two GPUs.
+pytestmark = [
+    pytest.mark.skipif(not env.has_multi_gpu(), reason="need multiple gpus"),
+]
+
+
+def _run_with_nccl_session(devices, func):
+    def run_and_check():
+        sess = di.ThreadedSession(num_workers=len(devices))
+        try:
+            sess.init_ccl("nccl", *devices)
+            return func(sess)
+        finally:
+            sess.shutdown()
+
+    return tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@register_global_func("tests.disco.shard_dim_0", override=True)
+def _shard_dim_0(src, num_shards, tgt):
+    s_0, s_1 = src.shape
+    tgt.copyfrom(src.numpy().reshape(num_shards, s_0 // num_shards, s_1))
+
+
+@register_global_func("tests.disco.shard_dim_1", override=True)
+def _shard_dim_1(src, num_shards, tgt):
+    s_0, s_1 = src.shape
+    tgt.copyfrom(src.numpy().reshape(s_0, num_shards, s_1 // num_shards).transpose(1, 0, 2))
+
+
+@register_global_func("tests.disco.shard_qkv_0", override=True)
+def _shard_qkv_0(src, num_shards, q_heads, kv_heads, tgt):
+    total_dim, hidden_size = src.shape
+    head_dim = total_dim // (q_heads + kv_heads + kv_heads)
+    q_dim = q_heads * head_dim
+    kv_dim = kv_heads * head_dim
+    w_q = src.numpy()[:q_dim, :].reshape(
+        num_shards,
+        q_heads // num_shards,
+        head_dim,
+        hidden_size,
+    )
+    w_k = src.numpy()[q_dim : q_dim + kv_dim, :].reshape(
+        num_shards,
+        kv_heads // num_shards,
+        head_dim,
+        hidden_size,
+    )
+    w_v = src.numpy()[q_dim + kv_dim :, :].reshape(
+        num_shards,
+        kv_heads // num_shards,
+        head_dim,
+        hidden_size,
+    )
+    w_qkv = np.concatenate([w_q, w_k, w_v], axis=1)
+    tgt.copyfrom(w_qkv)
+
+
+@register_global_func("tests.disco.shard_qkv_1", override=True)
+def _shard_qkv_1(src, tgt):
+    s, _, _, h = src.shape  # pylint: disable=invalid-name
+    tgt.copyfrom(src.numpy().reshape(s, -1, h))
+
+
+def _create_loader(sess, path, param_dict, shard_info):
+    path_tensor_cache = path + "/tensor-cache.json"
+    tvmjs.dump_tensor_cache(param_dict, path, encode_format="raw")
+    with open(path_tensor_cache, encoding="utf-8") as i_f:
+        tensor_cache = i_f.read()
+    loader_create = sess.get_global_func("runtime.disco.ShardLoader")
+    loader = loader_create(path_tensor_cache, tensor_cache, json.dumps(shard_info), None)
+    return loader
+
+
+def _simulate_presharded_weights(base_path, param_dict, num_shards, shard_info):
+    """Create fake weights to simulate those produced MLC-LLM's pre-sharding"""
+
+    sharded_params = {}
+
+    for key, ndarray in param_dict.items():
+        assert key in shard_info, f"ShardInfo lacks shard info about param: {key}"
+        shard_dim = shard_info[key]
+        sharded_params[key] = [
+            tvm.runtime.tensor(np_shard)
+            for np_shard in np.split(ndarray, num_shards, axis=shard_dim)
+        ]
+
+    # Re-order so that the parameter order is sorted first by shard,
+    # then by parameter.  This matches the ordering used by MLC-LLM,
+    # and avoids having *.bin files that must be accessed by more than
+    # one worker.
+    sharded_params = {
+        f"{key}_shard-{i + 1}-of-{num_shards}": shards[i]
+        for i in range(num_shards)
+        for key, shards in sharded_params.items()
+    }
+
+    tvmjs.dump_tensor_cache(
+        sharded_params,
+        base_path,
+        encode_format="raw",
+    )
+
+
+def test_load_shard():
+    devices = [0, 1]
+    num_shards = len(devices)
+    param_dict = {
+        "x_0": np.random.uniform(size=[64, 128]).astype("float16"),
+        "x_1": np.random.uniform(size=[32, 128]).astype("float32"),
+    }
+    shard_info = {
+        "x_0": [
+            [
+                "tests.disco.shard_dim_1",
+                [(num_shards, 64, 64), "float16"],
+                num_shards,
+            ],
+        ],
+        "x_1": [
+            [
+                "tests.disco.shard_dim_0",
+                [(num_shards, 16, 128), "float32"],
+                num_shards,
+            ]
+        ],
+    }
+    with tempfile.TemporaryDirectory() as path:
+
+        def run_test(sess):
+            loader = _create_loader(sess, path, param_dict, shard_info)
+            loader_load = sess.get_global_func("runtime.disco.ShardLoaderLoad")
+            d_0 = loader_load(loader, Shape([0]))
+            d_1 = loader_load(loader, Shape([1]))
+            np.testing.assert_equal(
+                param_dict["x_0"][:, 0:64],
+                d_0.debug_get_from_remote(0).numpy(),
+            )
+            np.testing.assert_equal(
+                param_dict["x_0"][:, 64:128],
+                d_0.debug_get_from_remote(1).numpy(),
+            )
+            np.testing.assert_equal(
+                param_dict["x_1"][0:16, :],
+                d_1.debug_get_from_remote(0).numpy(),
+            )
+            np.testing.assert_equal(
+                param_dict["x_1"][16:32, :],
+                d_1.debug_get_from_remote(1).numpy(),
+            )
+
+        _run_with_nccl_session(devices, run_test)
+
+
+def _create_presharded_loader(sess, path):
+    path_tensor_cache = path + "/tensor-cache.json"
+    with open(path_tensor_cache, encoding="utf-8") as i_f:
+        tensor_cache = i_f.read()
+    loader_create = sess.get_global_func("runtime.disco.ShardLoader")
+    loader = loader_create(path_tensor_cache, tensor_cache, json.dumps({}), None)
+    return loader
+
+
+def test_load_presharded():
+    devices = [0, 1]
+    param_dict = {
+        "x_0": np.random.uniform(size=[64, 128]).astype("float16"),
+        "x_1": np.random.uniform(size=[32, 128]).astype("float32"),
+    }
+    shard_info = {
+        "x_0": 1,
+        "x_1": 0,
+    }
+
+    with tempfile.TemporaryDirectory() as path:
+        _simulate_presharded_weights(path, param_dict, len(devices), shard_info)
+
+        def run_test(sess):
+            loader = _create_presharded_loader(sess, path)
+            loader_load = sess.get_global_func("runtime.disco.ShardLoaderLoadPresharded")
+
+            d_0 = loader_load(loader, Shape([0]))
+            d_1 = loader_load(loader, Shape([1]))
+
+            np.testing.assert_equal(
+                param_dict["x_0"][:, 0:64],
+                d_0.debug_get_from_remote(0).numpy(),
+            )
+            np.testing.assert_equal(
+                param_dict["x_0"][:, 64:128],
+                d_0.debug_get_from_remote(1).numpy(),
+            )
+            np.testing.assert_equal(
+                param_dict["x_1"][0:16, :],
+                d_1.debug_get_from_remote(0).numpy(),
+            )
+            np.testing.assert_equal(
+                param_dict["x_1"][16:32, :],
+                d_1.debug_get_from_remote(1).numpy(),
+            )
+
+        _run_with_nccl_session(devices, run_test)
+
+
+def test_load_shard_in_relax():
+    devices = [0, 1]
+    num_shards = len(devices)
+    param_dict = {
+        "x_0": np.random.uniform(size=[64, 128]).astype("float16"),
+        "x_1": np.random.uniform(size=[32, 128]).astype("float32"),
+    }
+    shard_info = {
+        "x_0": [
+            [
+                "tests.disco.shard_dim_1",
+                [(num_shards, 64, 64), "float16"],
+                num_shards,
+            ],
+        ],
+        "x_1": [
+            [
+                "tests.disco.shard_dim_0",
+                [(num_shards, 16, 128), "float32"],
+                num_shards,
+            ]
+        ],
+    }
+
+    # pylint: disable=invalid-name
+    @I.ir_module
+    class Module:  # pylint: disable=too-few-public-methods
+        @R.function
+        def main(
+            loader: R.Any,
+        ) -> R.Tuple(R.Tensor((64, 64), "float32"), R.Tensor((16, 128), "float32")):
+            R.func_attr({"global_symbol": "main"})
+            with R.dataflow():
+                lv0: R.Tensor((64, 64), "float32") = R.call_pure_packed(
+                    "runtime.disco.ShardLoaderLoad",
+                    loader,
+                    R.shape([0]),
+                    ty_args=R.Tensor((64, 64), "float32"),
+                )
+                lv1: R.Tensor((16, 128), "float32") = R.call_pure_packed(
+                    "runtime.disco.ShardLoaderLoad",
+                    loader,
+                    R.shape([1]),
+                    ty_args=R.Tensor((16, 128), "float32"),
+                )
+                lv2 = R.tuple(lv0, lv1)
+                R.output(lv2)
+            return lv2
+
+    # pylint: enable=invalid-name
+    def relax_build(mod, target):
+        with target:
+            mod = rx.get_pipeline("zero")(mod)  # pylint: disable=no-value-for-parameter
+            return tvm.compile(mod, target="cuda")
+
+    target = Target(
+        {
+            "kind": "cuda",
+            "max_shared_memory_per_block": 49152,
+            "max_threads_per_block": 1024,
+            "thread_warp_size": 32,
+            "registers_per_block": 65536,
+            "arch": "sm_80",
+        }
+    )
+    with tempfile.TemporaryDirectory() as tmpdir:
+        dso_path = tmpdir + "/test.so"
+        relax_build(Module, target).export_library(dso_path)
+
+        def run_test(sess):
+            mod = sess.load_vm_module(dso_path)
+            loader = _create_loader(sess, tmpdir, param_dict, shard_info)
+            result = mod["main"](loader)
+            np.testing.assert_equal(
+                param_dict["x_0"][:, 0:64],
+                result.debug_get_from_remote(0)[0].numpy(),
+            )
+            np.testing.assert_equal(
+                param_dict["x_0"][:, 64:128],
+                result.debug_get_from_remote(1)[0].numpy(),
+            )
+            np.testing.assert_equal(
+                param_dict["x_1"][0:16, :],
+                result.debug_get_from_remote(0)[1].numpy(),
+            )
+            np.testing.assert_equal(
+                param_dict["x_1"][16:32, :],
+                result.debug_get_from_remote(1)[1].numpy(),
+            )
+
+        _run_with_nccl_session(devices, run_test)
+
+
+def test_load_shard_all():
+    devices = [0, 1]
+    num_shards = len(devices)
+    param_dict = {
+        "param_0": np.random.uniform(size=[64, 128]).astype("float16"),
+        "param_1": np.random.uniform(size=[32, 128]).astype("float32"),
+    }
+    shard_info = {
+        "param_0": [
+            [
+                "tests.disco.shard_dim_1",
+                [(num_shards, 64, 64), "float16"],
+                num_shards,
+            ],
+        ],
+        "param_1": [
+            [
+                "tests.disco.shard_dim_0",
+                [(2, 16, 128), "float32"],
+                num_shards,
+            ]
+        ],
+    }
+    with tempfile.TemporaryDirectory() as path:
+
+        def run_test(sess):
+            loader = _create_loader(sess, path, param_dict, shard_info)
+            loader_load = sess.get_global_func("runtime.disco.ShardLoaderLoadAll")
+            params = loader_load(loader)
+            p_0 = params.debug_get_from_remote(0)
+            p_1 = params.debug_get_from_remote(1)
+            np.testing.assert_equal(param_dict["param_0"][:, 0:64], p_0[0].numpy())
+            np.testing.assert_equal(param_dict["param_0"][:, 64:128], p_1[0].numpy())
+            np.testing.assert_equal(param_dict["param_1"][0:16, :], p_0[1].numpy())
+            np.testing.assert_equal(param_dict["param_1"][16:32, :], p_1[1].numpy())
+
+        _run_with_nccl_session(devices, run_test)
+
+
+def test_load_all_presharded():
+    devices = [0, 1]
+    num_shards = len(devices)
+    param_dict = {
+        "param_0": np.random.uniform(size=[64, 128]).astype("float16"),
+        "param_1": np.random.uniform(size=[32, 128]).astype("float32"),
+    }
+    shard_info = {
+        "param_0": 0,
+        "param_1": 1,
+    }
+    with tempfile.TemporaryDirectory() as path:
+        _simulate_presharded_weights(path, param_dict, len(devices), shard_info)
+
+        def run_test(sess):
+            loader = _create_presharded_loader(sess, path)
+            loader_load = sess.get_global_func("runtime.disco.ShardLoaderLoadAllPresharded")
+            params = loader_load(loader)
+
+            p_0 = params.debug_get_from_remote(0)
+            p_1 = params.debug_get_from_remote(1)
+
+            np.testing.assert_equal(param_dict["param_0"][0:32, :], p_0[0].numpy())
+            np.testing.assert_equal(param_dict["param_0"][32:64, :], p_1[0].numpy())
+            np.testing.assert_equal(param_dict["param_1"][:, 0:64], p_0[1].numpy())
+            np.testing.assert_equal(param_dict["param_1"][:, 64:128], p_1[1].numpy())
+
+        _run_with_nccl_session(devices, run_test)
+
+
+def test_load_shard_broadcast():
+    devices = [0, 1]
+    param_dict = {
+        "param_0": np.random.uniform(size=[64, 128]).astype("float16"),
+        "param_1": np.random.uniform(size=[32, 128]).astype("float32"),
+    }
+    shard_info = {}
+    with tempfile.TemporaryDirectory() as path:
+
+        def run_test(sess):
+            loader = _create_loader(sess, path, param_dict, shard_info)
+            loader_load = sess.get_global_func("runtime.disco.ShardLoaderLoadAll")
+            params = loader_load(loader)
+            p_0 = params.debug_get_from_remote(0)
+            p_1 = params.debug_get_from_remote(1)
+            np.testing.assert_equal(param_dict["param_0"], p_0[0].numpy())
+            np.testing.assert_equal(param_dict["param_0"], p_1[0].numpy())
+            np.testing.assert_equal(param_dict["param_1"], p_0[1].numpy())
+            np.testing.assert_equal(param_dict["param_1"], p_1[1].numpy())
+
+        _run_with_nccl_session(devices, run_test)
+
+
+def test_load_qkv_proj_shard():  # pylint: disable=too-many-locals
+    devices = [0, 1]
+    num_shards = len(devices)
+    q_heads = 8
+    kv_heads = 10
+    head_dim = 10
+    hidden_size = 20
+    w_q = np.random.uniform(size=[q_heads * head_dim, hidden_size]).astype("float16")
+    w_k = np.random.uniform(size=[kv_heads * head_dim, hidden_size]).astype("float16")
+    w_v = np.random.uniform(size=[kv_heads * head_dim, hidden_size]).astype("float16")
+    w_qkv = np.concatenate([w_q, w_k, w_v], axis=0)
+    param_dict = {"w_qkv": w_qkv}
+    np_qkv = np.concatenate(
+        [
+            w_q.reshape((num_shards, q_heads // num_shards, head_dim, hidden_size)),
+            w_k.reshape((num_shards, kv_heads // num_shards, head_dim, hidden_size)),
+            w_v.reshape((num_shards, kv_heads // num_shards, head_dim, hidden_size)),
+        ],
+        axis=1,
+    ).reshape((num_shards, -1, hidden_size))
+
+    shard_info = {
+        "w_qkv": [
+            [
+                "tests.disco.shard_qkv_0",
+                [
+                    (num_shards, (q_heads + kv_heads * 2) // num_shards, head_dim, hidden_size),
+                    "float16",
+                ],
+                num_shards,
+                q_heads,
+                kv_heads,
+            ],
+            [
+                "tests.disco.shard_qkv_1",
+                [
+                    (num_shards, (q_heads + kv_heads * 2) // num_shards * head_dim, hidden_size),
+                    "float16",
+                ],
+            ],
+        ],
+    }
+
+    with tempfile.TemporaryDirectory() as path:
+
+        def run_test(sess):
+            loader = _create_loader(sess, path, param_dict, shard_info)
+            loader_load = sess.get_global_func("runtime.disco.ShardLoaderLoad")
+            d_0 = loader_load(loader, Shape([0]))
+            np.testing.assert_equal(
+                np_qkv[0],
+                d_0.debug_get_from_remote(0).numpy(),
+            )
+            np.testing.assert_equal(
+                np_qkv[1],
+                d_0.debug_get_from_remote(1).numpy(),
+            )
+
+        _run_with_nccl_session(devices, run_test)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/disco/test_nvshmem.py

@@ -0,0 +1,390 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Basic tests for a Disco nvshmem support"""
+
+# pylint: disable=missing-docstring
+import multiprocessing
+import os
+import shutil
+import socket
+import subprocess
+import sys
+import tempfile
+import threading
+
+import numpy as np
+import pytest
+from tvm_ffi import Shape
+
+import tvm
+import tvm.testing
+from tvm.runtime import disco as di
+from tvm.script import ir as I
+from tvm.script import relax as R
+from tvm.script import tirx as T
+from tvm.testing import env
+
+if di is None:
+    pytest.skip("disco runtime is not available", allow_module_level=True)
+
+pytestmark = [
+    pytest.mark.gpu,
+    pytest.mark.skipif(not env.has_nvshmem(), reason="need nvshmem"),
+]
+
+
+_SOCKET_SESSION_TESTER = None
+
+
+def _get_free_port():
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.bind(("", 0))
+    port = sock.getsockname()[1]
+    sock.close()
+    return port
+
+
+class SocketSessionTester:
+    """Run a disco SocketSession with one local node and remote nodes.
+
+    Each remote node is a `tvm.exec.disco_remote_socket_session` subprocess
+    launched with the current Python interpreter.
+    """
+
+    def __init__(self, num_workers, num_nodes=2, num_groups=1):
+        # Initialize the attributes used by __del__ first, so that teardown is
+        # safe even when __init__ raises below.
+        self.sess = None
+        self.remote_nodes = []
+        assert num_workers % num_nodes == 0
+        num_workers_per_node = num_workers // num_nodes
+        server_host = "localhost"
+        server_port = _get_free_port()
+        server_exc = []
+
+        def start_server():
+            try:
+                self.sess = di.SocketSession(
+                    num_nodes, num_workers_per_node, num_groups, server_host, server_port
+                )
+            except Exception as exc:  # pylint: disable=broad-except
+                server_exc.append(exc)
+
+        thread = threading.Thread(target=start_server)
+        thread.start()
+
+        cmd = "tvm.exec.disco_remote_socket_session"
+        for _i in range(num_nodes - 1):
+            self.remote_nodes.append(
+                subprocess.Popen(
+                    [
+                        sys.executable,
+                        "-m",
+                        cmd,
+                        server_host,
+                        str(server_port),
+                        str(num_workers_per_node),
+                    ],
+                    stdout=sys.stdout,
+                    stderr=sys.stderr,
+                )
+            )
+
+        thread.join()
+        if server_exc:
+            raise server_exc[0]
+
+    # Bound at class creation: module globals may already be cleared when
+    # __del__ runs during interpreter shutdown.
+    _TIMEOUT_EXPIRED = subprocess.TimeoutExpired
+
+    def __del__(self):
+        try:
+            # Shut down the session first so remote nodes can exit gracefully.
+            if self.sess is not None:
+                self.sess.shutdown()
+        finally:
+            for node in self.remote_nodes:
+                try:
+                    node.wait(timeout=10)
+                except self._TIMEOUT_EXPIRED:
+                    node.kill()
+                    node.wait()
+
+
+def create_socket_session(num_workers):
+    """Create a socket session backed by one local and one remote node.
+
+    The tester is kept alive in a module-level global so that the session
+    survives until the next call (or interpreter exit) replaces it.
+    """
+    global _SOCKET_SESSION_TESTER
+    # Rebind (not `del`) so the global stays defined if the constructor raises.
+    _SOCKET_SESSION_TESTER = None
+    _SOCKET_SESSION_TESTER = SocketSessionTester(num_workers)
+    assert _SOCKET_SESSION_TESTER.sess is not None
+    return _SOCKET_SESSION_TESTER.sess
+
+
+_all_session_kinds = [di.ProcessSession, create_socket_session]
+_all_num_workers = [2, 4]
+
+_SUBPROCESS_TIMEOUT_SEC = 600
+
+
+def _run_in_fresh_process(target, *args):
+    """Run a test body in a freshly spawned process.
+
+    After the first call to `nvshmem_init`, a subsequent call to `nvshmem_init`
+    or `nvshmem_init_thread` in the same program results in undefined behavior,
+    and worker-0 of a Disco session lives in the calling process. So each test
+    body must run in its own process. The 'spawn' start method avoids
+    inheriting CUDA state from this process.
+    """
+
+    def run_and_check():
+        proc = multiprocessing.get_context("spawn").Process(target=target, args=args)
+        proc.start()
+        proc.join(timeout=_SUBPROCESS_TIMEOUT_SEC)
+        if proc.is_alive():
+            proc.kill()
+            proc.join()
+            pytest.fail(
+                f"{target.__name__}{args} timed out after {_SUBPROCESS_TIMEOUT_SEC} seconds"
+            )
+        assert proc.exitcode == 0, f"{target.__name__}{args} failed with exit code {proc.exitcode}"
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+def _require_cuda_devices(num_workers):
+    # Each nvshmem worker binds its own CUDA device (cudaSetDevice(worker_id)).
+    if not all(tvm.cuda(i).exist for i in range(num_workers)):
+        pytest.skip(f"Requires {num_workers} CUDA devices")
+
+
+def _init_finalize(session_kind, num_workers):
+    sess = session_kind(num_workers=num_workers)
+    f_init_nvshmem_uid = tvm.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid")
+    uid = f_init_nvshmem_uid()
+    init_dfunc = sess.get_global_func("runtime.disco.nvshmem.init_nvshmem")
+    init_dfunc(uid, num_workers, 0)
+    sess.sync_worker_0()
+    finalize_dfunc = sess.get_global_func("runtime.disco.nvshmem.finalize_nvshmem")
+    finalize_dfunc()
+    sess.sync_worker_0()
+
+
+def _empty(session_kind, num_workers):
+    device = tvm.cuda()
+    sess = session_kind(num_workers=num_workers)
+    f_init_nvshmem_uid = tvm.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid")
+    uid = f_init_nvshmem_uid()
+    init_dfunc = sess.get_global_func("runtime.disco.nvshmem.init_nvshmem")
+    init_dfunc(uid, num_workers, 0)
+    sess.sync_worker_0()
+    empty_dfunc = sess.get_global_func("runtime.disco.nvshmem.empty")
+    _a = empty_dfunc(Shape((32, 64)), "float32", device)
+    _b = empty_dfunc(Shape((64, 32)), "float32", device)
+    sess.sync_worker_0()
+    finalize_dfunc = sess.get_global_func("runtime.disco.nvshmem.finalize_nvshmem")
+    finalize_dfunc()
+    sess.sync_worker_0()
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+@pytest.mark.parametrize("num_workers", _all_num_workers)
+def test_nvshmem_init_finalize(session_kind, num_workers: int):
+    _require_cuda_devices(num_workers)
+    _run_in_fresh_process(_init_finalize, session_kind, num_workers)
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+@pytest.mark.parametrize("num_workers", _all_num_workers)
+def test_nvshmem_empty(session_kind, num_workers: int):
+    _require_cuda_devices(num_workers)
+    _run_in_fresh_process(_empty, session_kind, num_workers)
+
+
+def _compile():
+    num_workers = 2
+    sess = di.ProcessSession(num_workers=num_workers)
+
+    f_init_nvshmem_uid = tvm.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid")
+    uid = f_init_nvshmem_uid()
+    init_dfunc = sess.get_global_func("runtime.disco.nvshmem.init_nvshmem")
+    init_dfunc(uid, num_workers, 0)
+    sess.sync_worker_0()
+
+    @T.prim_func(s_tir=True)
+    def main(A: T.Buffer((8, 16), "float32"), B: T.Buffer((16, 8), "float32")):
+        for i in T.thread_binding(T.int64(8), thread="threadIdx.y"):
+            for j in T.thread_binding(T.int64(16), thread="threadIdx.x"):
+                with T.sblock("T_transpose"):
+                    v0 = T.axis.spatial(T.int64(8), i)
+                    v1 = T.axis.spatial(T.int64(16), j)
+                    T.reads(A[v0, v1])
+                    T.writes(B[v1, v0])
+                    B[v1, v0] = A[v0, v1]
+
+    tmpdir = tempfile.mkdtemp()
+    try:
+        path = tmpdir + "/test.so"
+        A_np = np.arange(8 * 16).astype("float32").reshape([8, 16])
+        B_np = np.zeros((16, 8), dtype="float32")
+        A_array = sess.empty(A_np.shape, "float32")
+        B_array = sess.empty(B_np.shape, "float32")
+        A_array.debug_copy_from(0, A_np)
+
+        target = tvm.target.Target("cuda")
+        tvm.compile(main, target=target).export_library(path)
+        mod = sess.load_vm_module(path)
+        mod["main"](A_array, B_array)
+
+        B_res = B_array.debug_get_from_remote(0).numpy()
+        np.testing.assert_equal(B_res, A_np.T)
+
+        # sync all workers to make sure the temporary files are cleaned up after all workers
+        # finish the execution
+        sess._sync_all()
+
+        finalize_dfunc = sess.get_global_func("runtime.disco.nvshmem.finalize_nvshmem")
+        finalize_dfunc()
+        sess.sync_worker_0()
+    finally:
+        sess.shutdown()
+        shutil.rmtree(tmpdir, ignore_errors=True)
+
+
+def test_nvshmem_compile():
+    _require_cuda_devices(2)
+    _run_in_fresh_process(_compile)
+
+
+NVSHMEM_QUERY_KERNEL_SOURCE = """
+#include <nvshmem.h>
+
+extern "C" __global__ void nvshmem_query_kernel(int* my_pe_out, int* n_pes_out) {
+    my_pe_out[0] = nvshmem_my_pe();
+    n_pes_out[0] = nvshmem_n_pes();
+}
+"""
+
+
+def _kernel_compile(compile_mode):
+    """Compile and run a kernel that calls NVSHMEM functions.
+
+    Runs in a fresh process, so setting the env var is safe.
+    """
+    os.environ["TVM_CUDA_COMPILE_MODE"] = compile_mode
+
+    num_workers = 2
+    sess = di.ProcessSession(num_workers=num_workers)
+
+    f_init_nvshmem_uid = tvm.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid")
+    uid = f_init_nvshmem_uid()
+    init_dfunc = sess.get_global_func("runtime.disco.nvshmem.init_nvshmem")
+    init_dfunc(uid, num_workers, 0)
+    sess.sync_worker_0()
+
+    try:
+
+        @I.ir_module(s_tir=True)
+        class NvshmemQueryModule:
+            @T.prim_func(s_tir=True)
+            def query_pe(
+                my_pe_out: T.Buffer((1,), "int32"),
+                n_pes_out: T.Buffer((1,), "int32"),
+            ):
+                with T.sblock("root"):
+                    T.reads()
+                    T.writes(my_pe_out[0:1], n_pes_out[0:1])
+                    T.call_kernel(
+                        NVSHMEM_QUERY_KERNEL_SOURCE,
+                        ((1,), (1,)),  # grid=(1,), block=(1,)
+                        my_pe_out.data,
+                        n_pes_out.data,
+                        kernel_name="nvshmem_query_kernel",
+                    )
+
+            @R.function
+            def main() -> R.Tuple(R.Tensor((1,), "int32"), R.Tensor((1,), "int32")):
+                cls = NvshmemQueryModule
+                with R.dataflow():
+                    my_pe = R.call_tir(
+                        cls.query_pe,
+                        (),
+                        out_ty=[
+                            R.Tensor((1,), "int32"),
+                            R.Tensor((1,), "int32"),
+                        ],
+                    )
+                    R.output(my_pe)
+                return my_pe
+
+        tmpdir = tempfile.mkdtemp()
+        try:
+            path = tmpdir + "/test_nvshmem_kernel.so"
+
+            target = tvm.target.Target("cuda")
+            tvm.compile(NvshmemQueryModule, target=target).export_library(path)
+            mod = sess.load_vm_module(path)
+            result = mod["main"]()
+
+            # Verify results from each worker
+            for worker_id in range(num_workers):
+                my_pe_result, n_pes_result = result.debug_get_from_remote(worker_id)
+                my_pe_val = my_pe_result.numpy()[0]
+                n_pes_val = n_pes_result.numpy()[0]
+                assert my_pe_val == worker_id, (
+                    f"Worker {worker_id} reported my_pe={my_pe_val}, expected {worker_id}"
+                )
+                assert n_pes_val == num_workers, (
+                    f"Worker {worker_id} reported n_pes={n_pes_val}, expected {num_workers}"
+                )
+
+            # Sync all workers before cleanup
+            sess._sync_all()
+
+            finalize_dfunc = sess.get_global_func("runtime.disco.nvshmem.finalize_nvshmem")
+            finalize_dfunc()
+            sess.sync_worker_0()
+        finally:
+            shutil.rmtree(tmpdir, ignore_errors=True)
+    finally:
+        sess.shutdown()
+
+
+def test_nvshmem_kernel_compile_nvcc():
+    """Test NVSHMEM kernel compilation with nvcc."""
+    _require_cuda_devices(2)
+    _run_in_fresh_process(_kernel_compile, "nvcc")
+
+
+def test_nvshmem_kernel_compile_nvrtc():
+    """Test NVSHMEM kernel compilation with nvrtc."""
+    _require_cuda_devices(2)
+    try:
+        from cuda.bindings import nvrtc  # noqa: F401
+    except ImportError:
+        pytest.skip("cuda-python not available, skipping nvrtc test")
+
+    _run_in_fresh_process(_kernel_compile, "nvrtc")
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/disco/test_session.py

@@ -0,0 +1,347 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Basic tests for a Disco session"""
+
+# pylint: disable=missing-docstring
+import socket
+import subprocess
+import sys
+import tempfile
+import threading
+
+import numpy as np
+import pytest
+from tvm_ffi import Shape
+from tvm_ffi.core import String
+
+import tvm
+import tvm.testing
+
+# Imported for the side effect of registering the tests.disco.* worker functions.
+from tvm.exec import disco_worker as _  # noqa: F401  # pylint: disable=unused-import
+from tvm.runtime import disco as di
+from tvm.script import ir as I
+from tvm.script import relax as R
+from tvm.script import tirx as T
+
+if di is None:
+    pytest.skip("disco runtime is not available", allow_module_level=True)
+
+
+_SOCKET_SESSION_TESTER = None
+
+
+def _get_free_port():
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.bind(("", 0))
+    port = sock.getsockname()[1]
+    sock.close()
+    return port
+
+
+class SocketSessionTester:
+    """Run a disco SocketSession with one local node and remote nodes.
+
+    Each remote node is a `tvm.exec.disco_remote_socket_session` subprocess
+    launched with the current Python interpreter.
+    """
+
+    def __init__(self, num_workers, num_nodes=2, num_groups=1):
+        # Initialize the attributes used by __del__ first, so that teardown is
+        # safe even when __init__ raises below.
+        self.sess = None
+        self.remote_nodes = []
+        assert num_workers % num_nodes == 0
+        num_workers_per_node = num_workers // num_nodes
+        server_host = "localhost"
+        server_port = _get_free_port()
+        server_exc = []
+
+        def start_server():
+            try:
+                self.sess = di.SocketSession(
+                    num_nodes, num_workers_per_node, num_groups, server_host, server_port
+                )
+            except Exception as exc:  # pylint: disable=broad-except
+                server_exc.append(exc)
+
+        thread = threading.Thread(target=start_server)
+        thread.start()
+
+        cmd = "tvm.exec.disco_remote_socket_session"
+        for _i in range(num_nodes - 1):
+            self.remote_nodes.append(
+                subprocess.Popen(
+                    [
+                        sys.executable,
+                        "-m",
+                        cmd,
+                        server_host,
+                        str(server_port),
+                        str(num_workers_per_node),
+                    ],
+                    stdout=sys.stdout,
+                    stderr=sys.stderr,
+                )
+            )
+
+        thread.join()
+        if server_exc:
+            raise server_exc[0]
+
+    # Bound at class creation: module globals may already be cleared when
+    # __del__ runs during interpreter shutdown.
+    _TIMEOUT_EXPIRED = subprocess.TimeoutExpired
+
+    def __del__(self):
+        try:
+            # Shut down the session first so remote nodes can exit gracefully.
+            if self.sess is not None:
+                self.sess.shutdown()
+        finally:
+            for node in self.remote_nodes:
+                try:
+                    node.wait(timeout=10)
+                except self._TIMEOUT_EXPIRED:
+                    node.kill()
+                    node.wait()
+
+
+def create_socket_session(num_workers):
+    """Create a socket session backed by one local and one remote node.
+
+    The tester is kept alive in a module-level global so that the session
+    survives until the next call (or interpreter exit) replaces it.
+    """
+    global _SOCKET_SESSION_TESTER
+    # Rebind (not `del`) so the global stays defined if the constructor raises.
+    _SOCKET_SESSION_TESTER = None
+    _SOCKET_SESSION_TESTER = SocketSessionTester(num_workers)
+    assert _SOCKET_SESSION_TESTER.sess is not None
+    return _SOCKET_SESSION_TESTER.sess
+
+
+def _numpy_to_worker_0(sess: di.Session, np_array: np.array, device):
+    x_array = sess.empty(np_array.shape, "float32", device=device)
+    host_array = tvm.runtime.tensor(np_array, device=device)
+    sess.copy_to_worker_0(host_array, x_array)
+    return x_array
+
+
+def _numpy_from_worker_0(sess: di.Session, remote_array, shape, dtype):
+    host_array = tvm.runtime.empty(shape, dtype, device=tvm.cpu())
+    sess.copy_from_worker_0(host_array, remote_array)
+    sess.sync_worker_0()
+    return host_array.numpy()
+
+
+_all_session_kinds = [di.ThreadedSession, di.ProcessSession, create_socket_session]
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+def test_int(session_kind):  # pylint: disable=invalid-name
+    num_workers = 4
+    sess = session_kind(num_workers=num_workers)
+    func: di.DPackedFunc = sess.get_global_func("tests.disco.add_one")
+    result: di.DRef = func(1)
+    for i in range(num_workers):
+        assert result.debug_get_from_remote(i) == 2
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+def test_float(session_kind):
+    num_workers = 4
+    sess = session_kind(num_workers=num_workers)
+    func: di.DPackedFunc = sess.get_global_func("tests.disco.add_one_float")
+    result: di.DRef = func(1.5)
+
+    for i in range(num_workers):
+        assert result.debug_get_from_remote(i) == 2.0
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+def test_tensor(session_kind):
+    num_workers = 4
+    sess = session_kind(num_workers=num_workers)
+    device = tvm.cpu(0)
+    x_np = np.arange(6).astype("float32").reshape([2, 3])
+    y_np = np.arange(6).astype("float32").reshape([2, 3]) + 1
+    x_disc = _numpy_to_worker_0(sess, x_np, device=device)
+    y_disc = sess.get_global_func("tests.disco.add_one_tensor")(x_disc)
+    y_nd = _numpy_from_worker_0(sess, y_disc, shape=y_np.shape, dtype=y_np.dtype)
+    np.testing.assert_equal(y_nd, y_np)
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+def test_string(session_kind):
+    num_workers = 4
+    sess = session_kind(num_workers=num_workers)
+    func: di.DPackedFunc = sess.get_global_func("tests.disco.str")
+    result: di.DRef = func("hello")
+
+    for i in range(num_workers):
+        assert result.debug_get_from_remote(i) == "hello_suffix"
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+def test_string_obj(session_kind):
+    num_workers = 4
+    sess = session_kind(num_workers=num_workers)
+    func: di.DPackedFunc = sess.get_global_func("tests.disco.str_obj")
+    result: di.DRef = func(String("hello"))
+
+    for i in range(num_workers):
+        value = result.debug_get_from_remote(i)
+        assert isinstance(value, str)
+        assert value == "hello_suffix"
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+def test_shape_tuple(session_kind):
+    num_workers = 4
+    sess = session_kind(num_workers=num_workers)
+    func: di.DPackedFunc = sess.get_global_func("tests.disco.shape_tuple")
+    result: di.DRef = func(Shape([1, 2, 3]))
+    for i in range(num_workers):
+        value = result.debug_get_from_remote(i)
+        assert isinstance(value, Shape)
+        assert list(value) == [1, 2, 3, 4, 5]
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+def test_vm_module(session_kind):
+    num_workers = 4
+    sess = session_kind(num_workers=num_workers)
+
+    # pylint: disable=invalid-name
+    @I.ir_module(s_tir=True)
+    class TestMod:
+        @T.prim_func(s_tir=True)
+        def transpose(A: T.Buffer((8, 16), "float32"), B: T.Buffer((16, 8), "float32")):
+            for i, j in T.grid(16, 8):
+                with T.sblock("transpose"):
+                    vi, vj = T.axis.remap("SS", [i, j])
+                    B[vi, vj] = A[vj, vi]
+
+        @R.function
+        def main(A: R.Tensor((8, 16), dtype="float32")) -> R.Tensor((16, 8), dtype="float32"):
+            cls = TestMod
+            with R.dataflow():
+                B = R.call_tir(cls.transpose, (A,), out_ty=R.Tensor((16, 8), dtype="float32"))
+                R.output(B)
+            return B
+
+    # pylint: enable=invalid-name
+    with tempfile.TemporaryDirectory() as tmpdir:
+        path = tmpdir + "/test.so"
+        device = tvm.cpu()
+        x_np = np.arange(8 * 16).astype("float32").reshape([8, 16])
+        y_np = x_np.transpose()
+
+        tvm.compile(TestMod, target="llvm").export_library(path)
+        mod = sess.load_vm_module(path, device=device)
+
+        x_disc = _numpy_to_worker_0(sess, x_np, device=device)
+        y_disc = mod["main"](x_disc)
+        y_nd = _numpy_from_worker_0(sess, y_disc, shape=y_np.shape, dtype=y_np.dtype)
+        np.testing.assert_equal(y_nd, y_np)
+
+        # sync all workers to make sure the temporary files are cleaned up after all workers
+        # finish the execution
+        for i in range(num_workers):
+            sess._sync_worker(i)
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+def test_vm_multi_func(session_kind):
+    num_workers = 4
+    sess = session_kind(num_workers=num_workers)
+
+    # pylint: disable=invalid-name
+    @I.ir_module(s_tir=True)
+    class TestMod:
+        @T.prim_func(s_tir=True)
+        def t1(A: T.Buffer((8, 16), "float32"), B: T.Buffer((16, 8), "float32")):
+            for i, j in T.grid(16, 8):
+                with T.sblock("t1"):
+                    vi, vj = T.axis.remap("SS", [i, j])
+                    B[vi, vj] = A[vj, vi]
+
+        @T.prim_func(s_tir=True)
+        def t2(A: T.Buffer((16, 8), "float32"), B: T.Buffer((8, 16), "float32")):
+            for i, j in T.grid(8, 16):
+                with T.sblock("t2"):
+                    vi, vj = T.axis.remap("SS", [i, j])
+                    B[vi, vj] = A[vj, vi]
+
+        @R.function
+        def transpose_1(A: R.Tensor((8, 16), dtype="float32")) -> R.Tensor(
+            (16, 8), dtype="float32"
+        ):
+            R.func_attr({"global_symbol": "transpose_1"})
+            cls = TestMod
+            with R.dataflow():
+                B = R.call_tir(cls.t1, (A,), out_ty=R.Tensor((16, 8), dtype="float32"))
+                R.output(B)
+            return B
+
+        @R.function
+        def transpose_2(A: R.Tensor((16, 8), dtype="float32")) -> R.Tensor(
+            (8, 16), dtype="float32"
+        ):
+            R.func_attr({"global_symbol": "transpose_2"})
+            cls = TestMod
+            with R.dataflow():
+                B = R.call_tir(cls.t2, (A,), out_ty=R.Tensor((8, 16), dtype="float32"))
+                R.output(B)
+            return B
+
+    # pylint: enable=invalid-name
+    with tempfile.TemporaryDirectory() as tmpdir:
+        path = tmpdir + "/test.so"
+        device = tvm.cpu()
+        x_np = np.arange(8 * 16).astype("float32").reshape([8, 16])
+        y_np = x_np.transpose()
+
+        tvm.compile(TestMod, target="llvm").export_library(path)
+        mod = sess.load_vm_module(path, device=device)
+
+        x_disc = _numpy_to_worker_0(sess, x_np, device=device)
+        y_disc = mod["transpose_1"](x_disc)
+        z_disc = mod["transpose_2"](y_disc)
+        y_nd = _numpy_from_worker_0(sess, y_disc, shape=y_np.shape, dtype=y_np.dtype)
+        z_nd = _numpy_from_worker_0(sess, z_disc, shape=x_np.shape, dtype=x_np.dtype)
+        np.testing.assert_equal(y_nd, y_np)
+        np.testing.assert_equal(z_nd, x_np)
+
+        # sync all workers to make sure the temporary files are cleaned up after all workers
+        # finish the execution
+        for i in range(num_workers):
+            sess._sync_worker(i)
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+@pytest.mark.parametrize("num_workers", [1, 2, 4])
+def test_num_workers(session_kind, num_workers):
+    if session_kind == create_socket_session and num_workers < 2:
+        return
+    sess = session_kind(num_workers=num_workers)
+    assert sess.num_workers == num_workers
+
+
+if __name__ == "__main__":
+    tvm.testing.main()