chore: import upstream snapshot with attribution

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()