446 lines
18 KiB
Python
446 lines
18 KiB
Python
import math
|
||
import os
|
||
from collections import OrderedDict, defaultdict
|
||
from collections.abc import Iterator
|
||
from dataclasses import dataclass
|
||
from typing import Any
|
||
|
||
import torch
|
||
import torch.distributed as dist
|
||
from vllm.logger import logger
|
||
|
||
from vllm_ascend.distributed.parallel_state import get_p_tp_group
|
||
|
||
MAX_HCCL_REGISTER_REGIONS = 256
|
||
REGISTER_MERGE_GAP_BYTES = 4096
|
||
|
||
|
||
def kv_alltoall_and_rearrange(pd_tp_ratio: int, key: torch.Tensor, value: torch.TensorType):
|
||
if pd_tp_ratio <= 1:
|
||
return None, None
|
||
elif key is None or value is None:
|
||
raise ValueError("key or value is None")
|
||
k_output = alltoall_and_rearrange(pd_tp_ratio, key)
|
||
v_output = alltoall_and_rearrange(pd_tp_ratio, value)
|
||
return k_output, v_output
|
||
|
||
|
||
def alltoall_and_rearrange(tp_ratio: int, input_tensor: torch.Tensor):
|
||
num_kv_heads = input_tensor.size(1)
|
||
output_tensor = torch.zeros_like(input_tensor)
|
||
dist.all_to_all_single(output_tensor, input_tensor, group=get_p_tp_group().device_group)
|
||
input_tensor = 0
|
||
result = rearrange_output(output_tensor, tp_ratio, num_kv_heads)
|
||
output_tensor = 0
|
||
return result
|
||
|
||
|
||
def rearrange_output(base_output: torch.Tensor, cut_num: int, num_kv_heads: int):
|
||
size_0 = base_output.size(0)
|
||
if size_0 % cut_num != 0:
|
||
raise ValueError(f"The size of dim 0 [{size_0}] must be divisible by the cut_num [{cut_num}]")
|
||
chunk_size = size_0 // cut_num
|
||
reshaped = base_output.view(cut_num, chunk_size, -1)
|
||
transposed = reshaped.transpose(0, 1)
|
||
return transposed.contiguous().view(size_0, num_kv_heads, -1)
|
||
|
||
|
||
def align_memory(tensor: torch.Tensor, alignment: int) -> torch.Tensor:
|
||
data_ptr = tensor.data_ptr()
|
||
aligned_addr = (data_ptr + alignment - 1) // alignment * alignment
|
||
offset = (aligned_addr - data_ptr) // tensor.element_size()
|
||
return tensor[int(offset) :]
|
||
|
||
|
||
def get_transfer_timeout_value():
|
||
ascend_transfer_timeout = os.getenv("ASCEND_TRANSFER_TIMEOUT", "")
|
||
if len(ascend_transfer_timeout) > 0:
|
||
return int(ascend_transfer_timeout)
|
||
hccl_rdma_timeout = int(os.getenv("HCCL_RDMA_TIMEOUT", "20")) # type: ignore
|
||
hccl_rdma_retry_cnt = int(os.getenv("HCCL_RDMA_RETRY_CNT", "7")) # type: ignore
|
||
return int((4.096 * (2**hccl_rdma_timeout)) * hccl_rdma_retry_cnt // 1000 + 3000)
|
||
|
||
|
||
@dataclass
|
||
class parallel_info:
|
||
tp_size: int
|
||
pcp_size: int
|
||
dcp_size: int
|
||
use_mla: bool
|
||
pd_head_ratio: int
|
||
|
||
|
||
def get_cp_group(tp: int, heads: int, dcp: int):
|
||
# Partition the second dimension of [pcp][head_group][dcp] to obtain a complete head group
|
||
# head_group is all blocks for request in the same head
|
||
# tp8 dcp2 heads4 return[[0,1,2,3]]
|
||
# tp8 dcp1 heads4 return[[0,2,4,6],[1,3,5,7]]
|
||
step = tp // heads
|
||
if step == 0:
|
||
return [[i for i in range(tp // dcp)]]
|
||
else:
|
||
return [
|
||
set([k // dcp for h in range(heads) for k in range(h * step + i * dcp, h * step + (i + 1) * dcp)])
|
||
for i in range(step // dcp)
|
||
]
|
||
|
||
|
||
def context_parallel_parameters_check(
|
||
remote_pcp_size: int,
|
||
remote_dcp_size: int,
|
||
p_parallel_info: parallel_info,
|
||
d_parallel_info: parallel_info,
|
||
total_num_kv_heads: int,
|
||
):
|
||
# Check whether the pcp–dcp ratio is supported
|
||
assert (p_parallel_info.pcp_size * p_parallel_info.dcp_size) % (remote_pcp_size * remote_dcp_size) == 0
|
||
if not p_parallel_info.use_mla:
|
||
p_node_heads_per_rank = math.ceil(total_num_kv_heads / p_parallel_info.tp_size)
|
||
d_node_heads_per_rank = math.ceil(total_num_kv_heads / d_parallel_info.dcp_size)
|
||
assert d_node_heads_per_rank % p_node_heads_per_rank == 0
|
||
|
||
|
||
def get_tp_rank_head_mapping(num_key_value_heads: int, tp_size: int):
|
||
# Get the head_idx corresponding to the tp_rank, {tp_rank:[head_indx]}
|
||
mapping = {}
|
||
if tp_size <= num_key_value_heads:
|
||
if num_key_value_heads % tp_size != 0:
|
||
raise ValueError(f"Number of heads ({num_key_value_heads}) cannot be evenly divided by TP ({tp_size}).")
|
||
|
||
heads_per_rank = num_key_value_heads // tp_size
|
||
|
||
for rank in range(tp_size):
|
||
start_idx = rank * heads_per_rank
|
||
end_idx = start_idx + heads_per_rank
|
||
mapping[rank] = list(range(start_idx, end_idx))
|
||
else:
|
||
if tp_size % num_key_value_heads != 0:
|
||
raise ValueError(f"Number of heads ({num_key_value_heads}) cannot be evenly divided by TP ({tp_size}).")
|
||
ranks_per_head = tp_size // num_key_value_heads
|
||
for rank in range(tp_size):
|
||
head_idx = rank // ranks_per_head
|
||
mapping[rank] = [head_idx]
|
||
return mapping
|
||
|
||
|
||
def get_head_group_mapping(num_key_value_heads: int, tp_size: int, num_groups: int, select_cp_group: list[int]):
|
||
# Get the mapping dictionary, where the key is head_group_rank and the value is head_idx
|
||
if tp_size % num_groups != 0:
|
||
raise ValueError(
|
||
f"Total number of devices ({tp_size}) cannot be divided by the number of groups ({num_groups})."
|
||
)
|
||
ranks_per_group = tp_size // num_groups
|
||
tp_mapping = get_tp_rank_head_mapping(num_key_value_heads, tp_size)
|
||
group_mapping = {}
|
||
for group_rank in range(num_groups):
|
||
if group_rank in select_cp_group:
|
||
start_rank = group_rank * ranks_per_group
|
||
end_rank = start_rank + ranks_per_group
|
||
heads_set = set()
|
||
|
||
for rank in range(start_rank, end_rank):
|
||
heads_set.update(tp_mapping[rank])
|
||
group_mapping[group_rank] = sorted(list(heads_set))
|
||
return group_mapping
|
||
|
||
|
||
def get_local_remote_block_port_mappings(
|
||
to_trans_idx: int,
|
||
p_parallel_info: parallel_info,
|
||
d_parallel_info: parallel_info,
|
||
d_hosts: list[str],
|
||
d_port: int,
|
||
selected_p_cp_group: list[int],
|
||
selected_d_cp_group: list[int],
|
||
prompt_len: int,
|
||
block_size: int,
|
||
req_meta,
|
||
total_num_kv_heads: int,
|
||
req_id: str,
|
||
):
|
||
p_head_group_size = p_parallel_info.tp_size // p_parallel_info.dcp_size
|
||
d_head_group_size = d_parallel_info.tp_size // d_parallel_info.dcp_size
|
||
world_size = d_parallel_info.pcp_size * d_head_group_size * d_parallel_info.dcp_size
|
||
# Compute which logic_block_idx corresponds to each tp_rank
|
||
p_rank_block_mapping: list[list[list[list[int]]]] = [
|
||
[[[] for _ in range(p_parallel_info.dcp_size)] for _ in range(p_head_group_size)]
|
||
for _ in range(p_parallel_info.pcp_size)
|
||
]
|
||
for logic_block_idx in range(to_trans_idx):
|
||
pcp_rank = (logic_block_idx // p_parallel_info.dcp_size) % p_parallel_info.pcp_size
|
||
dcp_rank = logic_block_idx % p_parallel_info.dcp_size
|
||
for p_head_group_rank in range(p_head_group_size):
|
||
if p_head_group_rank in selected_p_cp_group:
|
||
p_rank_block_mapping[pcp_rank][p_head_group_rank][dcp_rank].append(logic_block_idx)
|
||
|
||
# Find the remote device that holds the logic_block_idx
|
||
d_block_rank_mapping: dict[int, dict[int, dict[str, Any]]] = defaultdict(lambda: defaultdict(dict))
|
||
for logic_block_idx in range(to_trans_idx):
|
||
pcp_rank = (logic_block_idx // d_parallel_info.dcp_size) % d_parallel_info.pcp_size
|
||
for d_head_group_rank in range(d_head_group_size):
|
||
if d_head_group_rank in selected_d_cp_group:
|
||
dcp_rank = logic_block_idx % d_parallel_info.dcp_size
|
||
world_rank = (
|
||
pcp_rank * d_head_group_size * d_parallel_info.dcp_size
|
||
+ d_head_group_rank * d_parallel_info.dcp_size
|
||
+ dcp_rank
|
||
)
|
||
world_size = d_parallel_info.pcp_size * d_head_group_size * d_parallel_info.dcp_size
|
||
host = d_hosts[(len(d_hosts) * world_rank) // world_size]
|
||
port = d_port + world_rank
|
||
block_idx = (logic_block_idx - (pcp_rank * d_parallel_info.pcp_size + dcp_rank)) // (
|
||
d_parallel_info.pcp_size * d_parallel_info.dcp_size
|
||
)
|
||
d_block_rank_mapping[logic_block_idx][d_head_group_rank] = {
|
||
"pcp_rank": pcp_rank,
|
||
"dcp_rank": dcp_rank,
|
||
"host": host,
|
||
"port": port,
|
||
"block_idx": block_idx,
|
||
}
|
||
# Get how many times each device should receive done_single for this request
|
||
d_trans_count_mapping = {}
|
||
trans_block_size = math.ceil(prompt_len / block_size) # Total number of blocks
|
||
transed_block_size = math.ceil(req_meta.remote_cache_tokens / block_size) # Number of prefix cache hit blocks
|
||
d_cp_size = d_parallel_info.pcp_size * d_parallel_info.dcp_size
|
||
for d_pcp_rank in range(d_parallel_info.pcp_size):
|
||
for d_head_group_rank in range(d_head_group_size):
|
||
for d_dcp_rank in range(d_parallel_info.dcp_size):
|
||
if trans_block_size >= (p_parallel_info.pcp_size * p_parallel_info.dcp_size):
|
||
trans_count = (p_parallel_info.pcp_size * p_parallel_info.dcp_size) // d_cp_size
|
||
else:
|
||
current_rank_idx = d_pcp_rank * d_parallel_info.dcp_size + d_dcp_rank
|
||
total_global_blocks = transed_block_size + trans_block_size
|
||
|
||
target_total_count = total_global_blocks // d_cp_size
|
||
if current_rank_idx < (total_global_blocks % d_cp_size):
|
||
target_total_count += 1
|
||
|
||
prev_processed_count = transed_block_size // d_cp_size
|
||
if current_rank_idx < (transed_block_size % d_cp_size):
|
||
prev_processed_count += 1
|
||
|
||
trans_count = target_total_count - prev_processed_count
|
||
world_rank = (
|
||
d_pcp_rank * d_head_group_size * d_parallel_info.dcp_size
|
||
+ d_head_group_rank * d_parallel_info.dcp_size
|
||
+ d_dcp_rank
|
||
)
|
||
host = d_hosts[(len(d_hosts) * world_rank) // world_size]
|
||
port = d_port + world_rank
|
||
d_trans_count_mapping[(host, port)] = trans_count * p_parallel_info.pd_head_ratio
|
||
|
||
# Compute the mapping between local and remote head_group_rank
|
||
p_tp_rank_head_mapping = get_head_group_mapping(
|
||
total_num_kv_heads, p_parallel_info.tp_size, p_head_group_size, selected_p_cp_group
|
||
)
|
||
d_tp_rank_head_mapping = get_head_group_mapping(
|
||
total_num_kv_heads, d_parallel_info.tp_size, d_head_group_size, selected_d_cp_group
|
||
)
|
||
head_to_d_groups = defaultdict(set)
|
||
for d_rank, heads in d_tp_rank_head_mapping.items():
|
||
for head in heads:
|
||
head_to_d_groups[head].add(d_rank)
|
||
pd_head_mapping = {}
|
||
for p_rank, p_heads in p_tp_rank_head_mapping.items():
|
||
target_d_ranks = set()
|
||
for head in p_heads:
|
||
if head in head_to_d_groups:
|
||
target_d_ranks.update(head_to_d_groups[head])
|
||
else:
|
||
logger.info("Warning: Head %s exists in P but not in D mapping.", head)
|
||
pd_head_mapping[p_rank] = sorted(list(target_d_ranks))
|
||
logger.debug(
|
||
"MooncakeLayerwiseConnector _get_kv_split_metadata req_id=%r "
|
||
"P-side logic_block to rank mapping: %s, "
|
||
"D-side logic_block to rank mapping: %s, "
|
||
"P&D head_group_rank mapping: %s",
|
||
req_id,
|
||
p_rank_block_mapping,
|
||
d_block_rank_mapping,
|
||
pd_head_mapping,
|
||
)
|
||
return p_rank_block_mapping, d_block_rank_mapping, pd_head_mapping, d_trans_count_mapping
|
||
|
||
|
||
def get_transfer_mappings(
|
||
p_rank_block_mapping: list[list[list[list[int]]]],
|
||
d_block_rank_mapping: dict[int, dict[int, dict[str, Any]]],
|
||
pd_head_mapping: dict[int, set],
|
||
d_trans_count_mapping: dict[tuple[str, int], int],
|
||
req_meta,
|
||
block_group_idx: int,
|
||
p_parallel_info: parallel_info,
|
||
req_id: str,
|
||
transed_idx: int,
|
||
to_trans_idx: int,
|
||
tp_rank: int,
|
||
pcp_rank: int,
|
||
dcp_rank: int,
|
||
):
|
||
transfer_mappings: dict[tuple[str, int], dict[str, Any]] = {}
|
||
p_head_group_rank = (tp_rank - dcp_rank) // p_parallel_info.dcp_size
|
||
p_block_idxs: list[int] = p_rank_block_mapping[pcp_rank][p_head_group_rank][dcp_rank]
|
||
p_block_ids = req_meta.local_block_ids[block_group_idx]
|
||
d_block_ids = req_meta.remote_block_ids[block_group_idx]
|
||
for p_block_idx, logic_block_idx in enumerate(p_block_idxs):
|
||
if logic_block_idx < transed_idx or logic_block_idx >= to_trans_idx:
|
||
continue
|
||
for d_head_group_rank in pd_head_mapping[p_head_group_rank]:
|
||
p_block_id = p_block_ids[p_block_idx]
|
||
remote_host = d_block_rank_mapping[logic_block_idx][d_head_group_rank]["host"]
|
||
remote_port = d_block_rank_mapping[logic_block_idx][d_head_group_rank]["port"]
|
||
d_block_idx = d_block_rank_mapping[logic_block_idx][d_head_group_rank]["block_idx"]
|
||
d_block_id = d_block_ids[d_block_idx]
|
||
if (remote_host, remote_port) not in transfer_mappings:
|
||
transfer_mappings[(remote_host, remote_port)] = {
|
||
"local_block_ids": [],
|
||
"remote_block_ids": [],
|
||
"trans_count": 0,
|
||
}
|
||
transfer_mappings[(remote_host, remote_port)]["local_block_ids"].append(p_block_id)
|
||
transfer_mappings[(remote_host, remote_port)]["remote_block_ids"].append(d_block_id)
|
||
for (host, port), block_dict in transfer_mappings.items():
|
||
block_dict["trans_count"] = d_trans_count_mapping[(host, port)]
|
||
logger.debug("MooncakeLayerwiseConnector Request %s transfer tasks: %s", req_id, transfer_mappings)
|
||
return transfer_mappings
|
||
|
||
|
||
@dataclass
|
||
class RegisterRange:
|
||
start: int
|
||
end: int
|
||
|
||
|
||
@dataclass
|
||
class RegisterRegions:
|
||
ptrs: list[int]
|
||
lengths: list[int]
|
||
logical_tensor_count: int | None = None
|
||
logical_total_bytes: int | None = None
|
||
|
||
@property
|
||
def registered_bytes(self) -> int:
|
||
return sum(self.lengths)
|
||
|
||
|
||
def iter_kv_cache_tensors(obj: Any) -> Iterator[torch.Tensor]:
|
||
"""Flatten kv_caches into tensors without materializing new tensors."""
|
||
if obj is None:
|
||
return
|
||
|
||
if isinstance(obj, torch.Tensor):
|
||
yield obj
|
||
return
|
||
|
||
if isinstance(obj, (tuple, list)):
|
||
for item in obj:
|
||
yield from iter_kv_cache_tensors(item)
|
||
return
|
||
|
||
if isinstance(obj, dict):
|
||
for item in obj.values():
|
||
yield from iter_kv_cache_tensors(item)
|
||
return
|
||
|
||
|
||
def tensor_storage_key(tensor: torch.Tensor) -> int:
|
||
"""Return a stable grouping key for tensors sharing the same storage.
|
||
|
||
Do NOT use this key as the register address directly. For aligned KV cache
|
||
views, tensor.untyped_storage().data_ptr() may point to the original raw
|
||
allocation, whose address can be unaligned. We only use it to group views.
|
||
"""
|
||
try:
|
||
return tensor.untyped_storage().data_ptr()
|
||
except Exception:
|
||
try:
|
||
return tensor.storage().data_ptr()
|
||
except Exception:
|
||
return tensor.data_ptr()
|
||
|
||
|
||
def collect_storage_merged_register_regions(
|
||
kv_caches: dict[str, Any],
|
||
) -> RegisterRegions:
|
||
"""Collect HCCL/Mooncake register regions with storage-aware merging.
|
||
|
||
Metadata should still use each logical tensor's own data_ptr().
|
||
register_buffer should use the merged memory ranges returned here.
|
||
"""
|
||
ranges_by_storage: OrderedDict[int, list[RegisterRange]] = OrderedDict()
|
||
logical_tensor_count = 0
|
||
logical_total_bytes = 0
|
||
|
||
for tensor in iter_kv_cache_tensors(kv_caches):
|
||
if tensor is None or tensor.numel() == 0:
|
||
continue
|
||
|
||
if not tensor.is_contiguous():
|
||
logger.warning(
|
||
"Mooncake register_buffer got a non-contiguous KV cache "
|
||
"tensor: shape=%s, dtype=%s, data_ptr=%s. "
|
||
"Registration will use logical numel * element_size.",
|
||
tuple(tensor.shape),
|
||
tensor.dtype,
|
||
hex(tensor.data_ptr()),
|
||
)
|
||
|
||
nbytes = tensor.nbytes
|
||
start = tensor.data_ptr()
|
||
end = start + nbytes
|
||
storage_key = tensor_storage_key(tensor)
|
||
|
||
logical_tensor_count += 1
|
||
logical_total_bytes += nbytes
|
||
|
||
ranges_by_storage.setdefault(storage_key, []).append(RegisterRange(start, end))
|
||
|
||
register_ptrs: list[int] = []
|
||
register_lengths: list[int] = []
|
||
|
||
for ranges in ranges_by_storage.values():
|
||
ranges.sort(key=lambda r: r.start)
|
||
|
||
merged_start = ranges[0].start
|
||
merged_end = ranges[0].end
|
||
|
||
for region in ranges[1:]:
|
||
if region.start <= merged_end + REGISTER_MERGE_GAP_BYTES:
|
||
merged_end = max(merged_end, region.end)
|
||
else:
|
||
register_ptrs.append(merged_start)
|
||
register_lengths.append(merged_end - merged_start)
|
||
merged_start = region.start
|
||
merged_end = region.end
|
||
|
||
register_ptrs.append(merged_start)
|
||
register_lengths.append(merged_end - merged_start)
|
||
|
||
return RegisterRegions(
|
||
ptrs=register_ptrs,
|
||
lengths=register_lengths,
|
||
logical_tensor_count=logical_tensor_count,
|
||
logical_total_bytes=logical_total_bytes,
|
||
)
|
||
|
||
|
||
def validate_register_region_count(regions: RegisterRegions) -> None:
|
||
region_count = len(regions.ptrs)
|
||
if region_count <= MAX_HCCL_REGISTER_REGIONS:
|
||
return
|
||
|
||
detail = f"registered_bytes={regions.registered_bytes}"
|
||
if regions.logical_tensor_count is not None:
|
||
detail += f", logical_tensors={regions.logical_tensor_count}, logical_bytes={regions.logical_total_bytes}"
|
||
|
||
raise RuntimeError(
|
||
"Mooncake register_buffer region count "
|
||
f"{region_count} exceeds HCCL per-process limit "
|
||
f"{MAX_HCCL_REGISTER_REGIONS}. "
|
||
"KV cache registration would fail. "
|
||
f"{detail}. "
|
||
"Please reduce KV cache allocation fragmentation or merge "
|
||
"k/v/dsa/scale allocations further."
|
||
)
|