format code

Author: JiaqingFu

examples/distributed/example_overlapping_allgather.py

@@ -1,22 +1,19 @@
-import argparse
 import torch
 import torch.distributed as dist
 import pynvshmem
 import tilelang
 import tilelang.language as T
 import os
-from tilelang.distributed.utils import init_distributed, dtype_map, perf_fn
-from tilelang.distributed.utils import init_dist
+from tilelang.distributed.utils import init_distributed
 from tilelang.env import env
 from packaging import version
 import importlib.metadata
 
 cuda_python_version = importlib.metadata.version("cuda-python")
 if version.parse(cuda_python_version) >= version.parse("12.8.0"):
-    from cuda.bindings import driver as cuda
     from cuda.bindings import runtime as cudart
 else:
-    from cuda import cuda, cudart
+    from cuda import cudart
 # NODES=2 NODE_RANK=0 ARNOLD_WORKER_0_HOST=ip0 bash tilelang/distributed/launch.sh ./examples/distributed/example_overlapping_allgather.py
 # NODES=2 NODE_RANK=1 ARNOLD_WORKER_0_HOST=ip0 bash tilelang/distributed/launch.sh ./examples/distributed/example_overlapping_allgather.py
 

tilelang/distributed/utils.py

@@ -82,7 +82,7 @@ def init_distributed(return_tp_group=False, init_nvshmem=True, return_lc_group=F
         base = (RANK // local_world_size) * local_world_size
         LC_GROUP = torch.distributed.new_group(
             list(range(base, base + local_world_size)), backend="nccl")
-        
+
         return WORLD_SIZE, RANK, LOCAL_RANK, TP_GROUP, LC_GROUP
     elif return_tp_group:
         return WORLD_SIZE, RANK, LOCAL_RANK, TP_GROUP