[Feature] Hybrid EP and TP (#8590)

2025-07-31 02:53:25 -07:00
parent 51c38163c1
commit 7a1f7fc504
14 changed files with 142 additions and 39 deletions
@@ -56,6 +56,7 @@ class TpModelWorker:
        server_args: ServerArgs,
        gpu_id: int,
        tp_rank: int,
+        moe_ep_rank: int,
        pp_rank: int,
        dp_rank: Optional[int],
        nccl_port: int,
@@ -66,6 +67,7 @@ class TpModelWorker:
        # Parse args
        self.tp_size = server_args.tp_size
        self.tp_rank = tp_rank
+        self.moe_ep_rank = moe_ep_rank
        self.pp_rank = pp_rank

        # Init model and tokenizer
@@ -85,6 +87,8 @@ class TpModelWorker:
            gpu_id=gpu_id,
            tp_rank=tp_rank,
            tp_size=server_args.tp_size,
+            moe_ep_rank=moe_ep_rank,
+            moe_ep_size=server_args.ep_size,
            pp_rank=pp_rank,
            pp_size=server_args.pp_size,
            nccl_port=nccl_port,