        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
        for t in [inp_tv, weight_tv]:
            t.set_device_mesh(mesh)
            t.outer_split(-1, d)
            t.axis(-2).parallelize(nvfuser.ParallelType.mesh_x)

    b, s = 2, 1024
    inp_ref = torch.randn(b, s, d * e)
    weight_ref = torch.randn(e, d * e)
    out_ref = torch.nn.functional.linear(inp_ref, weight_ref)

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    weight = multidevice_test.shard_tensor(weight_ref, weight_tv)
    (out,) = fd.execute([inp, weight])

    # rtol is the same as the default for fp32. atol is slightly increased.
    torch.testing.assert_close(out.cpu(), out_ref, rtol=1.3e-6, atol=1e-3)

@pytest.mark.mpi
def test_row_parallel_linear_with_bias(multidevice_test):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
    e = 5

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor([-1, -1, d * e])
        weight_tv = fd.define_tensor([e, d * e])
        bias_tv = fd.define_tensor([e])
        out = fd.ops.linear(inp_tv, weight_tv, bias_tv)
        fd.add_output(out)

        for t in [inp_tv, weight_tv]:
            t.set_device_mesh(mesh)
            t.outer_split(-1, d)
            t.axis(-2).parallelize(nvfuser.ParallelType.mesh_x)
        bias_tv.set_device_mesh(mesh)

    b, s = 2, 3
    inp_ref = torch.randn(b, s, d * e)
    weight_ref = torch.randn(e, d * e)
    bias_ref = torch.randn(e)
    out_ref = torch.nn.functional.linear(inp_ref, weight_ref, bias_ref)

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    weight = multidevice_test.shard_tensor(weight_ref, weight_tv)
    bias = multidevice_test.shard_tensor(bias_ref, bias_tv)
    (out,) = fd.execute([inp, weight, bias])

    # rtol is the same as the default for fp32. atol is slightly increased.
    torch.testing.assert_close(out.cpu(), out_ref, rtol=1.3e-6, atol=1e-3)

@pytest.mark.mpi
def test_linear_reduce_scatter(multidevice_test):
    d = multidevice_test.size
    b, s, e = 3, 5, 7

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor([-1, d * s, d * e], dtype=DataType.BFloat16)
        weight_tv = fd.define_tensor([-1, d * e], dtype=DataType.BFloat16)
        bias_tv = fd.define_tensor([e], dtype=DataType.BFloat16)
        out_tv = fd.ops.linear(inp_tv, weight_tv, bias_tv)
        fd.add_output(out_tv)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
        bias_tv.set_device_mesh(mesh)
        for tv in [inp_tv, weight_tv, out_tv]:
            tv.set_device_mesh(mesh)
            tv.split(-1, d, inner_split=False)
            tv.axis(-2).parallelize(nvfuser.ParallelType.mesh_x)


            t.outer_split(-1, d)
            t.axis(-2).parallelize(nvfuser.ParallelType.mesh_x)
        bias_tv.set_device_mesh(mesh)

    b, s = 2, 3
    inp_ref = torch.randn(b, s, d * e)
    weight_ref = torch.randn(e, d * e)
    bias_ref = torch.randn(e)
    out_ref = torch.nn.functional.linear(inp_ref, weight_ref, bias_ref)

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    weight = multidevice_test.shard_tensor(weight_ref, weight_tv)
    bias = multidevice_test.shard_tensor(bias_ref, bias_tv)
    (out,) = fd.execute([inp, weight, bias])

    # rtol is the same as the default for fp32. atol is slightly increased.
    torch.testing.assert_close(out.cpu(), out_ref, rtol=1.3e-6, atol=1e-3)

@pytest.mark.mpi
def test_linear_reduce_scatter(multidevice_test):
    d = multidevice_test.size
    b, s, e = 3, 5, 7

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor([-1, d * s, d * e], dtype=DataType.BFloat16)
        weight_tv = fd.define_tensor([-1, d * e], dtype=DataType.BFloat16)
        bias_tv = fd.define_tensor([e], dtype=DataType.BFloat16)
        out_tv = fd.ops.linear(inp_tv, weight_tv, bias_tv)
        fd.add_output(out_tv)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
        bias_tv.set_device_mesh(mesh)
        for tv in [inp_tv, weight_tv, out_tv]:
            tv.set_device_mesh(mesh)
            tv.split(-1, d, inner_split=False)
            tv.axis(-2).parallelize(nvfuser.ParallelType.mesh_x)

        out_tv.outer_split(1, d)
        out_tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

    inp_ref = torch.randint(-2, 3, (b, d * s, d * e)).to(torch.bfloat16)
    weight_ref = torch.randint(-2, 3, (e, d * e)).to(torch.bfloat16)
    bias_ref = torch.randint(-2, 3, (e,)).to(torch.bfloat16)
    out_ref = torch.nn.functional.linear(inp_ref, weight_ref, bias_ref)

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    weight = multidevice_test.shard_tensor(weight_ref, weight_tv)
    bias = multidevice_test.shard_tensor(bias_ref, bias_tv)

    with PythonProfiler() as prof:
        (out,) = fd.execute([inp, weight, bias])

    # Only one reduce scatter kernel should be scheduled.
    assert len(
        [kp for kp in prof.profile.kernel_profiles if kp.scheduler == "communication"]
    ) == (1 if d > 1 else 0)

    torch.testing.assert_close(
        out,
        multidevice_test.shard_tensor(out_ref, out_tv),
    )

@pytest.mark.mpi
def test_column_parallel_matmul(multidevice_test):
    d = multidevice_test.size
    e = 768

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor([-1, -1, e])  # [b, s, e]
        weight_tv = fd.define_tensor([e, d * e])
        out = fd.ops.matmul(inp_tv, weight_tv)  # [b, s, d*e]
        fd.add_output(out)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
        for t in [inp_tv, weight_tv, out]:
            t.set_device_mesh(mesh)

        # Shard N for weight (K, N)
        weight_tv.outer_split(-1, d)
        weight_tv.axis(-2).parallelize(nvfuser.ParallelType.mesh_x)


    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    weight = multidevice_test.shard_tensor(weight_ref, weight_tv)
    bias = multidevice_test.shard_tensor(bias_ref, bias_tv)

    with PythonProfiler() as prof:
        (out,) = fd.execute([inp, weight, bias])

    # Only one reduce scatter kernel should be scheduled.
    assert len(
        [kp for kp in prof.profile.kernel_profiles if kp.scheduler == "communication"]
    ) == (1 if d > 1 else 0)

    torch.testing.assert_close(
        out,
        multidevice_test.shard_tensor(out_ref, out_tv),
    )

@pytest.mark.mpi
def test_column_parallel_matmul(multidevice_test):
    d = multidevice_test.size
    e = 768

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor([-1, -1, e])  # [b, s, e]
        weight_tv = fd.define_tensor([e, d * e])
        out = fd.ops.matmul(inp_tv, weight_tv)  # [b, s, d*e]
        fd.add_output(out)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
        for t in [inp_tv, weight_tv, out]:
            t.set_device_mesh(mesh)

        # Shard N for weight (K, N)
        weight_tv.outer_split(-1, d)
        weight_tv.axis(-2).parallelize(nvfuser.ParallelType.mesh_x)

        # Output of linear: {.., i{M}, i{N}, r{K}}
        # Shard N -> axis(-2)
        out.outer_split(-2, d)
        out.axis(-3).parallelize(nvfuser.ParallelType.mesh_x)

    b, s = 2, 1024
    inp_ref = torch.randn(b, s, e)
    weight_ref = torch.randn(e, d * e)
    out_ref = torch.matmul(inp_ref, weight_ref)

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    weight = multidevice_test.shard_tensor(weight_ref, weight_tv)
    (out,) = fd.execute([inp, weight])

    # rtol is the same as the default for fp32. atol is slightly increased.
    torch.testing.assert_close(
        out, multidevice_test.shard_tensor_1d(out_ref, 2, mesh), rtol=1.3e-6, atol=1e-3
    )

@pytest.mark.mpi
def test_row_parallel_matmul(multidevice_test):
    d = multidevice_test.size
    e = 8

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor([-1, d * e], contiguity=True, dtype=DataType.Half)
        weight_tv = fd.define_tensor([d * e, e], contiguity=True, dtype=DataType.Half)
        out = fd.ops.matmul(inp_tv, weight_tv)
        fd.add_output(out)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
        for t in [inp_tv, weight_tv, out]:
            t.set_device_mesh(mesh)

        # Shard K for inp (M, K)
        inp_tv.outer_split(-1, d)
        inp_tv.axis(-2).parallelize(nvfuser.ParallelType.mesh_x)


        # Shard N -> axis(-2)
        out.outer_split(-2, d)
        out.axis(-3).parallelize(nvfuser.ParallelType.mesh_x)

    b, s = 2, 1024
    inp_ref = torch.randn(b, s, e)
    weight_ref = torch.randn(e, d * e)
    out_ref = torch.matmul(inp_ref, weight_ref)

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    weight = multidevice_test.shard_tensor(weight_ref, weight_tv)
    (out,) = fd.execute([inp, weight])

    # rtol is the same as the default for fp32. atol is slightly increased.
    torch.testing.assert_close(
        out, multidevice_test.shard_tensor_1d(out_ref, 2, mesh), rtol=1.3e-6, atol=1e-3
    )

@pytest.mark.mpi
def test_row_parallel_matmul(multidevice_test):
    d = multidevice_test.size
    e = 8

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor([-1, d * e], contiguity=True, dtype=DataType.Half)
        weight_tv = fd.define_tensor([d * e, e], contiguity=True, dtype=DataType.Half)
        out = fd.ops.matmul(inp_tv, weight_tv)
        fd.add_output(out)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
        for t in [inp_tv, weight_tv, out]:
            t.set_device_mesh(mesh)

        # Shard K for inp (M, K)
        inp_tv.outer_split(-1, d)
        inp_tv.axis(-2).parallelize(nvfuser.ParallelType.mesh_x)

        # Shard K for weight (K, N)
        weight_tv.outer_split(0, d)
        weight_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

        # I don't think the manual rFactor is necessary any more.
        # DecomposeReshardingsPass should handle it.
        #
        # [i{M}, i{N}, r{K}]
        out.outer_split(-1, d)
        # [i{M}, i{N}, r{d}, r{K//d}]
        local_out = out.rfactor(axes=[-1])
        # local_out = [i{M}, i{N}, i{d}, r{K//d}]
        # out = [i{M}, i{N}, r{d}]
        local_out.set_device_mesh(mesh)
        local_out.axis(-2).parallelize(nvfuser.ParallelType.mesh_x)

    b, s = 1, 4
    inp_ref = torch.randn(b * s, d * e, dtype=torch.half)
    weight_ref = torch.randn(d * e, e, dtype=torch.half)
    out_ref = torch.matmul(inp_ref, weight_ref)

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    weight = multidevice_test.shard_tensor(weight_ref, weight_tv)
    (out,) = fd.execute([inp, weight])

    torch.testing.assert_close(out.cpu(), out_ref, rtol=1e-3, atol=1e-2)

@pytest.mark.mpi
def test_column_parallel_grouped_mm(multidevice_test):
    d = multidevice_test.size
    g = 4
    k = 16
    n = 64 * d

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor([-1, k], dtype=DataType.BFloat16, contiguity=True)
        w_tv = fd.define_tensor([g, n, k], dtype=DataType.BFloat16, contiguity=True)
        w_t = fd.ops.permute(w_tv, [0, 2, 1])
        offsets_tv = fd.define_tensor([g], dtype=DataType.Int32, contiguity=True)
        out = fd.ops.grouped_mm(inp_tv, w_t, offsets_tv)
        fd.add_output(out)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
        for t in [inp_tv, w_tv, offsets_tv]:
            t.set_device_mesh(mesh)


        # [i{M}, i{N}, r{d}, r{K//d}]
        local_out = out.rfactor(axes=[-1])
        # local_out = [i{M}, i{N}, i{d}, r{K//d}]
        # out = [i{M}, i{N}, r{d}]
        local_out.set_device_mesh(mesh)
        local_out.axis(-2).parallelize(nvfuser.ParallelType.mesh_x)

    b, s = 1, 4
    inp_ref = torch.randn(b * s, d * e, dtype=torch.half)
    weight_ref = torch.randn(d * e, e, dtype=torch.half)
    out_ref = torch.matmul(inp_ref, weight_ref)

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    weight = multidevice_test.shard_tensor(weight_ref, weight_tv)
    (out,) = fd.execute([inp, weight])

    torch.testing.assert_close(out.cpu(), out_ref, rtol=1e-3, atol=1e-2)

@pytest.mark.mpi
def test_column_parallel_grouped_mm(multidevice_test):
    d = multidevice_test.size
    g = 4
    k = 16
    n = 64 * d

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor([-1, k], dtype=DataType.BFloat16, contiguity=True)
        w_tv = fd.define_tensor([g, n, k], dtype=DataType.BFloat16, contiguity=True)
        w_t = fd.ops.permute(w_tv, [0, 2, 1])
        offsets_tv = fd.define_tensor([g], dtype=DataType.Int32, contiguity=True)
        out = fd.ops.grouped_mm(inp_tv, w_t, offsets_tv)
        fd.add_output(out)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
        for t in [inp_tv, w_tv, offsets_tv]:
            t.set_device_mesh(mesh)

        w_tv.outer_split(1, d)
        w_tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

    m = 32
    inp_ref = torch.randn(m, k, dtype=torch.bfloat16)
    w_ref = torch.randn(g, n, k, dtype=torch.bfloat16)
    group_sizes = [5, 7, 9, 11]
    assert sum(group_sizes) == m
    group_outs = [
        group_in @ group_w.T
        for group_in, group_w in zip(inp_ref.split(group_sizes), w_ref.unbind())
    ]
    out_ref = torch.cat(group_outs, dim=0)

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    w = multidevice_test.shard_tensor(w_ref, w_tv)
    offsets = torch.cumsum(torch.tensor(group_sizes), 0, dtype=torch.int32).cuda()
    (out,) = fd.execute([inp, w, offsets])

    torch.testing.assert_close(out, multidevice_test.shard_tensor_1d(out_ref, 1, mesh))

@pytest.mark.mpi
def test_row_parallel_grouped_mm(multidevice_test):
    d = multidevice_test.size
    g = 4
    k = 16 * d
    n = 64

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor([-1, k], dtype=DataType.BFloat16, contiguity=True)
        w_tv = fd.define_tensor([g, n, k], dtype=DataType.BFloat16, contiguity=True)
        w_t = fd.ops.permute(w_tv, [0, 2, 1])
        offsets_tv = fd.define_tensor([g], dtype=DataType.Int32, contiguity=True)
        out_tv = fd.ops.grouped_mm(inp_tv, w_t, offsets_tv)
        fd.add_output(out_tv)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
        for t in [inp_tv, w_tv, offsets_tv]:
            t.set_device_mesh(mesh)


    m = 32
    inp_ref = torch.randn(m, k, dtype=torch.bfloat16)
    w_ref = torch.randn(g, n, k, dtype=torch.bfloat16)
    group_sizes = [5, 7, 9, 11]
    assert sum(group_sizes) == m
    group_outs = [
        group_in @ group_w.T
        for group_in, group_w in zip(inp_ref.split(group_sizes), w_ref.unbind())
    ]
    out_ref = torch.cat(group_outs, dim=0)

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    w = multidevice_test.shard_tensor(w_ref, w_tv)
    offsets = torch.cumsum(torch.tensor(group_sizes), 0, dtype=torch.int32).cuda()
    (out,) = fd.execute([inp, w, offsets])

    torch.testing.assert_close(out, multidevice_test.shard_tensor_1d(out_ref, 1, mesh))

@pytest.mark.mpi
def test_row_parallel_grouped_mm(multidevice_test):
    d = multidevice_test.size
    g = 4
    k = 16 * d
    n = 64

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor([-1, k], dtype=DataType.BFloat16, contiguity=True)
        w_tv = fd.define_tensor([g, n, k], dtype=DataType.BFloat16, contiguity=True)
        w_t = fd.ops.permute(w_tv, [0, 2, 1])
        offsets_tv = fd.define_tensor([g], dtype=DataType.Int32, contiguity=True)
        out_tv = fd.ops.grouped_mm(inp_tv, w_t, offsets_tv)
        fd.add_output(out_tv)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
        for t in [inp_tv, w_tv, offsets_tv]:
            t.set_device_mesh(mesh)

        inp_tv.outer_split(-1, d)
        inp_tv.axis(-2).parallelize(nvfuser.ParallelType.mesh_x)

        w_tv.outer_split(-1, d)
        w_tv.axis(-2).parallelize(nvfuser.ParallelType.mesh_x)

    m = 32
    inp_ref = torch.randint(-2, 3, (m, k), dtype=torch.bfloat16)
    w_ref = torch.randint(-2, 3, (g, n, k), dtype=torch.bfloat16)
    group_sizes = [5, 7, 9, 11]
    assert sum(group_sizes) == m
    group_outs = [
        group_in @ group_w.T
        for group_in, group_w in zip(inp_ref.split(group_sizes), w_ref.unbind())
    ]
    out_ref = torch.cat(group_outs, dim=0)

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    w = multidevice_test.shard_tensor(w_ref, w_tv)
    offsets = torch.cumsum(torch.tensor(group_sizes), 0, dtype=torch.int32).cuda()
    (out,) = fd.execute([inp, w, offsets])

    torch.testing.assert_close(out.cpu(), out_ref)

@pytest.mark.mpi
def test_issue4729(multidevice_test):
    d = multidevice_test.size

    with FusionDefinition() as fd:
        x_tv = fd.define_tensor([1, 1, d * 3], dtype=DataType.BFloat16, contiguity=True)
        y_tv = fd.define_tensor([1, 1, d * 3], dtype=DataType.BFloat16, contiguity=True)
        w_tv = fd.define_tensor([-1, d * 3], dtype=DataType.BFloat16, contiguity=True)
        x = fd.ops.cast(x_tv, DataType.Float)
        y = fd.ops.cast(y_tv, DataType.Float)
        xy = fd.ops.mul(x, y)
        xy = fd.ops.cast(xy, DataType.BFloat16)
        z = fd.ops.linear(xy, w_tv)
        fd.add_output(z)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
        for t in [x_tv, y_tv, w_tv]:
            t.set_device_mesh(mesh)


    m = 32
    inp_ref = torch.randint(-2, 3, (m, k), dtype=torch.bfloat16)
    w_ref = torch.randint(-2, 3, (g, n, k), dtype=torch.bfloat16)
    group_sizes = [5, 7, 9, 11]
    assert sum(group_sizes) == m
    group_outs = [
        group_in @ group_w.T
        for group_in, group_w in zip(inp_ref.split(group_sizes), w_ref.unbind())
    ]
    out_ref = torch.cat(group_outs, dim=0)

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    w = multidevice_test.shard_tensor(w_ref, w_tv)
    offsets = torch.cumsum(torch.tensor(group_sizes), 0, dtype=torch.int32).cuda()
    (out,) = fd.execute([inp, w, offsets])

    torch.testing.assert_close(out.cpu(), out_ref)

@pytest.mark.mpi
def test_issue4729(multidevice_test):
    d = multidevice_test.size

    with FusionDefinition() as fd:
        x_tv = fd.define_tensor([1, 1, d * 3], dtype=DataType.BFloat16, contiguity=True)
        y_tv = fd.define_tensor([1, 1, d * 3], dtype=DataType.BFloat16, contiguity=True)
        w_tv = fd.define_tensor([-1, d * 3], dtype=DataType.BFloat16, contiguity=True)
        x = fd.ops.cast(x_tv, DataType.Float)
        y = fd.ops.cast(y_tv, DataType.Float)
        xy = fd.ops.mul(x, y)
        xy = fd.ops.cast(xy, DataType.BFloat16)
        z = fd.ops.linear(xy, w_tv)
        fd.add_output(z)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
        for t in [x_tv, y_tv, w_tv]:
            t.set_device_mesh(mesh)
            t.outer_split(-1, d)
            t.axis(-2).parallelize(nvfuser.ParallelType.mesh_x)

    x_ref = torch.randint(-2, 3, (1, 1, d * 3), dtype=torch.bfloat16)
    y_ref = torch.randint(-2, 3, (1, 1, d * 3), dtype=torch.bfloat16)
    w_ref = torch.randint(-2, 3, (2, d * 3), dtype=torch.bfloat16)
    z_ref = torch.nn.functional.linear(x_ref * y_ref, w_ref)

    x = multidevice_test.shard_tensor(x_ref, x_tv)
    y = multidevice_test.shard_tensor(y_ref, y_tv)
    w = multidevice_test.shard_tensor(w_ref, w_tv)
    (z,) = fd.execute([x, y, w])

    torch.testing.assert_close(z.cpu(), z_ref)

@pytest.mark.mpi
def test_sequence_parallel_linear(multidevice_test):
    d = multidevice_test.size
    b, s, e = 2, 1024, 768
    assert (
        s % d == 0
    ), f"Sequence length {s} must be divisible by the number of devices {d}"
    assert e % d == 0, f"Hidden size {e} must be divisible by the number of devices {d}"

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor(shape=[-1, -1, -1], contiguity=True)  # [b, s, e]
        weight_tv = fd.define_tensor(shape=[-1, -1], contiguity=True)  # [e, e]
        bias_tv = fd.define_tensor(shape=[-1], contiguity=True)  # [e]
        out_tv = fd.ops.linear(inp_tv, weight_tv, bias_tv)  # [b, s, e]
        fd.add_output(out_tv)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
        for t in [inp_tv, weight_tv, bias_tv, out_tv]:


        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
        for t in [x_tv, y_tv, w_tv]:
            t.set_device_mesh(mesh)
            t.outer_split(-1, d)
            t.axis(-2).parallelize(nvfuser.ParallelType.mesh_x)

    x_ref = torch.randint(-2, 3, (1, 1, d * 3), dtype=torch.bfloat16)
    y_ref = torch.randint(-2, 3, (1, 1, d * 3), dtype=torch.bfloat16)
    w_ref = torch.randint(-2, 3, (2, d * 3), dtype=torch.bfloat16)
    z_ref = torch.nn.functional.linear(x_ref * y_ref, w_ref)

    x = multidevice_test.shard_tensor(x_ref, x_tv)
    y = multidevice_test.shard_tensor(y_ref, y_tv)
    w = multidevice_test.shard_tensor(w_ref, w_tv)
    (z,) = fd.execute([x, y, w])

    torch.testing.assert_close(z.cpu(), z_ref)

@pytest.mark.mpi
def test_sequence_parallel_linear(multidevice_test):
    d = multidevice_test.size
    b, s, e = 2, 1024, 768
    assert (
        s % d == 0
    ), f"Sequence length {s} must be divisible by the number of devices {d}"
    assert e % d == 0, f"Hidden size {e} must be divisible by the number of devices {d}"

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor(shape=[-1, -1, -1], contiguity=True)  # [b, s, e]
        weight_tv = fd.define_tensor(shape=[-1, -1], contiguity=True)  # [e, e]
        bias_tv = fd.define_tensor(shape=[-1], contiguity=True)  # [e]
        out_tv = fd.ops.linear(inp_tv, weight_tv, bias_tv)  # [b, s, e]
        fd.add_output(out_tv)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
        for t in [inp_tv, weight_tv, bias_tv, out_tv]:
            t.set_device_mesh(mesh)

        inp_tv.outer_split(1, d)
        inp_tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)
        weight_tv.outer_split(0, d)
        weight_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)
        bias_tv.outer_split(0, d)
        bias_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)
        out_tv.outer_split(1, d)
        out_tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

    inp_ref = torch.randn(b, s, e)
    weight_ref = torch.randn(e, e)
    bias_ref = torch.randn(e)
    out_ref = torch.nn.functional.linear(inp_ref, weight_ref, bias_ref)

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    weight = multidevice_test.shard_tensor(weight_ref, weight_tv)
    bias = multidevice_test.shard_tensor(bias_ref, bias_tv)
    (out,) = fd.execute([inp, weight, bias])
    torch.testing.assert_close(
        out, multidevice_test.shard_tensor(out_ref, out_tv), rtol=1e-3, atol=1e-2
    )

@pytest.mark.mpi
def test_data_and_tensor_parallel_mlp(multidevice_test):
    # Build a GPT-3 style MLP with data and tensor parallelism.
    d = multidevice_test.size
    tp_size = 2

    if d % tp_size != 0:
        pytest.skip(f"Number of devices ({d}) must be divisible by tp_size ({tp_size})")

    dp_size = d // tp_size

    e = 3
    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor([-1, -1, e], contiguity=True)
        up_w_tv = fd.define_tensor([4 * e, e], contiguity=True)
        down_w_tv = fd.define_tensor([e, 4 * e], contiguity=True)
        up_out = fd.ops.linear(inp_tv, up_w_tv)
        down_out = fd.ops.linear(up_out, down_w_tv)


        bias_tv.outer_split(0, d)
        bias_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)
        out_tv.outer_split(1, d)
        out_tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

    inp_ref = torch.randn(b, s, e)
    weight_ref = torch.randn(e, e)
    bias_ref = torch.randn(e)
    out_ref = torch.nn.functional.linear(inp_ref, weight_ref, bias_ref)

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    weight = multidevice_test.shard_tensor(weight_ref, weight_tv)
    bias = multidevice_test.shard_tensor(bias_ref, bias_tv)
    (out,) = fd.execute([inp, weight, bias])
    torch.testing.assert_close(
        out, multidevice_test.shard_tensor(out_ref, out_tv), rtol=1e-3, atol=1e-2
    )

@pytest.mark.mpi
def test_data_and_tensor_parallel_mlp(multidevice_test):
    # Build a GPT-3 style MLP with data and tensor parallelism.
    d = multidevice_test.size
    tp_size = 2

    if d % tp_size != 0:
        pytest.skip(f"Number of devices ({d}) must be divisible by tp_size ({tp_size})")

    dp_size = d // tp_size

    e = 3
    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor([-1, -1, e], contiguity=True)
        up_w_tv = fd.define_tensor([4 * e, e], contiguity=True)
        down_w_tv = fd.define_tensor([e, 4 * e], contiguity=True)
        up_out = fd.ops.linear(inp_tv, up_w_tv)
        down_out = fd.ops.linear(up_out, down_w_tv)
        fd.add_output(down_out)

        # mesh_y (dim 0) for data parallelism, mesh_x (dim 1) for tensor parallelism
        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d).reshape(dp_size, tp_size))
        for t in [inp_tv, up_w_tv, down_w_tv]:
            t.set_device_mesh(mesh)

        inp_tv.outer_split(0, dp_size)
        inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_y)
        up_w_tv.outer_split(0, tp_size)
        up_w_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)
        down_w_tv.outer_split(-1, tp_size)
        down_w_tv.axis(-2).parallelize(nvfuser.ParallelType.mesh_x)

    batch_per_rank = 7
    b, s = dp_size * batch_per_rank, 5

    inp_ref = torch.testing.make_tensor(b, s, e, dtype=torch.int32, device="cpu").to(
        torch.float
    )
    up_w_ref = torch.testing.make_tensor(4 * e, e, dtype=torch.int32, device="cpu").to(
        torch.float
    )
    down_w_ref = torch.testing.make_tensor(
        e, 4 * e, dtype=torch.int32, device="cpu"
    ).to(torch.float)
    up_out_ref = torch.nn.functional.linear(inp_ref, up_w_ref)
    down_out_ref = torch.nn.functional.linear(up_out_ref, down_w_ref)

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    up_w = multidevice_test.shard_tensor(up_w_ref, up_w_tv)
    down_w = multidevice_test.shard_tensor(down_w_ref, down_w_tv)
    (out,) = fd.execute([inp, up_w, down_w])

    dp_rank = multidevice_test.rank // tp_size
    torch.testing.assert_close(
        out.cpu(),
        down_out_ref[dp_rank * batch_per_rank : (dp_rank + 1) * batch_per_rank],
    )


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import torch

def get_benchmark_fn(func, /, profile: bool):
    def wrapper(*args, **kwargs):
        if profile:
            torch.cuda.cudart().cudaProfilerStart()
        result = func(*args, **kwargs)
        torch.cuda.synchronize()
        if profile:
            torch.cuda.cudart().cudaProfilerStop()
        return result

    return wrapper

# Returns two functors, the first with profiler off and the second with profiler
# on. The first functor is usually used for warmup and the second for actual
# benchmarking. This way, one can collect stats of non-warmup
# benchmark iterations using `nsys profile --capture-range=cudaProfilerApi`.
#
# https://docs.nvidia.com/nsight-systems/UserGuide/index.html#handling-application-launchers-mpirun-deepspeed-etc
# has described several ways to profile multi-process applications.
#
# For single-node profiling, I recommend putting `nsys profile` before
# `mpirun`, e.g., `nsys profile ... mpirun -np 8 ...` instead of `mpirun -np 8
# nsys profile ...` or `mpirun -np 1 nsys profile ... : -np 7 ...`. This config
# tries to collect and align traces on different GPUs so it gives the most
# complete picture.  See
# https://github.com/NVIDIA/Fuser/pull/5751/files#r2663586669 for my
# experiment.
def get_benchmark_fns(func):
    return get_benchmark_fn(func, profile=False), get_benchmark_fn(func, profile=True)


        return result

    return wrapper

# Returns two functors, the first with profiler off and the second with profiler
# on. The first functor is usually used for warmup and the second for actual
# benchmarking. This way, one can collect stats of non-warmup
# benchmark iterations using `nsys profile --capture-range=cudaProfilerApi`.
#
# https://docs.nvidia.com/nsight-systems/UserGuide/index.html#handling-application-launchers-mpirun-deepspeed-etc
# has described several ways to profile multi-process applications.
#
# For single-node profiling, I recommend putting `nsys profile` before
# `mpirun`, e.g., `nsys profile ... mpirun -np 8 ...` instead of `mpirun -np 8
# nsys profile ...` or `mpirun -np 1 nsys profile ... : -np 7 ...`. This config
# tries to collect and align traces on different GPUs so it gives the most
# complete picture.  See
# https://github.com/NVIDIA/Fuser/pull/5751/files#r2663586669 for my
# experiment.
def get_benchmark_fns(func):
    return get_benchmark_fn(func, profile=False), get_benchmark_fn(func, profile=True)


wrapper

# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import torch

def get_benchmark_fn(func, /, profile: bool):
    def wrapper(*args, **kwargs):
        if profile:
            torch.cuda.cudart().cudaProfilerStart()
        result = func(*args, **kwargs)
        torch.cuda.synchronize()
        if profile:
            torch.cuda.cudart().cudaProfilerStop()
        return result

    return wrapper

# Returns two functors, the first with profiler off and the second with profiler
# on. The first functor is usually used for warmup and the second for actual
# benchmarking. This way, one can collect stats of non-warmup
# benchmark iterations using `nsys profile --capture-range=cudaProfilerApi`.
#
# https://docs.nvidia.com/nsight-systems/UserGuide/index.html#handling-application-launchers-mpirun-deepspeed-etc
# has described several ways to profile multi-process applications.
#
# For single-node profiling, I recommend putting `nsys profile` before
# `mpirun`, e.g., `nsys profile ... mpirun -np 8 ...` instead of `mpirun -np 8
# nsys profile ...` or `mpirun -np 1 nsys profile ... : -np 7 ...`. This config
# tries to collect and align traces on different GPUs so it gives the most
# complete picture.  See
# https://github.com/NVIDIA/Fuser/pull/5751/files#r2663586669 for my
# experiment.
def get_benchmark_fns(func):


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import pytest
import torch
from enum import Enum, auto
from torch.nn.attention import SDPBackend

import nvfuser_direct as nvfuser
from nvfuser_direct import DataType, FusionDefinition
from python.direct_utils import is_pre_ampere

@pytest.mark.mpi
def test_sizes_and_ranks(multidevice_test):
    size, rank, local_size, local_rank = (
        multidevice_test.size,
        multidevice_test.rank,
        multidevice_test.local_size,
        multidevice_test.local_rank,
    )
    assert size > 0
    assert rank >= 0 and rank < size
    assert local_size > 0
    assert local_rank >= 0 and local_rank < local_size

@pytest.mark.mpi
def test_pointwise(multidevice_test):
    num_devices = multidevice_test.size

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor((-1, -1), contiguity=False, dtype=DataType.Float)
        tv1 = fd.ops.relu(inp_tv)
        tv2 = fd.ops.add(tv1, tv1)
        fd.add_output(tv2)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(num_devices))
        for tv in [inp_tv, tv1, tv2]:
            tv.set_device_mesh(mesh)

        inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

    inp_ref = torch.randn(num_devices, 4)
    out_ref = inp_ref.relu() * 2

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    (out,) = fd.execute([inp])
    torch.testing.assert_close(out.cpu(), out_ref)

@pytest.mark.mpi
def test_pointwise_vectorization(multidevice_test):
    # Repro for https://github.com/NVIDIA/Fuser/issues/5920
    d = multidevice_test.size
    cp_size = 2
    if d % (cp_size * cp_size) != 0:
        pytest.skip(
            f"We only support even split, so {d=} has to be divisible by {cp_size * cp_size} for {cp_size=}."
        )
    dp_size = d // (cp_size * cp_size)

    c = 128
    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor(
            (-1, c, -1, -1, cp_size), contiguity=True, dtype=DataType.BFloat16
        )
        out_tv = fd.ops.set(inp_tv)
        fd.add_output(out_tv)

        mesh = nvfuser.multidevice.DeviceMesh(
            torch.arange(d).reshape(dp_size, cp_size, cp_size)
        )
        for tv in [inp_tv, out_tv]:
            tv.set_device_mesh(mesh)

        inp_tv.axis(4).parallelize(nvfuser.ParallelType.mesh_y)
        inp_tv.outer_split(3, cp_size)
        inp_tv.axis(3).parallelize(nvfuser.ParallelType.mesh_x)
        inp_tv.outer_split(0, dp_size)
        inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_z)

        out_tv.axis(4).parallelize(nvfuser.ParallelType.mesh_y)
        out_tv.outer_split(3, cp_size)
        out_tv.axis(3).parallelize(nvfuser.ParallelType.mesh_x)
        out_tv.outer_split(0, dp_size)
        out_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_z)
        out_tv.set_allocation_domain(
            (
                out_tv.axis(3),
                out_tv.axis(0),
                out_tv.axis(1),
                out_tv.axis(2),
                out_tv.axis(4),
                out_tv.axis(5),
                out_tv.axis(6),
            ),
            True,
        )

    b = dp_size * 3
    s = cp_size * 5
    inp_ref = torch.randn(b, c, s, s, cp_size, dtype=torch.bfloat16)
    out_ref = inp_ref

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    (out,) = fd.execute([inp])
    torch.testing.assert_close(out, multidevice_test.shard_tensor(out_ref, out_tv))

class QkvFormat(Enum):
    BHSE = auto()
    BSHE = auto()

@pytest.mark.skipif(
    is_pre_ampere(),
    reason="Flash Attention is only supported on Ampere and newer devices.",
)
@pytest.mark.parametrize("qkv_format", [QkvFormat.BHSE, QkvFormat.BSHE])
@pytest.mark.mpi
def test_sdpa(multidevice_test, qkv_format: QkvFormat):
    d, b, s, h, e = multidevice_test.size, 2, 1024, 12, 768

    if h % d != 0:
        pytest.skip(f"We only support even split, so {h} has to be divisible by {d}.")

    def _definition(fd: FusionDefinition, qkv_format: QkvFormat) -> None:
        match qkv_format:
            case QkvFormat.BHSE:
                stride_order = [4, 3, 2, 1, 0]
            case QkvFormat.BSHE:
                stride_order = [4, 3, 1, 2, 0]

        q, k, v, out_grad = [
            fd.define_tensor(
                shape=[d, b, h // d, s, e // h],
                contiguity=True,
                dtype=DataType.BFloat16,
                stride_order=stride_order,
            )
            for _ in range(4)
        ]

        # TODO(#3123): support sharded dropout and change this to a
        # positive probability.
        dropout_p = fd.define_scalar(0.0, dtype=DataType.Double)
        is_causal = fd.define_scalar(True, dtype=DataType.Bool)
        attn, logsumexp, seed, offset = fd.ops.sdpfa_fwd(
            q, k, v, dropout_p=dropout_p, is_causal=is_causal
        )

        q_grad, k_grad, v_grad = fd.ops.sdpfa_bwd(
            out_grad,
            q,
            k,
            v,
            attn,
            logsumexp,
            dropout_p,
            is_causal,
            seed,
            offset,
            scale=None,
        )

        fd.add_output(attn)
        for grad in [q_grad, k_grad, v_grad]:
            fd.add_output(grad)

    def _multidevice_schedule(fd: FusionDefinition) -> None:
        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
        for t in fd.fusion.inputs():
            t.set_device_mesh(mesh)
            t.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

    def make_unsharded_tensor() -> torch.Tensor:
        return torch.randn(b, h, s, e // h, dtype=torch.bfloat16, device="cuda")

    q, k, v = [make_unsharded_tensor().requires_grad_() for _ in range(3)]
    out_grad = make_unsharded_tensor()

    with torch.nn.attention.sdpa_kernel(SDPBackend.FLASH_ATTENTION):
        expected_out = torch.nn.functional.scaled_dot_product_attention(
            q, k, v, dropout_p=0.0, is_causal=True, scale=None
        )
        expected_out.backward(out_grad)
        expected_q_grad, expected_k_grad, expected_v_grad = q.grad, k.grad, v.grad

    rank = multidevice_test.rank

    # Head-parallelize Q, K, V or the attention output of an SDPA.
    def head_parallelize(t: torch.Tensor) -> torch.Tensor:
        assert t.shape == torch.Size([b, h, s, e // h])
        # The input `t` may require gradient. We .detach() so the new `t`
        # doesn't require gradient.
        t = t.detach().view([b, d, h // d, s, e // h]).narrow(1, rank, 1)
        match qkv_format:
            case QkvFormat.BHSE:
                return t.transpose(0, 1).contiguous()
            case QkvFormat.BSHE:
                return t.permute(1, 0, 3, 2, 4).contiguous().transpose(2, 3)

    with FusionDefinition() as fd:
        _definition(fd, qkv_format)
        _multidevice_schedule(fd)

    out, q_grad, k_grad, v_grad = fd.execute(
        [
            head_parallelize(q).requires_grad_(),
            head_parallelize(k).requires_grad_(),
            head_parallelize(v).requires_grad_(),
            head_parallelize(out_grad),
        ]
    )

    def assert_close(actual: torch.Tensor, expected: torch.Tensor) -> None:
        match qkv_format:
            case QkvFormat.BHSE:
                assert actual.is_contiguous()
            case QkvFormat.BSHE:
                assert actual.transpose(2, 3).is_contiguous()

        # Use the default rtol for bfloat16 and a relaxed atol.
        torch.testing.assert_close(actual, expected, rtol=1.6e-2, atol=1e-2)

    assert_close(out, head_parallelize(expected_out))
    assert_close(q_grad, head_parallelize(expected_q_grad))
    assert_close(k_grad, head_parallelize(expected_k_grad))
    assert_close(v_grad, head_parallelize(expected_v_grad))

@pytest.mark.skipif(
    is_pre_ampere(),
    reason="Flash Attention is only supported on Ampere and newer devices.",
)
@pytest.mark.parametrize("qkv_format", [QkvFormat.BHSE, QkvFormat.BSHE])
@pytest.mark.mpi
def test_sdpa_loop_split(multidevice_test, qkv_format: QkvFormat):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
    h, e = 12, 768
    if h % d != 0:
        pytest.skip(f"We only support even split, so {h} has to be divisible by {d}.")

    def _definition(fd: FusionDefinition, qkv_format: QkvFormat) -> None:
        match qkv_format:
            case QkvFormat.BHSE:
                stride_order = [3, 2, 1, 0]
            case QkvFormat.BSHE:
                stride_order = [3, 1, 2, 0]

        q, k, v, out_grad = [
            fd.define_tensor(
                shape=[-1, h, -1, e // h],
                dtype=DataType.BFloat16,
                stride_order=stride_order,
            )
            for _ in range(4)
        ]

        # TODO(#3123): support sharded dropout and change this to a
        # positive probability.
        dropout_p = fd.define_scalar(0.0, dtype=DataType.Double)
        is_causal = fd.define_scalar(True, dtype=DataType.Bool)
        attn, logsumexp, seed, offset = fd.ops.sdpfa_fwd(
            q, k, v, dropout_p=dropout_p, is_causal=is_causal
        )

        q_grad, k_grad, v_grad = fd.ops.sdpfa_bwd(
            out_grad,
            q,
            k,
            v,
            attn,
            logsumexp,
            dropout_p,
            is_causal,
            seed,
            offset,
            scale=None,
        )

        fd.add_output(attn)
        for grad in [q_grad, k_grad, v_grad]:
            fd.add_output(grad)

    def _multidevice_schedule(fd: FusionDefinition) -> None:
        for t in fd.fusion.inputs():
            t.set_device_mesh(mesh)
            t.outer_split(1, d)
            t.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

    b, s = 2, 1024

    def make_unsharded_tensor() -> torch.Tensor:
        return torch.randn(b, h, s, e // h, dtype=torch.bfloat16, device="cpu")

    q, k, v = [make_unsharded_tensor().requires_grad_() for _ in range(3)]
    out_grad = make_unsharded_tensor()
    sharded_q, sharded_k, sharded_v, sharded_out_grad = [
        multidevice_test.shard_tensor_1d(t, 1, mesh) for t in [q, k, v, out_grad]
    ]

    with torch.nn.attention.sdpa_kernel(SDPBackend.FLASH_ATTENTION):
        expected_out = torch.nn.functional.scaled_dot_product_attention(
            q, k, v, dropout_p=0.0, is_causal=True, scale=None
        )
        expected_out.backward(out_grad)
        expected_q_grad, expected_k_grad, expected_v_grad = q.grad, k.grad, v.grad

    with FusionDefinition() as fd:
        _definition(fd, qkv_format)
        _multidevice_schedule(fd)

    def reformat_tensor(t: torch.Tensor) -> torch.Tensor:
        match qkv_format:
            case QkvFormat.BHSE:
                return t
            case QkvFormat.BSHE:
                return t.transpose(1, 2).contiguous().transpose(1, 2)

    attn, q_grad, k_grad, v_grad = fd.execute(
        [
            reformat_tensor(sharded_q).requires_grad_(),
            reformat_tensor(sharded_k).requires_grad_(),
            reformat_tensor(sharded_v).requires_grad_(),
            reformat_tensor(sharded_out_grad),
        ]
    )

    def assert_close(actual, expected):
        match qkv_format:
            case QkvFormat.BHSE:
                assert actual.is_contiguous()
            case QkvFormat.BSHE:
                assert actual.transpose(1, 2).is_contiguous()

        # Use the default rtol for bfloat16 and a relaxed atol.
        torch.testing.assert_close(
            actual,
            multidevice_test.shard_tensor_1d(expected, 1, mesh),
            rtol=1.6e-2,
            atol=1e-2,
        )

    for actual, expected in zip(
        [attn, q_grad, k_grad, v_grad],
        [expected_out, expected_q_grad, expected_k_grad, expected_v_grad],
    ):
        assert_close(actual, expected)

@pytest.mark.mpi
def test_privatize_squeeze(multidevice_test):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor((-1, -1, -1), dtype=DataType.BFloat16)
        tv1 = fd.ops.broadcast(inp_tv, [True, False, False, False])
        tv2 = fd.ops.squeeze(tv1, dims=[0])
        tv3 = fd.ops.cast(tv2, dtype=DataType.Float)
        tv4 = fd.ops.sum(tv3, dims=[0])
        tv5 = fd.ops.sum(tv3, dims=[1])
        fd.add_output(tv4)
        fd.add_output(tv5)

        inp_tv.set_device_mesh(mesh)
        inp_tv.outer_split(0, d)
        inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

    inp_ref = torch.randn(d * 3, 5, 6, dtype=torch.bfloat16)
    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)

    out1, out2 = fd.execute([inp])
    torch.testing.assert_close(out1.cpu(), inp_ref.to(torch.float).sum(0))
    torch.testing.assert_close(out2, inp.to(torch.float).sum(1))

@pytest.mark.mpi
def test_inner_reduction(multidevice_test):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor((-1, -1), dtype=DataType.Float)
        out = fd.ops.sum(inp_tv, [1])
        fd.add_output(out)

        inp_tv.set_device_mesh(mesh)
        inp_tv.outer_split(0, d)
        inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

    inp_ref = torch.ones(d * 3, 5)

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    (out,) = fd.execute([inp])
    torch.testing.assert_close(out, inp.sum(1))

@pytest.mark.mpi
def test_insert_resharding_after(multidevice_test):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor((-1, -1), contiguity=True)
        out = fd.ops.relu(inp_tv)
        fd.add_output(out)

        inp_tv.set_device_mesh(mesh)
        inp_tv.outer_split(0, d)
        inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

        out.set_device_mesh(mesh)

    inp_ref = torch.randn(d * 3, 5)
    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)

    (out,) = fd.execute([inp])
    torch.testing.assert_close(out.cpu(), inp_ref.relu())

@pytest.mark.mpi
def test_welford(multidevice_test):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
    b, s, e = 1, 2048, 12288

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor(shape=[b, s, e], contiguity=True)
        var, mean = fd.ops.var_mean(inp_tv, dims=[2], correction=0, keepdim=False)
        fd.add_output(var)
        fd.add_output(mean)

        inp_tv.set_device_mesh(mesh)
        inp_tv.outer_split(1, d)
        inp_tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

    inp_ref = torch.randn(b, s, e)
    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)

    var, mean = fd.execute([inp])

    torch.testing.assert_close(var, inp.var(2), rtol=1e-3, atol=1e-3)
    torch.testing.assert_close(mean, inp.mean(2), rtol=1e-3, atol=1e-3)

@pytest.mark.mpi
def test_binary(multidevice_test):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))

    with FusionDefinition() as fd:
        x_tv = fd.define_tensor((-1, -1), contiguity=True, dtype=DataType.Half)
        y_tv = fd.define_tensor((-1, -1), contiguity=True, dtype=DataType.Half)
        z = fd.ops.add(x_tv, y_tv)
        fd.add_output(z)

        x_tv.set_device_mesh(mesh)
        y_tv.set_device_mesh(mesh)
        y_tv.outer_split(0, d)
        y_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

    x_ref = torch.randn(d * 2, 3, dtype=torch.float16)
    y_ref = torch.randn(d * 2, 3, dtype=torch.float16)
    z_ref = x_ref.float() + y_ref.float()

    x = multidevice_test.shard_tensor(x_ref, x_tv)
    y = multidevice_test.shard_tensor(y_ref, y_tv)
    (z,) = fd.execute([x, y])

    torch.testing.assert_close(z, multidevice_test.shard_tensor_1d(z_ref, 0, mesh))

@pytest.mark.mpi
def test_reduction_with_2d_mesh(multidevice_test):
    d = multidevice_test.size
    tp_size = 2

    # Skip if d is not divisible by tp_size
    if d % tp_size != 0:
        pytest.skip(f"Number of devices ({d}) must be divisible by tp_size ({tp_size})")

    dp_size = d // tp_size

    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d).reshape(dp_size, tp_size))

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor([-1, -1], dtype=DataType.Float, contiguity=True)
        out = fd.ops.sum(inp_tv, [1])
        fd.add_output(out)

        inp_tv.set_device_mesh(mesh)
        inp_tv.outer_split(0, dp_size)
        inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_y)
        inp_tv.outer_split(-1, tp_size)
        inp_tv.axis(-2).parallelize(nvfuser.ParallelType.mesh_x)

    rank = multidevice_test.rank
    rows_per_rank, cols_per_rank = 2, 3
    rows, cols = dp_size * rows_per_rank, tp_size * cols_per_rank

    inp_ref = torch.arange(rows * cols, dtype=torch.float).reshape(rows, cols)
    out_ref = inp_ref.sum([-1])

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    (out,) = fd.execute([inp])

    dp_rank = rank // tp_size
    torch.testing.assert_close(
        out.cpu(), out_ref[dp_rank * rows_per_rank : (dp_rank + 1) * rows_per_rank]
    )


# SPDX-FileCopyrightText: Copyright (c) 2024-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import pytest
import torch
from enum import Enum, auto
from torch.nn.attention import SDPBackend

import nvfuser_direct as nvfuser
from nvfuser_direct import DataType, FusionDefinition
from python.direct_utils import is_pre_ampere

@pytest.mark.mpi
def test_sizes_and_ranks(multidevice_test):
    size, rank, local_size, local_rank = (
        multidevice_test.size,
        multidevice_test.rank,
        multidevice_test.local_size,
        multidevice_test.local_rank,
    )
    assert size > 0
    assert rank >= 0 and rank < size
    assert local_size > 0
    assert local_rank >= 0 and local_rank < local_size

@pytest.mark.mpi
def test_pointwise(multidevice_test):
    num_devices = multidevice_test.size

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor((-1, -1), contiguity=False, dtype=DataType.Float)
        tv1 = fd.ops.relu(inp_tv)
        tv2 = fd.ops.add(tv1, tv1)
        fd.add_output(tv2)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(num_devices))
        for tv in [inp_tv, tv1, tv2]:
            tv.set_device_mesh(mesh)

        inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

    inp_ref = torch.randn(num_devices, 4)
    out_ref = inp_ref.relu() * 2


import nvfuser_direct as nvfuser
from nvfuser_direct import DataType, FusionDefinition
from python.direct_utils import is_pre_ampere

@pytest.mark.mpi
def test_sizes_and_ranks(multidevice_test):
    size, rank, local_size, local_rank = (
        multidevice_test.size,
        multidevice_test.rank,
        multidevice_test.local_size,
        multidevice_test.local_rank,
    )
    assert size > 0
    assert rank >= 0 and rank < size
    assert local_size > 0
    assert local_rank >= 0 and local_rank < local_size

@pytest.mark.mpi
def test_pointwise(multidevice_test):
    num_devices = multidevice_test.size

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor((-1, -1), contiguity=False, dtype=DataType.Float)
        tv1 = fd.ops.relu(inp_tv)
        tv2 = fd.ops.add(tv1, tv1)
        fd.add_output(tv2)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(num_devices))
        for tv in [inp_tv, tv1, tv2]:
            tv.set_device_mesh(mesh)

        inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

    inp_ref = torch.randn(num_devices, 4)
    out_ref = inp_ref.relu() * 2

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    (out,) = fd.execute([inp])
    torch.testing.assert_close(out.cpu(), out_ref)

@pytest.mark.mpi
def test_pointwise_vectorization(multidevice_test):
    # Repro for https://github.com/NVIDIA/Fuser/issues/5920
    d = multidevice_test.size
    cp_size = 2
    if d % (cp_size * cp_size) != 0:
        pytest.skip(
            f"We only support even split, so {d=} has to be divisible by {cp_size * cp_size} for {cp_size=}."
        )
    dp_size = d // (cp_size * cp_size)

    c = 128
    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor(
            (-1, c, -1, -1, cp_size), contiguity=True, dtype=DataType.BFloat16
        )
        out_tv = fd.ops.set(inp_tv)
        fd.add_output(out_tv)


        inp_tv = fd.define_tensor((-1, -1), contiguity=False, dtype=DataType.Float)
        tv1 = fd.ops.relu(inp_tv)
        tv2 = fd.ops.add(tv1, tv1)
        fd.add_output(tv2)

        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(num_devices))
        for tv in [inp_tv, tv1, tv2]:
            tv.set_device_mesh(mesh)

        inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

    inp_ref = torch.randn(num_devices, 4)
    out_ref = inp_ref.relu() * 2

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    (out,) = fd.execute([inp])
    torch.testing.assert_close(out.cpu(), out_ref)

@pytest.mark.mpi
def test_pointwise_vectorization(multidevice_test):
    # Repro for https://github.com/NVIDIA/Fuser/issues/5920
    d = multidevice_test.size
    cp_size = 2
    if d % (cp_size * cp_size) != 0:
        pytest.skip(
            f"We only support even split, so {d=} has to be divisible by {cp_size * cp_size} for {cp_size=}."
        )
    dp_size = d // (cp_size * cp_size)

    c = 128
    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor(
            (-1, c, -1, -1, cp_size), contiguity=True, dtype=DataType.BFloat16
        )
        out_tv = fd.ops.set(inp_tv)
        fd.add_output(out_tv)

        mesh = nvfuser.multidevice.DeviceMesh(
            torch.arange(d).reshape(dp_size, cp_size, cp_size)
        )
        for tv in [inp_tv, out_tv]:
            tv.set_device_mesh(mesh)

        inp_tv.axis(4).parallelize(nvfuser.ParallelType.mesh_y)
        inp_tv.outer_split(3, cp_size)
        inp_tv.axis(3).parallelize(nvfuser.ParallelType.mesh_x)
        inp_tv.outer_split(0, dp_size)
        inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_z)

        out_tv.axis(4).parallelize(nvfuser.ParallelType.mesh_y)
        out_tv.outer_split(3, cp_size)
        out_tv.axis(3).parallelize(nvfuser.ParallelType.mesh_x)
        out_tv.outer_split(0, dp_size)
        out_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_z)
        out_tv.set_allocation_domain(
            (
                out_tv.axis(3),
                out_tv.axis(0),
                out_tv.axis(1),
                out_tv.axis(2),
                out_tv.axis(4),
                out_tv.axis(5),
                out_tv.axis(6),
            ),
            True,
        )

    b = dp_size * 3
    s = cp_size * 5
    inp_ref = torch.randn(b, c, s, s, cp_size, dtype=torch.bfloat16)
    out_ref = inp_ref

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    (out,) = fd.execute([inp])
    torch.testing.assert_close(out, multidevice_test.shard_tensor(out_ref, out_tv))

class QkvFormat(Enum):
    BHSE = auto()
    BSHE = auto()

@pytest.mark.skipif(
    is_pre_ampere(),
    reason="Flash Attention is only supported on Ampere and newer devices.",
)
@pytest.mark.parametrize("qkv_format", [QkvFormat.BHSE, QkvFormat.BSHE])
@pytest.mark.mpi
def test_sdpa(multidevice_test, qkv_format: QkvFormat):
    d, b, s, h, e = multidevice_test.size, 2, 1024, 12, 768

    if h % d != 0:
        pytest.skip(f"We only support even split, so {h} has to be divisible by {d}.")

    def _definition(fd: FusionDefinition, qkv_format: QkvFormat) -> None:


QkvFormat

                out_tv.axis(0),
                out_tv.axis(1),
                out_tv.axis(2),
                out_tv.axis(4),
                out_tv.axis(5),
                out_tv.axis(6),
            ),
            True,
        )

    b = dp_size * 3
    s = cp_size * 5
    inp_ref = torch.randn(b, c, s, s, cp_size, dtype=torch.bfloat16)
    out_ref = inp_ref

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    (out,) = fd.execute([inp])
    torch.testing.assert_close(out, multidevice_test.shard_tensor(out_ref, out_tv))

class QkvFormat(Enum):
    BHSE = auto()
    BSHE = auto()

@pytest.mark.skipif(
    is_pre_ampere(),
    reason="Flash Attention is only supported on Ampere and newer devices.",
)
@pytest.mark.parametrize("qkv_format", [QkvFormat.BHSE, QkvFormat.BSHE])
@pytest.mark.mpi
def test_sdpa(multidevice_test, qkv_format: QkvFormat):
    d, b, s, h, e = multidevice_test.size, 2, 1024, 12, 768

    if h % d != 0:
        pytest.skip(f"We only support even split, so {h} has to be divisible by {d}.")

    def _definition(fd: FusionDefinition, qkv_format: QkvFormat) -> None:
        match qkv_format:
            case QkvFormat.BHSE:
                stride_order = [4, 3, 2, 1, 0]
            case QkvFormat.BSHE:
                stride_order = [4, 3, 1, 2, 0]


    s = cp_size * 5
    inp_ref = torch.randn(b, c, s, s, cp_size, dtype=torch.bfloat16)
    out_ref = inp_ref

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    (out,) = fd.execute([inp])
    torch.testing.assert_close(out, multidevice_test.shard_tensor(out_ref, out_tv))

class QkvFormat(Enum):
    BHSE = auto()
    BSHE = auto()

@pytest.mark.skipif(
    is_pre_ampere(),
    reason="Flash Attention is only supported on Ampere and newer devices.",
)
@pytest.mark.parametrize("qkv_format", [QkvFormat.BHSE, QkvFormat.BSHE])
@pytest.mark.mpi
def test_sdpa(multidevice_test, qkv_format: QkvFormat):
    d, b, s, h, e = multidevice_test.size, 2, 1024, 12, 768

    if h % d != 0:
        pytest.skip(f"We only support even split, so {h} has to be divisible by {d}.")

    def _definition(fd: FusionDefinition, qkv_format: QkvFormat) -> None:
        match qkv_format:
            case QkvFormat.BHSE:
                stride_order = [4, 3, 2, 1, 0]
            case QkvFormat.BSHE:
                stride_order = [4, 3, 1, 2, 0]

        q, k, v, out_grad = [
            fd.define_tensor(
                shape=[d, b, h // d, s, e // h],
                contiguity=True,
                dtype=DataType.BFloat16,
                stride_order=stride_order,
            )
            for _ in range(4)
        ]

        # TODO(#3123): support sharded dropout and change this to a
        # positive probability.
        dropout_p = fd.define_scalar(0.0, dtype=DataType.Double)
        is_causal = fd.define_scalar(True, dtype=DataType.Bool)
        attn, logsumexp, seed, offset = fd.ops.sdpfa_fwd(
            q, k, v, dropout_p=dropout_p, is_causal=is_causal
        )

        q_grad, k_grad, v_grad = fd.ops.sdpfa_bwd(
            out_grad,
            q,
            k,
            v,
            attn,
            logsumexp,
            dropout_p,
            is_causal,
            seed,
            offset,
            scale=None,
        )

        fd.add_output(attn)
        for grad in [q_grad, k_grad, v_grad]:
            fd.add_output(grad)

    def _multidevice_schedule(fd: FusionDefinition) -> None:
        mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
        for t in fd.fusion.inputs():
            t.set_device_mesh(mesh)
            t.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

    def make_unsharded_tensor() -> torch.Tensor:
        return torch.randn(b, h, s, e // h, dtype=torch.bfloat16, device="cuda")

    q, k, v = [make_unsharded_tensor().requires_grad_() for _ in range(3)]
    out_grad = make_unsharded_tensor()

    with torch.nn.attention.sdpa_kernel(SDPBackend.FLASH_ATTENTION):
        expected_out = torch.nn.functional.scaled_dot_product_attention(
            q, k, v, dropout_p=0.0, is_causal=True, scale=None
        )
        expected_out.backward(out_grad)
        expected_q_grad, expected_k_grad, expected_v_grad = q.grad, k.grad, v.grad

    rank = multidevice_test.rank

    # Head-parallelize Q, K, V or the attention output of an SDPA.
    def head_parallelize(t: torch.Tensor) -> torch.Tensor:
        assert t.shape == torch.Size([b, h, s, e // h])
        # The input `t` may require gradient. We .detach() so the new `t`
        # doesn't require gradient.
        t = t.detach().view([b, d, h // d, s, e // h]).narrow(1, rank, 1)
        match qkv_format:
            case QkvFormat.BHSE:
                return t.transpose(0, 1).contiguous()
            case QkvFormat.BSHE:
                return t.permute(1, 0, 3, 2, 4).contiguous().transpose(2, 3)

    with FusionDefinition() as fd:
        _definition(fd, qkv_format)
        _multidevice_schedule(fd)

    out, q_grad, k_grad, v_grad = fd.execute(
        [
            head_parallelize(q).requires_grad_(),
            head_parallelize(k).requires_grad_(),
            head_parallelize(v).requires_grad_(),
            head_parallelize(out_grad),
        ]
    )

    def assert_close(actual: torch.Tensor, expected: torch.Tensor) -> None:
        match qkv_format:
            case QkvFormat.BHSE:
                assert actual.is_contiguous()
            case QkvFormat.BSHE:
                assert actual.transpose(2, 3).is_contiguous()

        # Use the default rtol for bfloat16 and a relaxed atol.
        torch.testing.assert_close(actual, expected, rtol=1.6e-2, atol=1e-2)

    assert_close(out, head_parallelize(expected_out))
    assert_close(q_grad, head_parallelize(expected_q_grad))
    assert_close(k_grad, head_parallelize(expected_k_grad))
    assert_close(v_grad, head_parallelize(expected_v_grad))

@pytest.mark.skipif(
    is_pre_ampere(),
    reason="Flash Attention is only supported on Ampere and newer devices.",
)
@pytest.mark.parametrize("qkv_format", [QkvFormat.BHSE, QkvFormat.BSHE])
@pytest.mark.mpi
def test_sdpa_loop_split(multidevice_test, qkv_format: QkvFormat):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
    h, e = 12, 768
    if h % d != 0:
        pytest.skip(f"We only support even split, so {h} has to be divisible by {d}.")

    def _definition(fd: FusionDefinition, qkv_format: QkvFormat) -> None:
        match qkv_format:
            case QkvFormat.BHSE:
                stride_order = [3, 2, 1, 0]
            case QkvFormat.BSHE:


            case QkvFormat.BHSE:
                assert actual.is_contiguous()
            case QkvFormat.BSHE:
                assert actual.transpose(2, 3).is_contiguous()

        # Use the default rtol for bfloat16 and a relaxed atol.
        torch.testing.assert_close(actual, expected, rtol=1.6e-2, atol=1e-2)

    assert_close(out, head_parallelize(expected_out))
    assert_close(q_grad, head_parallelize(expected_q_grad))
    assert_close(k_grad, head_parallelize(expected_k_grad))
    assert_close(v_grad, head_parallelize(expected_v_grad))

@pytest.mark.skipif(
    is_pre_ampere(),
    reason="Flash Attention is only supported on Ampere and newer devices.",
)
@pytest.mark.parametrize("qkv_format", [QkvFormat.BHSE, QkvFormat.BSHE])
@pytest.mark.mpi
def test_sdpa_loop_split(multidevice_test, qkv_format: QkvFormat):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
    h, e = 12, 768
    if h % d != 0:
        pytest.skip(f"We only support even split, so {h} has to be divisible by {d}.")

    def _definition(fd: FusionDefinition, qkv_format: QkvFormat) -> None:
        match qkv_format:
            case QkvFormat.BHSE:
                stride_order = [3, 2, 1, 0]
            case QkvFormat.BSHE:
                stride_order = [3, 1, 2, 0]

        q, k, v, out_grad = [
            fd.define_tensor(
                shape=[-1, h, -1, e // h],
                dtype=DataType.BFloat16,
                stride_order=stride_order,
            )
            for _ in range(4)
        ]

        # TODO(#3123): support sharded dropout and change this to a
        # positive probability.
        dropout_p = fd.define_scalar(0.0, dtype=DataType.Double)
        is_causal = fd.define_scalar(True, dtype=DataType.Bool)
        attn, logsumexp, seed, offset = fd.ops.sdpfa_fwd(
            q, k, v, dropout_p=dropout_p, is_causal=is_causal
        )

        q_grad, k_grad, v_grad = fd.ops.sdpfa_bwd(
            out_grad,
            q,
            k,
            v,
            attn,
            logsumexp,
            dropout_p,
            is_causal,
            seed,
            offset,
            scale=None,
        )

        fd.add_output(attn)
        for grad in [q_grad, k_grad, v_grad]:
            fd.add_output(grad)

    def _multidevice_schedule(fd: FusionDefinition) -> None:
        for t in fd.fusion.inputs():
            t.set_device_mesh(mesh)
            t.outer_split(1, d)
            t.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

    b, s = 2, 1024

    def make_unsharded_tensor() -> torch.Tensor:
        return torch.randn(b, h, s, e // h, dtype=torch.bfloat16, device="cpu")

    q, k, v = [make_unsharded_tensor().requires_grad_() for _ in range(3)]
    out_grad = make_unsharded_tensor()
    sharded_q, sharded_k, sharded_v, sharded_out_grad = [
        multidevice_test.shard_tensor_1d(t, 1, mesh) for t in [q, k, v, out_grad]
    ]

    with torch.nn.attention.sdpa_kernel(SDPBackend.FLASH_ATTENTION):
        expected_out = torch.nn.functional.scaled_dot_product_attention(
            q, k, v, dropout_p=0.0, is_causal=True, scale=None
        )
        expected_out.backward(out_grad)
        expected_q_grad, expected_k_grad, expected_v_grad = q.grad, k.grad, v.grad

    with FusionDefinition() as fd:
        _definition(fd, qkv_format)
        _multidevice_schedule(fd)

    def reformat_tensor(t: torch.Tensor) -> torch.Tensor:
        match qkv_format:
            case QkvFormat.BHSE:
                return t
            case QkvFormat.BSHE:
                return t.transpose(1, 2).contiguous().transpose(1, 2)

    attn, q_grad, k_grad, v_grad = fd.execute(
        [
            reformat_tensor(sharded_q).requires_grad_(),
            reformat_tensor(sharded_k).requires_grad_(),
            reformat_tensor(sharded_v).requires_grad_(),
            reformat_tensor(sharded_out_grad),
        ]
    )

    def assert_close(actual, expected):
        match qkv_format:
            case QkvFormat.BHSE:
                assert actual.is_contiguous()
            case QkvFormat.BSHE:
                assert actual.transpose(1, 2).is_contiguous()

        # Use the default rtol for bfloat16 and a relaxed atol.
        torch.testing.assert_close(
            actual,
            multidevice_test.shard_tensor_1d(expected, 1, mesh),
            rtol=1.6e-2,
            atol=1e-2,
        )

    for actual, expected in zip(
        [attn, q_grad, k_grad, v_grad],
        [expected_out, expected_q_grad, expected_k_grad, expected_v_grad],
    ):
        assert_close(actual, expected)

@pytest.mark.mpi
def test_privatize_squeeze(multidevice_test):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor((-1, -1, -1), dtype=DataType.BFloat16)
        tv1 = fd.ops.broadcast(inp_tv, [True, False, False, False])
        tv2 = fd.ops.squeeze(tv1, dims=[0])
        tv3 = fd.ops.cast(tv2, dtype=DataType.Float)
        tv4 = fd.ops.sum(tv3, dims=[0])
        tv5 = fd.ops.sum(tv3, dims=[1])
        fd.add_output(tv4)
        fd.add_output(tv5)

        inp_tv.set_device_mesh(mesh)
        inp_tv.outer_split(0, d)
        inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)


                assert actual.is_contiguous()
            case QkvFormat.BSHE:
                assert actual.transpose(1, 2).is_contiguous()

        # Use the default rtol for bfloat16 and a relaxed atol.
        torch.testing.assert_close(
            actual,
            multidevice_test.shard_tensor_1d(expected, 1, mesh),
            rtol=1.6e-2,
            atol=1e-2,
        )

    for actual, expected in zip(
        [attn, q_grad, k_grad, v_grad],
        [expected_out, expected_q_grad, expected_k_grad, expected_v_grad],
    ):
        assert_close(actual, expected)

@pytest.mark.mpi
def test_privatize_squeeze(multidevice_test):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor((-1, -1, -1), dtype=DataType.BFloat16)
        tv1 = fd.ops.broadcast(inp_tv, [True, False, False, False])
        tv2 = fd.ops.squeeze(tv1, dims=[0])
        tv3 = fd.ops.cast(tv2, dtype=DataType.Float)
        tv4 = fd.ops.sum(tv3, dims=[0])
        tv5 = fd.ops.sum(tv3, dims=[1])
        fd.add_output(tv4)
        fd.add_output(tv5)

        inp_tv.set_device_mesh(mesh)
        inp_tv.outer_split(0, d)
        inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

    inp_ref = torch.randn(d * 3, 5, 6, dtype=torch.bfloat16)
    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)

    out1, out2 = fd.execute([inp])
    torch.testing.assert_close(out1.cpu(), inp_ref.to(torch.float).sum(0))
    torch.testing.assert_close(out2, inp.to(torch.float).sum(1))

@pytest.mark.mpi
def test_inner_reduction(multidevice_test):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor((-1, -1), dtype=DataType.Float)
        out = fd.ops.sum(inp_tv, [1])
        fd.add_output(out)

        inp_tv.set_device_mesh(mesh)
        inp_tv.outer_split(0, d)
        inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

    inp_ref = torch.ones(d * 3, 5)

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    (out,) = fd.execute([inp])
    torch.testing.assert_close(out, inp.sum(1))


        tv2 = fd.ops.squeeze(tv1, dims=[0])
        tv3 = fd.ops.cast(tv2, dtype=DataType.Float)
        tv4 = fd.ops.sum(tv3, dims=[0])
        tv5 = fd.ops.sum(tv3, dims=[1])
        fd.add_output(tv4)
        fd.add_output(tv5)

        inp_tv.set_device_mesh(mesh)
        inp_tv.outer_split(0, d)
        inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

    inp_ref = torch.randn(d * 3, 5, 6, dtype=torch.bfloat16)
    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)

    out1, out2 = fd.execute([inp])
    torch.testing.assert_close(out1.cpu(), inp_ref.to(torch.float).sum(0))
    torch.testing.assert_close(out2, inp.to(torch.float).sum(1))

@pytest.mark.mpi
def test_inner_reduction(multidevice_test):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor((-1, -1), dtype=DataType.Float)
        out = fd.ops.sum(inp_tv, [1])
        fd.add_output(out)

        inp_tv.set_device_mesh(mesh)
        inp_tv.outer_split(0, d)
        inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

    inp_ref = torch.ones(d * 3, 5)

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    (out,) = fd.execute([inp])
    torch.testing.assert_close(out, inp.sum(1))

@pytest.mark.mpi
def test_insert_resharding_after(multidevice_test):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor((-1, -1), contiguity=True)
        out = fd.ops.relu(inp_tv)
        fd.add_output(out)

        inp_tv.set_device_mesh(mesh)
        inp_tv.outer_split(0, d)
        inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

        out.set_device_mesh(mesh)

    inp_ref = torch.randn(d * 3, 5)
    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)


def test_inner_reduction(multidevice_test):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor((-1, -1), dtype=DataType.Float)
        out = fd.ops.sum(inp_tv, [1])
        fd.add_output(out)

        inp_tv.set_device_mesh(mesh)
        inp_tv.outer_split(0, d)
        inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

    inp_ref = torch.ones(d * 3, 5)

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    (out,) = fd.execute([inp])
    torch.testing.assert_close(out, inp.sum(1))

@pytest.mark.mpi
def test_insert_resharding_after(multidevice_test):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor((-1, -1), contiguity=True)
        out = fd.ops.relu(inp_tv)
        fd.add_output(out)

        inp_tv.set_device_mesh(mesh)
        inp_tv.outer_split(0, d)
        inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

        out.set_device_mesh(mesh)

    inp_ref = torch.randn(d * 3, 5)
    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)

    (out,) = fd.execute([inp])
    torch.testing.assert_close(out.cpu(), inp_ref.relu())

@pytest.mark.mpi
def test_welford(multidevice_test):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
    b, s, e = 1, 2048, 12288

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor(shape=[b, s, e], contiguity=True)
        var, mean = fd.ops.var_mean(inp_tv, dims=[2], correction=0, keepdim=False)
        fd.add_output(var)
        fd.add_output(mean)

        inp_tv.set_device_mesh(mesh)
        inp_tv.outer_split(1, d)
        inp_tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

    inp_ref = torch.randn(b, s, e)
    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)


    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor((-1, -1), contiguity=True)
        out = fd.ops.relu(inp_tv)
        fd.add_output(out)

        inp_tv.set_device_mesh(mesh)
        inp_tv.outer_split(0, d)
        inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

        out.set_device_mesh(mesh)

    inp_ref = torch.randn(d * 3, 5)
    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)

    (out,) = fd.execute([inp])
    torch.testing.assert_close(out.cpu(), inp_ref.relu())

@pytest.mark.mpi
def test_welford(multidevice_test):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
    b, s, e = 1, 2048, 12288

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor(shape=[b, s, e], contiguity=True)
        var, mean = fd.ops.var_mean(inp_tv, dims=[2], correction=0, keepdim=False)
        fd.add_output(var)
        fd.add_output(mean)

        inp_tv.set_device_mesh(mesh)
        inp_tv.outer_split(1, d)
        inp_tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

    inp_ref = torch.randn(b, s, e)
    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)

    var, mean = fd.execute([inp])

    torch.testing.assert_close(var, inp.var(2), rtol=1e-3, atol=1e-3)
    torch.testing.assert_close(mean, inp.mean(2), rtol=1e-3, atol=1e-3)

@pytest.mark.mpi
def test_binary(multidevice_test):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))

    with FusionDefinition() as fd:
        x_tv = fd.define_tensor((-1, -1), contiguity=True, dtype=DataType.Half)
        y_tv = fd.define_tensor((-1, -1), contiguity=True, dtype=DataType.Half)
        z = fd.ops.add(x_tv, y_tv)
        fd.add_output(z)

        x_tv.set_device_mesh(mesh)
        y_tv.set_device_mesh(mesh)
        y_tv.outer_split(0, d)
        y_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

    x_ref = torch.randn(d * 2, 3, dtype=torch.float16)
    y_ref = torch.randn(d * 2, 3, dtype=torch.float16)


    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor(shape=[b, s, e], contiguity=True)
        var, mean = fd.ops.var_mean(inp_tv, dims=[2], correction=0, keepdim=False)
        fd.add_output(var)
        fd.add_output(mean)

        inp_tv.set_device_mesh(mesh)
        inp_tv.outer_split(1, d)
        inp_tv.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

    inp_ref = torch.randn(b, s, e)
    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)

    var, mean = fd.execute([inp])

    torch.testing.assert_close(var, inp.var(2), rtol=1e-3, atol=1e-3)
    torch.testing.assert_close(mean, inp.mean(2), rtol=1e-3, atol=1e-3)

@pytest.mark.mpi
def test_binary(multidevice_test):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))

    with FusionDefinition() as fd:
        x_tv = fd.define_tensor((-1, -1), contiguity=True, dtype=DataType.Half)
        y_tv = fd.define_tensor((-1, -1), contiguity=True, dtype=DataType.Half)
        z = fd.ops.add(x_tv, y_tv)
        fd.add_output(z)

        x_tv.set_device_mesh(mesh)
        y_tv.set_device_mesh(mesh)
        y_tv.outer_split(0, d)
        y_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

    x_ref = torch.randn(d * 2, 3, dtype=torch.float16)
    y_ref = torch.randn(d * 2, 3, dtype=torch.float16)
    z_ref = x_ref.float() + y_ref.float()

    x = multidevice_test.shard_tensor(x_ref, x_tv)
    y = multidevice_test.shard_tensor(y_ref, y_tv)
    (z,) = fd.execute([x, y])

    torch.testing.assert_close(z, multidevice_test.shard_tensor_1d(z_ref, 0, mesh))

@pytest.mark.mpi
def test_reduction_with_2d_mesh(multidevice_test):
    d = multidevice_test.size
    tp_size = 2

    # Skip if d is not divisible by tp_size
    if d % tp_size != 0:
        pytest.skip(f"Number of devices ({d}) must be divisible by tp_size ({tp_size})")

    dp_size = d // tp_size

    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d).reshape(dp_size, tp_size))

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor([-1, -1], dtype=DataType.Float, contiguity=True)
        out = fd.ops.sum(inp_tv, [1])
        fd.add_output(out)


        z = fd.ops.add(x_tv, y_tv)
        fd.add_output(z)

        x_tv.set_device_mesh(mesh)
        y_tv.set_device_mesh(mesh)
        y_tv.outer_split(0, d)
        y_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

    x_ref = torch.randn(d * 2, 3, dtype=torch.float16)
    y_ref = torch.randn(d * 2, 3, dtype=torch.float16)
    z_ref = x_ref.float() + y_ref.float()

    x = multidevice_test.shard_tensor(x_ref, x_tv)
    y = multidevice_test.shard_tensor(y_ref, y_tv)
    (z,) = fd.execute([x, y])

    torch.testing.assert_close(z, multidevice_test.shard_tensor_1d(z_ref, 0, mesh))

@pytest.mark.mpi
def test_reduction_with_2d_mesh(multidevice_test):
    d = multidevice_test.size
    tp_size = 2

    # Skip if d is not divisible by tp_size
    if d % tp_size != 0:
        pytest.skip(f"Number of devices ({d}) must be divisible by tp_size ({tp_size})")

    dp_size = d // tp_size

    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d).reshape(dp_size, tp_size))

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor([-1, -1], dtype=DataType.Float, contiguity=True)
        out = fd.ops.sum(inp_tv, [1])
        fd.add_output(out)

        inp_tv.set_device_mesh(mesh)
        inp_tv.outer_split(0, dp_size)
        inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_y)
        inp_tv.outer_split(-1, tp_size)
        inp_tv.axis(-2).parallelize(nvfuser.ParallelType.mesh_x)

    rank = multidevice_test.rank
    rows_per_rank, cols_per_rank = 2, 3
    rows, cols = dp_size * rows_per_rank, tp_size * cols_per_rank

    inp_ref = torch.arange(rows * cols, dtype=torch.float).reshape(rows, cols)
    out_ref = inp_ref.sum([-1])

    inp = multidevice_test.shard_tensor(inp_ref, inp_tv)
    (out,) = fd.execute([inp])

    dp_rank = rank // tp_size
    torch.testing.assert_close(
        out.cpu(), out_ref[dp_rank * rows_per_rank : (dp_rank + 1) * rows_per_rank]
    )


    assert_close(out, head_parallelize(expected_out))
    assert_close(q_grad, head_parallelize(expected_q_grad))
    assert_close(k_grad, head_parallelize(expected_k_grad))
    assert_close(v_grad, head_parallelize(expected_v_grad))

@pytest.mark.skipif(
    is_pre_ampere(),
    reason="Flash Attention is only supported on Ampere and newer devices.",
)
@pytest.mark.parametrize("qkv_format", [QkvFormat.BHSE, QkvFormat.BSHE])
@pytest.mark.mpi
def test_sdpa_loop_split(multidevice_test, qkv_format: QkvFormat):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))
    h, e = 12, 768
    if h % d != 0:
        pytest.skip(f"We only support even split, so {h} has to be divisible by {d}.")

    def _definition(fd: FusionDefinition, qkv_format: QkvFormat) -> None:
        match qkv_format:
            case QkvFormat.BHSE:
                stride_order = [3, 2, 1, 0]
            case QkvFormat.BSHE:
                stride_order = [3, 1, 2, 0]

        q, k, v, out_grad = [
            fd.define_tensor(
                shape=[-1, h, -1, e // h],
                dtype=DataType.BFloat16,
                stride_order=stride_order,
            )
            for _ in range(4)
        ]

        # TODO(#3123): support sharded dropout and change this to a
        # positive probability.
        dropout_p = fd.define_scalar(0.0, dtype=DataType.Double)
        is_causal = fd.define_scalar(True, dtype=DataType.Bool)
        attn, logsumexp, seed, offset = fd.ops.sdpfa_fwd(
            q, k, v, dropout_p=dropout_p, is_causal=is_causal
        )

        q_grad, k_grad, v_grad = fd.ops.sdpfa_bwd(
            out_grad,
            q,
            k,
            v,
            attn,
            logsumexp,
            dropout_p,
            is_causal,
            seed,
            offset,
            scale=None,
        )

        fd.add_output(attn)
        for grad in [q_grad, k_grad, v_grad]:
            fd.add_output(grad)

    def _multidevice_schedule(fd: FusionDefinition) -> None:
        for t in fd.fusion.inputs():
            t.set_device_mesh(mesh)
            t.outer_split(1, d)
            t.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

    b, s = 2, 1024

    def make_unsharded_tensor() -> torch.Tensor:
        return torch.randn(b, h, s, e // h, dtype=torch.bfloat16, device="cpu")

    q, k, v = [make_unsharded_tensor().requires_grad_() for _ in range(3)]
    out_grad = make_unsharded_tensor()
    sharded_q, sharded_k, sharded_v, sharded_out_grad = [
        multidevice_test.shard_tensor_1d(t, 1, mesh) for t in [q, k, v, out_grad]
    ]

    with torch.nn.attention.sdpa_kernel(SDPBackend.FLASH_ATTENTION):
        expected_out = torch.nn.functional.scaled_dot_product_attention(


        )

        q_grad, k_grad, v_grad = fd.ops.sdpfa_bwd(
            out_grad,
            q,
            k,
            v,
            attn,
            logsumexp,
            dropout_p,
            is_causal,
            seed,
            offset,
            scale=None,
        )

        fd.add_output(attn)
        for grad in [q_grad, k_grad, v_grad]:
            fd.add_output(grad)

    def _multidevice_schedule(fd: FusionDefinition) -> None:
        for t in fd.fusion.inputs():
            t.set_device_mesh(mesh)
            t.outer_split(1, d)
            t.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

    b, s = 2, 1024

    def make_unsharded_tensor() -> torch.Tensor:
        return torch.randn(b, h, s, e // h, dtype=torch.bfloat16, device="cpu")

    q, k, v = [make_unsharded_tensor().requires_grad_() for _ in range(3)]
    out_grad = make_unsharded_tensor()
    sharded_q, sharded_k, sharded_v, sharded_out_grad = [
        multidevice_test.shard_tensor_1d(t, 1, mesh) for t in [q, k, v, out_grad]
    ]

    with torch.nn.attention.sdpa_kernel(SDPBackend.FLASH_ATTENTION):
        expected_out = torch.nn.functional.scaled_dot_product_attention(
            q, k, v, dropout_p=0.0, is_causal=True, scale=None
        )
        expected_out.backward(out_grad)
        expected_q_grad, expected_k_grad, expected_v_grad = q.grad, k.grad, v.grad

    with FusionDefinition() as fd:


            logsumexp,
            dropout_p,
            is_causal,
            seed,
            offset,
            scale=None,
        )

        fd.add_output(attn)
        for grad in [q_grad, k_grad, v_grad]:
            fd.add_output(grad)

    def _multidevice_schedule(fd: FusionDefinition) -> None:
        for t in fd.fusion.inputs():
            t.set_device_mesh(mesh)
            t.outer_split(1, d)
            t.axis(1).parallelize(nvfuser.ParallelType.mesh_x)

    b, s = 2, 1024

    def make_unsharded_tensor() -> torch.Tensor:
        return torch.randn(b, h, s, e // h, dtype=torch.bfloat16, device="cpu")

    q, k, v = [make_unsharded_tensor().requires_grad_() for _ in range(3)]
    out_grad = make_unsharded_tensor()
    sharded_q, sharded_k, sharded_v, sharded_out_grad = [
        multidevice_test.shard_tensor_1d(t, 1, mesh) for t in [q, k, v, out_grad]
    ]

    with torch.nn.attention.sdpa_kernel(SDPBackend.FLASH_ATTENTION):
        expected_out = torch.nn.functional.scaled_dot_product_attention(
            q, k, v, dropout_p=0.0, is_causal=True, scale=None
        )
        expected_out.backward(out_grad)
        expected_q_grad, expected_k_grad, expected_v_grad = q.grad, k.grad, v.grad

    with FusionDefinition() as fd:
        _definition(fd, qkv_format)
        _multidevice_schedule(fd)

    def reformat_tensor(t: torch.Tensor) -> torch.Tensor:
        match qkv_format:


        for t in fd.fusion.inputs():
            t.set_device_mesh(mesh)
            t.axis(0).parallelize(nvfuser.ParallelType.mesh_x)

    def make_unsharded_tensor() -> torch.Tensor:
        return torch.randn(b, h, s, e // h, dtype=torch.bfloat16, device="cuda")

    q, k, v = [make_unsharded_tensor().requires_grad_() for _ in range(3)]
    out_grad = make_unsharded_tensor()

    with torch.nn.attention.sdpa_kernel(SDPBackend.FLASH_ATTENTION):
        expected_out = torch.nn.functional.scaled_dot_product_attention(
            q, k, v, dropout_p=0.0, is_causal=True, scale=None
        )
        expected_out.backward(out_grad)
        expected_q_grad, expected_k_grad, expected_v_grad = q.grad, k.grad, v.grad

    rank = multidevice_test.rank

    # Head-parallelize Q, K, V or the attention output of an SDPA.
    def head_parallelize(t: torch.Tensor) -> torch.Tensor:
        assert t.shape == torch.Size([b, h, s, e // h])
        # The input `t` may require gradient. We .detach() so the new `t`
        # doesn't require gradient.
        t = t.detach().view([b, d, h // d, s, e // h]).narrow(1, rank, 1)
        match qkv_format:
            case QkvFormat.BHSE:
                return t.transpose(0, 1).contiguous()
            case QkvFormat.BSHE:
                return t.permute(1, 0, 3, 2, 4).contiguous().transpose(2, 3)

    with FusionDefinition() as fd:
        _definition(fd, qkv_format)
        _multidevice_schedule(fd)

    out, q_grad, k_grad, v_grad = fd.execute(
        [
            head_parallelize(q).requires_grad_(),
            head_parallelize(k).requires_grad_(),
            head_parallelize(v).requires_grad_(),
            head_parallelize(out_grad),
        ]
    )

    def assert_close(actual: torch.Tensor, expected: torch.Tensor) -> None:
        match qkv_format:
            case QkvFormat.BHSE:
                assert actual.is_contiguous()
            case QkvFormat.BSHE:
                assert actual.transpose(2, 3).is_contiguous()


    with FusionDefinition() as fd:
        _definition(fd, qkv_format)
        _multidevice_schedule(fd)

    def reformat_tensor(t: torch.Tensor) -> torch.Tensor:
        match qkv_format:
            case QkvFormat.BHSE:
                return t
            case QkvFormat.BSHE:
                return t.transpose(1, 2).contiguous().transpose(1, 2)

    attn, q_grad, k_grad, v_grad = fd.execute(
        [
            reformat_tensor(sharded_q).requires_grad_(),
            reformat_tensor(sharded_k).requires_grad_(),
            reformat_tensor(sharded_v).requires_grad_(),
            reformat_tensor(sharded_out_grad),
        ]
    )

    def assert_close(actual, expected):
        match qkv_format:
            case QkvFormat.BHSE:
                assert actual.is_contiguous()
            case QkvFormat.BSHE:
                assert actual.transpose(1, 2).is_contiguous()

        # Use the default rtol for bfloat16 and a relaxed atol.
        torch.testing.assert_close(
            actual,
            multidevice_test.shard_tensor_1d(expected, 1, mesh),
            rtol=1.6e-2,
            atol=1e-2,
        )

    for actual, expected in zip(
        [attn, q_grad, k_grad, v_grad],
        [expected_out, expected_q_grad, expected_k_grad, expected_v_grad],
    ):
        assert_close(actual, expected)

@pytest.mark.mpi
def test_privatize_squeeze(multidevice_test):
    d = multidevice_test.size
    mesh = nvfuser.multidevice.DeviceMesh(torch.arange(d))

    with FusionDefinition() as fd:
        inp_tv = fd.define_tensor((-1, -1, -1), dtype=DataType.BFloat16)
        tv1 = fd.ops.broadcast(inp_tv, [True, False, False, False])
        tv2 = fd.ops.squeeze(tv1, dims=[0])
        tv3 = fd.ops.cast(tv2, dtype=DataType.Float)
        tv4 = fd.ops.sum(tv3, dims=[0])
        tv5 = fd.ops.sum(tv3, dims=[1])


    def make_unsharded_tensor() -> torch.Tensor:
        return torch.randn(b, h, s, e // h, dtype=torch.bfloat16, device="cpu")

    q, k, v = [make_unsharded_tensor().requires_grad_() for _ in range(3)]
    out_grad = make_unsharded_tensor()
    sharded_q, sharded_k, sharded_v, sharded_out_grad = [
        multidevice_test.shard_tensor_1d(t, 1, mesh) for t in [q, k, v, out_grad]
    ]

    with torch.nn.attention.sdpa_kernel(SDPBackend.FLASH_ATTENTION):
        expected_out = torch.nn.functional.scaled_dot_product_attention(
            q, k, v, dropout_p=0.0, is_causal=True, scale=None
        )
        expected_out.backward(out_grad)
        expected_q_grad, expected_k_grad, expected_v_grad = q.grad, k.grad, v.grad

    with FusionDefinition() as fd:
        _definition(fd, qkv_format)
        _multidevice_schedule(fd)

    def reformat_tensor(t: torch.Tensor) -> torch.Tensor:
        match qkv_format:
            case QkvFormat.BHSE:
                return t
            case QkvFormat.BSHE:
                return t.transpose(1, 2).contiguous().transpose(1, 2)

    attn, q_grad, k_grad, v_grad = fd.execute(
        [
            reformat_tensor(sharded_q).requires_grad_(),
            reformat_tensor(sharded_k).requires_grad_(),
            reformat_tensor(sharded_v).requires_grad_(),
            reformat_tensor(sharded_out_grad),
        ]
    )

    def assert_close(actual, expected):
        match qkv_format:
            case QkvFormat.BHSE:
                assert actual.is_contiguous()
            case QkvFormat.BSHE:
                assert actual.transpose(1, 2).is_contiguous()

        # Use the default rtol for bfloat16 and a relaxed atol.
        torch.testing.assert_close(
            actual,


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import torch

FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_EPS = torch.finfo(torch.float8_e4m3fn).tiny
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max

# restore swizzled on block scaling factor:
# 1. restore swizzle
# 2. removes padding via slicing to [:mn, :k]
def swizzled_to_linear_128_4(a_sf_swizzled: torch.Tensor, mn, k):
    mn_padded, sf_k_padded = a_sf_swizzled.shape
    m_tiles = mn_padded // 128
    k_tiles = sf_k_padded // 4
    tmp = torch.reshape(a_sf_swizzled, (m_tiles, k_tiles, 32, 4, 4))
    return tmp.transpose(1, 3).reshape(mn_padded, sf_k_padded)[:mn, :k]

# apply swizzled on block scaling factor:
# 1. apply padding to [mn_t * 128 , k_t * 4]
# 2. apply swizzle
def linear_to_swizzled_128_4(a_sf_linear: torch.Tensor):
    mn, sf_k = a_sf_linear.shape
    m_tiles = (mn + 128 - 1) // 128
    mn_padded = m_tiles * 128
    k_tiles = (sf_k + 4 - 1) // 4
    k_padded = k_tiles * 4
    if mn_padded != mn or k_padded != sf_k:
        a_sf_padded = torch.empty(
            mn_padded, k_padded, dtype=a_sf_linear.dtype, device=a_sf_linear.device
        )
        a_sf_padded[0:mn, 0:sf_k] = a_sf_linear
    else:
        a_sf_padded = a_sf_linear
    # details about layout requirement on block-wise scaling factor
    # https://docs.nvidia.com/cutlass/media/docs/cpp/blackwell_functionality.html#scale-factor-layouts
    tmp = torch.reshape(a_sf_padded, (m_tiles, 4, 32, k_tiles, 4))
    return tmp.transpose(1, 3).reshape(mn_padded, k_padded)

def unpack_fp4(x: torch.Tensor) -> torch.Tensor:
    repeated = x.repeat_interleave(2, dim=1)
    repeated[:, 0::2] &= 0x0F
    repeated[:, 1::2] >>= 4
    return repeated

_FP4_LUT = torch.tensor(
    [
        0.0,  # 0: 0000 - zero
        0.5,  # 1: 0001 - smallest positive normal
        1.0,  # 2: 0010
        1.5,  # 3: 0011
        2.0,  # 4: 0100
        3.0,  # 5: 0101
        4.0,  # 6: 0110
        6.0,  # 7: 0111 - largest positive normal
        -0.0,  # 8: 1000 - negative zero
        -0.5,  # 9: 1001 - smallest negative normal
        -1.0,  # 10: 1010
        -1.5,  # 11: 1011
        -2.0,  # 12: 1100
        -3.0,  # 13: 1101
        -4.0,  # 14: 1110
        -6.0,  # 15: 1111 - largest negative normal
    ],
    dtype=torch.float32,
)

def fp4_to_fp32(fp4: torch.Tensor) -> torch.Tensor:
    # Convert FP4 indices to their corresponding floating point values
    # Each index (0-15) represents a 4-bit FP4 value in E2M1 format
    # Values based on the FP4 E2M1 specification
    fp4_lut = _FP4_LUT.to(fp4.device)
    return fp4_lut[fp4.to(torch.long)]

def dequantize_fp4(
    qx: torch.Tensor, sx: torch.Tensor, amax: torch.Tensor
) -> torch.Tensor:
    sf = sx.repeat_interleave(16, dim=1).view(torch.float8_e4m3fn).to(torch.float32)
    dqx = fp4_to_fp32(unpack_fp4(qx))
    sf = sf[: dqx.shape[0], : dqx.shape[1]]
    dequant = dqx * sf * (amax / (6.0 * 448))
    return dequant

def dequantize_to_dtype(
    tensor_fp4, tensor_sf, global_scale, dtype, device, block_size=16
):
    """Dequantize the fp4 tensor back to high precision."""
    # Two fp4 values are packed into one uint8.
    m, packed_k = tensor_fp4.shape
    k = packed_k * 2
    tensor_sf = tensor_sf.view(torch.float8_e4m3fn)
    tensor_sf = swizzled_to_linear_128_4(tensor_sf, m, k)
    out = dequantize_fp4(
        tensor_fp4.view(torch.uint8), tensor_sf, (6.0 * 448.0) / global_scale
    )
    return out.reshape(m, k)

# NOTE: This is from pytorch nvfp4 gemm tests.
def to_fp4(x):
    def down_size(size):
        assert size[-1] % 2 == 0, f"{size} last dim not divisible by two"
        return (*size[:-1], size[-1] // 2)

    def pack_uint4(uint8_data) -> torch.Tensor:
        # converting to uint8 for operations
        shape = uint8_data.shape
        assert shape[-1] % 2 == 0
        uint8_data = uint8_data.contiguous().view(-1)
        return (uint8_data[1::2] << 4 | uint8_data[::2]).view(down_size(shape))

    from torch.testing._internal.common_quantized import _f32_to_floatx_unpacked

    x = _f32_to_floatx_unpacked(x.float(), ebits=2, mbits=1)
    x = pack_uint4(x)
    x = x.view(torch.float4_e2m1fn_x2)
    return x

def pytorch_nvfp4_quantize(a, a_global_scale):
    BLOCK_SIZE = 16
    assert (
        a.size(-1) % BLOCK_SIZE == 0
    ), "The inner-most dim must be divisible by block_size; Padding is not implemented."
    assert a.is_contiguous(), "Only contiguous tensors are supported."

    original_shape = a.shape
    a_fp32 = a.float().reshape(original_shape[0], -1, BLOCK_SIZE)

    # Find absolute maximum along blockwise dimension
    max_abs = torch.amax(torch.abs(a_fp32), dim=-1)
    block_scale_fp32 = (max_abs / FLOAT4_E2M1_MAX).float()

    scaled_block_scale_fp32 = block_scale_fp32 * a_global_scale
    scaled_block_scale_fp32 = torch.clamp(
        scaled_block_scale_fp32, min=FLOAT8_E4M3_EPS, max=FLOAT8_E4M3_MAX
    )
    scaled_block_scale_fp8 = scaled_block_scale_fp32.to(torch.float8_e4m3fn)
    total_scale = scaled_block_scale_fp32 / a_global_scale
    a_scaled = a_fp32 / total_scale.unsqueeze(-1)
    a_scaled = torch.clamp(a_scaled, -FLOAT4_E2M1_MAX, FLOAT4_E2M1_MAX)
    a_scaled = a_scaled.view(original_shape)
    return to_fp4(a_scaled), scaled_block_scale_fp8

def round_up(x, y):
    return (x + y - 1) // y * y

def activation_scale_to_nvfp4(x, g_sf, offsets, blockscale_offsets, block_size):
    m = x.size(0)
    k = x.size(1)
    g = g_sf.size(0)
    padded_m_size = blockscale_offsets[g - 1] + round_up(m - offsets[g - 1], 128)
    block_scale = torch.empty(
        (padded_m_size, k // block_size), dtype=torch.float8_e4m3fn, device="cuda:0"
    )
    v_scaled = torch.empty((m, k // 2), dtype=torch.float4_e2m1fn_x2, device="cuda:0")
    for i in range(len(g_sf)):
        l = offsets[i]
        if i == g - 1:
            r = m
        else:
            r = offsets[i + 1]
        l_sf = blockscale_offsets[i]
        r_sf = l_sf + (r - l + 127) // 128 * 128
        v, b_sf = pytorch_nvfp4_quantize(x[l:r], g_sf[i])
        v_scaled[l:r] = v
        block_scale[l_sf:r_sf] = linear_to_swizzled_128_4(b_sf)

    return v_scaled, block_scale


# SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Owner(s): ["module: nvfuser"]

import torch

FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_EPS = torch.finfo(torch.float8_e4m3fn).tiny
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max

# restore swizzled on block scaling factor:
# 1. restore swizzle
# 2. removes padding via slicing to [:mn, :k]
def swizzled_to_linear_128_4(a_sf_swizzled: torch.Tensor, mn, k):
    mn_padded, sf_k_padded = a_sf_swizzled.shape
    m_tiles = mn_padded // 128
    k_tiles = sf_k_padded // 4
    tmp = torch.reshape(a_sf_swizzled, (m_tiles, k_tiles, 32, 4, 4))
    return tmp.transpose(1, 3).reshape(mn_padded, sf_k_padded)[:mn, :k]

# apply swizzled on block scaling factor:
# 1. apply padding to [mn_t * 128 , k_t * 4]
# 2. apply swizzle
def linear_to_swizzled_128_4(a_sf_linear: torch.Tensor):
    mn, sf_k = a_sf_linear.shape
    m_tiles = (mn + 128 - 1) // 128
    mn_padded = m_tiles * 128
    k_tiles = (sf_k + 4 - 1) // 4
    k_padded = k_tiles * 4
    if mn_padded != mn or k_padded != sf_k:
        a_sf_padded = torch.empty(
            mn_padded, k_padded, dtype=a_sf_linear.dtype, device=a_sf_linear.device
        )
        a_sf_padded[0:mn, 0:sf_k] = a_sf_linear
    else:
        a_sf_padded = a_sf_linear
    # details about layout requirement on block-wise scaling factor
    # https://docs.nvidia.com/cutlass/media/docs/cpp/blackwell_functionality.html#scale-factor-layouts


FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_EPS = torch.finfo(torch.float8_e4m3fn).tiny
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max

# restore swizzled on block scaling factor:
# 1. restore swizzle
# 2. removes padding via slicing to [:mn, :k]
def swizzled_to_linear_128_4(a_sf_swizzled: torch.Tensor, mn, k):
    mn_padded, sf_k_padded = a_sf_swizzled.shape
    m_tiles = mn_padded // 128
    k_tiles = sf_k_padded // 4
    tmp = torch.reshape(a_sf_swizzled, (m_tiles, k_tiles, 32, 4, 4))
    return tmp.transpose(1, 3).reshape(mn_padded, sf_k_padded)[:mn, :k]

# apply swizzled on block scaling factor:
# 1. apply padding to [mn_t * 128 , k_t * 4]
# 2. apply swizzle
def linear_to_swizzled_128_4(a_sf_linear: torch.Tensor):
    mn, sf_k = a_sf_linear.shape
    m_tiles = (mn + 128 - 1) // 128
    mn_padded = m_tiles * 128
    k_tiles = (sf_k + 4 - 1) // 4
    k_padded = k_tiles * 4
    if mn_padded != mn or k_padded != sf_k:
        a_sf_padded = torch.empty(
            mn_padded, k_padded, dtype=a_sf_linear.dtype, device=a_sf_linear.device
        )
        a_sf_padded[0:mn, 0:sf_k] = a_sf_linear
    else:
        a_sf_padded = a_sf_linear
    # details about layout requirement on block-wise scaling factor
    # https://docs.nvidia.com/cutlass/media/docs/cpp/blackwell_functionality.html#scale-factor-layouts
    tmp = torch.reshape(a_sf_padded, (m_tiles, 4, 32, k_tiles, 4))
    return tmp.transpose(1, 3).reshape(mn_padded, k_padded)

def unpack_fp4(x: torch.Tensor) -> torch.Tensor:
    repeated = x.repeat_interleave(2, dim=1)
    repeated[:, 0::2] &= 0x0F
    repeated[:, 1::2] >>= 4
    return repeated

_FP4_LUT = torch.tensor(
    [
        0.0,  # 0: 0000 - zero
        0.5,  # 1: 0001 - smallest positive normal
        1.0,  # 2: 0010
        1.5,  # 3: 0011
        2.0,  # 4: 0100
        3.0,  # 5: 0101
        4.0,  # 6: 0110
        6.0,  # 7: 0111 - largest positive normal
        -0.0,  # 8: 1000 - negative zero