flashinfer.comm.Mapping

class flashinfer.comm.Mapping(world_size=1, rank=0, gpus_per_node=8, *, cp_size=1, cp_config=None, tp_size=1, pp_size=1, moe_cluster_size=-1, moe_tp_size=-1, moe_ep_size=-1, attn_tp_size=-1, attn_cp_size=-1, auto_parallel=False, enable_attention_dp=False)

一个包含 8 个 GPU 的节点，tp_size = 4，cp_size = 1，pp_size = 2

2 个 tp 组

• [0, 1, 2, 3]

• [4, 5, 6, 7]

4 个 pp 组

• [0, 4]

• [1, 5]

• [2, 6]

• [3, 7]

一个包含 8 个 GPU 的节点，tp_size = 4，cp_size = 2，pp_size = 1

2 个 tp 组

• [0, 1, 2, 3]

• [4, 5, 6, 7]

4 个 cp 组

• [0, 4]

• [1, 5]

• [2, 6]

• [3, 7]

一个包含 8 个 GPU 的节点，moe_tp_size = 2，moe_ep_size = 4

4 个 moe_tp 组

• [0, 4]

• [1, 5]

• [2, 6]

• [3, 7]

2 个 moe_ep 组

• [0, 1, 2, 3]

• [4, 5, 6, 7]

2 个节点，包含 16 个 GPU，moe_tp_size = 2，moe_ep_size = 4，pp_size = 2

8 个 moe_tp 组

• [0 4]

• [1 5]

• [2 6]

• [3 7]

• [8 12]

• [9 13]

• [10 14]

• [11 15]

4 个 moe_ep 组

• [0, 1, 2, 3]

• [4, 5, 6, 7]

• [8, 9, 10, 11]

• [12, 13, 14, 15]

8 个 pp 组

• [0 8]

• [1 9]

• [2 10]

• [3 11]

• [4 12]

• [5 13]

• [6 14]

• [7 15]

2 个节点，包含 8 个 GPU，tp_size 2，pp_size 2，cp_size 2

4 个 tp 组：- [0, 1] - [2, 3] - [4, 5] - [6, 7]

4 个 pp 组：- [0, 4] - [1, 5] - [2, 6] - [3, 7]

4 个 cp 组：- [0, 2] - [1, 3] - [4, 6] - [5, 7]

__init__(world_size=1, rank=0, gpus_per_node=8, *, cp_size=1, cp_config=None, tp_size=1, pp_size=1, moe_cluster_size=-1, moe_tp_size=-1, moe_ep_size=-1, attn_tp_size=-1, attn_cp_size=-1, auto_parallel=False, enable_attention_dp=False)

方法

__init__([world_size, rank, gpus_per_node, ...])
ep_experts(num_experts)
from_dict(mapping)
get_local_rank(rank)
get_node_rank(rank)
has_cp()
has_moe_cluster()
has_moe_ep()
has_moe_tp()
has_pp()
has_tp()
is_first_pp_rank()
is_last_pp_rank()
is_multi_node()
is_second_last_pp_rank()
next_pp_rank()
pp_layers(num_layers)
prev_pp_rank()
to_dict()

属性

cp_group
cp_rank
local_rank
moe_cluster_group
moe_cluster_rank
moe_ep_group
moe_ep_rank
moe_tp_group
moe_tp_rank
node_rank
pp_group
pp_rank
rank
tp_group
tp_rank