b6ff8a2f50
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
93 lines
3.0 KiB
Python
93 lines
3.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import torch
|
|
|
|
from vllm.distributed.parallel_state import (
|
|
get_pcp_group,
|
|
get_tp_group,
|
|
)
|
|
|
|
|
|
def _gather_prefill_cache_inputs(
|
|
tensors: tuple[torch.Tensor, ...],
|
|
slot_mapping: torch.Tensor,
|
|
num_decode_tokens: int,
|
|
) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
|
|
"""Keep replicated decode writes local and gather partitioned prefills."""
|
|
local_num_tokens = tensors[0].shape[0]
|
|
assert all(tensor.shape[0] == local_num_tokens for tensor in tensors)
|
|
assert 0 <= num_decode_tokens <= local_num_tokens
|
|
|
|
if num_decode_tokens == local_num_tokens:
|
|
return tensors, slot_mapping[:num_decode_tokens]
|
|
|
|
pcp_group = get_pcp_group()
|
|
gathered_prefills = tuple(
|
|
pcp_group.all_gather(tensor[num_decode_tokens:].contiguous(), dim=0)
|
|
for tensor in tensors
|
|
)
|
|
pcp_size = pcp_group.world_size
|
|
gathered_slot_mapping = slot_mapping[: pcp_size * local_num_tokens]
|
|
if num_decode_tokens == 0:
|
|
return gathered_prefills, gathered_slot_mapping
|
|
|
|
cache_inputs = tuple(
|
|
torch.cat((tensor[:num_decode_tokens], gathered_prefill), dim=0)
|
|
for tensor, gathered_prefill in zip(tensors, gathered_prefills)
|
|
)
|
|
rank_slot_mappings = gathered_slot_mapping.view(pcp_size, local_num_tokens)
|
|
cache_slot_mapping = torch.cat(
|
|
(
|
|
rank_slot_mappings[0, :num_decode_tokens],
|
|
rank_slot_mappings[:, num_decode_tokens:].flatten(),
|
|
)
|
|
)
|
|
return cache_inputs, cache_slot_mapping
|
|
|
|
|
|
def maybe_gather_mla_latent_cache_inputs(
|
|
kv_c_normed: torch.Tensor,
|
|
k_pe: torch.Tensor,
|
|
slot_mapping: torch.Tensor | None,
|
|
num_decode_tokens: int | None,
|
|
use_pcp: bool,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
|
|
if not use_pcp or num_decode_tokens is None:
|
|
return kv_c_normed, k_pe, slot_mapping
|
|
assert slot_mapping is not None
|
|
num_tokens = kv_c_normed.shape[0]
|
|
k_pe_flat = k_pe.reshape(num_tokens, -1)
|
|
(cache_kv_c, cache_k_pe_flat), cache_slot_mapping = _gather_prefill_cache_inputs(
|
|
(kv_c_normed, k_pe_flat),
|
|
slot_mapping,
|
|
num_decode_tokens,
|
|
)
|
|
cache_k_pe = cache_k_pe_flat.view(-1, *k_pe.shape[1:])
|
|
return cache_kv_c, cache_k_pe, cache_slot_mapping
|
|
|
|
|
|
def maybe_gather_indexer_k(
|
|
k: torch.Tensor,
|
|
slot_mapping: torch.Tensor,
|
|
num_decode_tokens: int,
|
|
use_pcp: bool,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
if not use_pcp:
|
|
return k, slot_mapping
|
|
(cache_k,), cache_slot_mapping = _gather_prefill_cache_inputs(
|
|
(k,), slot_mapping, num_decode_tokens
|
|
)
|
|
return cache_k, cache_slot_mapping
|
|
|
|
|
|
def finalize_mla_pcp_decode(
|
|
output: torch.Tensor,
|
|
num_heads: int,
|
|
) -> torch.Tensor:
|
|
if output.shape[1] < num_heads:
|
|
output = get_pcp_group().all_gather(output, dim=1)
|
|
elif output.shape[1] > num_heads:
|
|
head_start = get_tp_group().rank_in_group * num_heads
|
|
output = output[:, head_start : head_start + num_heads]
|
|
return output
|