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)¶

A node with 8 GPUs, tp_size = 4, cp_size = 1, pp_size = 2

2 tp groups:

[0, 1, 2, 3]
[4, 5, 6, 7]

4 pp groups:

[0, 4]
[1, 5]
[2, 6]
[3, 7]

A node with 8 GPUs, tp_size = 4, cp_size = 2, pp_size = 1

2 tp groups:

[0, 1, 2, 3]
[4, 5, 6, 7]

4 cp groups:

[0, 4]
[1, 5]
[2, 6]
[3, 7]

A node with 8 GPUs, moe_tp_size = 2, moe_ep_size = 4

4 moe_tp groups:

[0, 4]
[1, 5]
[2, 6]
[3, 7]

2 moe_ep groups:

[0, 1, 2, 3]
[4, 5, 6, 7]

2 nodes with 16 GPUs, moe_tp_size = 2, moe_ep_size = 4, pp_size = 2

8 moe_tp groups:

[0 4]
[1 5]
[2 6]
[3 7]
[8 12]
[9 13]
[10 14]
[11 15]

4 moe_ep groups:

[0, 1, 2, 3]
[4, 5, 6, 7]
[8, 9, 10, 11]
[12, 13, 14, 15]

8 pp groups:

[0 8]
[1 9]
[2 10]
[3 11]
[4 12]
[5 13]
[6 14]
[7 15]

2 nodes with 8 GPUs, tp_size 2, pp_size 2, cp_size 2

4 tp groups: - [0, 1] - [2, 3] - [4, 5] - [6, 7]

4 pp groups: - [0, 4] - [1, 5] - [2, 6] - [3, 7]

4 cp groups: - [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)¶

Methods

`__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`()

Attributes

`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`