| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import torch |
| 4 | |
| 5 | from vllm.distributed.parallel_state import ( |
| 6 | get_pcp_group, |
| 7 | get_tp_group, |
| 8 | ) |
| 9 | |
| 10 | |
| 11 | def _gather_prefill_cache_inputs( |
| 12 | tensors: tuple[torch.Tensor, ...], |
| 13 | slot_mapping: torch.Tensor, |
| 14 | num_decode_tokens: int, |
| 15 | shard_decode_requests: bool = False, |
| 16 | ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]: |
| 17 | """Gather PCP cache inputs while preserving replicated KV-cache state. |
| 18 | |
| 19 | PCP-only execution shards decode requests across ranks. In that mode each |
| 20 | rank must gather the other owners' decode KV as well as partitioned prefill |
| 21 | KV so every PCP rank retains a complete cache replica. DCP execution keeps |
| 22 | decode requests replicated and uses the legacy prefill-only gather. |
| 23 | """ |
| 24 | local_num_tokens = tensors[0].shape[0] |
| 25 | assert all(tensor.shape[0] == local_num_tokens for tensor in tensors) |
| 26 | assert 0 <= num_decode_tokens <= local_num_tokens |
| 27 | |
| 28 | # Replicated draft decodes use unexpanded slot mappings, even with DP padding. |
| 29 | if slot_mapping.shape[0] <= local_num_tokens: |
| 30 | return ( |
| 31 | tuple(tensor[:num_decode_tokens] for tensor in tensors), |
| 32 | slot_mapping[:num_decode_tokens], |
| 33 | ) |
| 34 | |
| 35 | pcp_group = get_pcp_group() |
| 36 | pcp_size = pcp_group.world_size |
| 37 | gathered_slot_mapping = slot_mapping[: pcp_size * local_num_tokens] |
| 38 | if shard_decode_requests: |
| 39 | gathered_inputs = tuple( |
| 40 | pcp_group.all_gather(tensor.contiguous(), dim=0) for tensor in tensors |
| 41 | ) |
| 42 | return gathered_inputs, gathered_slot_mapping |
| 43 | |
| 44 | if num_decode_tokens == local_num_tokens: |
| 45 | return tensors, slot_mapping[:num_decode_tokens] |
| 46 | |
| 47 | gathered_prefills = tuple( |
| 48 | pcp_group.all_gather(tensor[num_decode_tokens:].contiguous(), dim=0) |
| 49 | for tensor in tensors |
| 50 | ) |
| 51 | if num_decode_tokens == 0: |
| 52 | return gathered_prefills, gathered_slot_mapping |
| 53 | |
| 54 | cache_inputs = tuple( |
| 55 | torch.cat((tensor[:num_decode_tokens], gathered_prefill), dim=0) |
| 56 | for tensor, gathered_prefill in zip(tensors, gathered_prefills) |
| 57 | ) |
| 58 | rank_slot_mappings = gathered_slot_mapping.view(pcp_size, local_num_tokens) |
| 59 | cache_slot_mapping = torch.cat( |
| 60 | ( |
| 61 | rank_slot_mappings[0, :num_decode_tokens], |
| 62 | rank_slot_mappings[:, num_decode_tokens:].flatten(), |
| 63 | ) |
| 64 | ) |
| 65 | return cache_inputs, cache_slot_mapping |
| 66 | |
| 67 | |
| 68 | def maybe_gather_mla_latent_cache_inputs( |
| 69 | kv_c_normed: torch.Tensor, |
| 70 | k_pe: torch.Tensor, |
| 71 | slot_mapping: torch.Tensor | None, |
| 72 | num_decode_tokens: int | None, |
| 73 | use_pcp: bool, |
| 74 | pcp_shard_decode_requests: bool = False, |
| 75 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: |
| 76 | if not use_pcp or num_decode_tokens is None: |
| 77 | return kv_c_normed, k_pe, slot_mapping |
| 78 | assert slot_mapping is not None |
| 79 | num_tokens = kv_c_normed.shape[0] |
| 80 | k_pe_flat = k_pe.reshape(num_tokens, -1) |
| 81 | (cache_kv_c, cache_k_pe_flat), cache_slot_mapping = _gather_prefill_cache_inputs( |
| 82 | (kv_c_normed, k_pe_flat), |
| 83 | slot_mapping, |
| 84 | num_decode_tokens, |
| 85 | pcp_shard_decode_requests, |
| 86 | ) |
| 87 | cache_k_pe = cache_k_pe_flat.view(-1, *k_pe.shape[1:]) |
| 88 | return cache_kv_c, cache_k_pe, cache_slot_mapping |
| 89 | |
| 90 | |
| 91 | def maybe_gather_indexer_k( |
| 92 | k: torch.Tensor, |
| 93 | slot_mapping: torch.Tensor, |
| 94 | num_decode_tokens: int, |
| 95 | use_pcp: bool, |
| 96 | pcp_shard_decode_requests: bool = False, |
| 97 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 98 | if not use_pcp: |
| 99 | return k, slot_mapping |
| 100 | (cache_k,), cache_slot_mapping = _gather_prefill_cache_inputs( |
| 101 | (k,), slot_mapping, num_decode_tokens, pcp_shard_decode_requests |
| 102 | ) |
| 103 | return cache_k, cache_slot_mapping |
| 104 | |
| 105 | |
| 106 | def finalize_mla_pcp_decode( |
| 107 | output: torch.Tensor, |
| 108 | num_heads: int, |
| 109 | ) -> torch.Tensor: |
| 110 | if output.shape[1] < num_heads: |
| 111 | output = get_pcp_group().all_gather(output, dim=1) |
| 112 | elif output.shape[1] > num_heads: |
| 113 | head_start = get_tp_group().rank_in_group * num_heads |
| 114 | output = output[:, head_start : head_start + num_heads] |
| 115 | return output |