pytorch - 💡(How to fix) Fix torch.compile doesn't use async all_to_all

def matmul_and_reshard_heads(*, xs, ws, device_mesh):
    num_heads = 16
    ys = []
    for x, w in zip(xs, ws):
        y = torch.matmul(x, w)
        batch_size, num_tokens, out_size = y.shape
        y = y.view(batch_size, num_tokens, num_heads, out_size // num_heads)
        y = DTensor.from_local(
            y, device_mesh=device_mesh, placements=(Shard(1),)
        )
        y = y.redistribute(
            # I tried async_op=True and async_op=False, both settings produce the same trace
            device_mesh, placements=(Shard(-2),), async_op=True
        )
        y = y.to_local()
        ys.append(y)
    return ys

# Run this function on 4 GPUs
def repro():
    device = torch.device("cuda")
    local_rank = int(os.environ.get("LOCAL_RANK", -1))
    global_rank = dist.get_rank()
    num_gpus = dist.get_world_size()

    mesh = dist.init_device_mesh(
        "cuda",
        (1, 4),
        mesh_dim_names=("dp", "sp"),
    )

    xs = [torch.randn((16, 16192, 4096), dtype=torch.bfloat16, device="cuda") for _ in range(8)]
    ws = [torch.randn((16, 4096, 8192), dtype=torch.bfloat16, device="cuda") for _ in range(8)]
    
    matmul_and_reshard_heads_compiled = torch.compile(matmul_and_reshard_heads, fullgraph=True)

    for i in range(100):
        ys = matmul_and_reshard_heads_compiled(xs=xs, ws=ws, device_mesh=mesh["sp"])
        [y.sum() for y in ys]
    torch.cuda.synchronize()