| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import torch |
| 4 | |
| 5 | from vllm.triton_utils import tl, triton |
| 6 | |
| 7 | |
| 8 | def prepare_dcp_local_seq_lens( |
| 9 | dcp_local_seq_lens: torch.Tensor, |
| 10 | seq_lens: torch.Tensor, |
| 11 | num_reqs: int, |
| 12 | dcp_size: int, |
| 13 | dcp_rank: int, |
| 14 | cp_interleave: int, |
| 15 | *, |
| 16 | num_reqs_padded: int | None = None, |
| 17 | ) -> torch.Tensor | None: |
| 18 | """Populate caller-owned storage and return its padded view, or None without DCP.""" |
| 19 | assert dcp_size > 1 |
| 20 | |
| 21 | max_num_reqs = dcp_local_seq_lens.shape[0] |
| 22 | BLOCK_SIZE = 128 |
| 23 | num_blocks = triton.cdiv(max_num_reqs, BLOCK_SIZE) |
| 24 | _dcp_local_seq_lens_kernel[(num_blocks,)]( |
| 25 | dcp_local_seq_lens, |
| 26 | seq_lens, |
| 27 | dcp_size, |
| 28 | dcp_rank, |
| 29 | cp_interleave, |
| 30 | num_reqs, |
| 31 | max_num_reqs, |
| 32 | BLOCK_SIZE, |
| 33 | ) |
| 34 | return dcp_local_seq_lens[:num_reqs_padded] |
| 35 | |
| 36 | |
| 37 | @triton.jit |
| 38 | def _dcp_local_seq_lens_kernel( |
| 39 | out_ptr, |
| 40 | seq_lens_ptr, |
| 41 | dcp_size, |
| 42 | dcp_rank, |
| 43 | cp_interleave, |
| 44 | num_reqs, |
| 45 | max_num_reqs, |
| 46 | BLOCK_SIZE: tl.constexpr, |
| 47 | ): |
| 48 | pid = tl.program_id(0) |
| 49 | block = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 50 | |
| 51 | seq_lens = tl.load(seq_lens_ptr + block, mask=block < num_reqs) |
| 52 | |
| 53 | # Distribute KV cache among different ranks, in a round-robin manner. |
| 54 | rounds = seq_lens // (dcp_size * cp_interleave) |
| 55 | remainder = seq_lens % (dcp_size * cp_interleave) |
| 56 | |
| 57 | remainder = tl.maximum(remainder - dcp_rank * cp_interleave, 0) |
| 58 | remainder = tl.minimum(remainder, cp_interleave) |
| 59 | local_seq_lens = rounds * cp_interleave + remainder |
| 60 | |
| 61 | # For [num_reqs, max_num_reqs), pad with 0 |
| 62 | local_seq_lens = tl.where(block < num_reqs, local_seq_lens, 0) |
| 63 | tl.store(out_ptr + block, local_seq_lens, mask=block < max_num_reqs) |
| 64 | |
| 65 | |
| 66 | @triton.jit |
| 67 | def cp_local_slot( |
| 68 | positions, |
| 69 | block_numbers, |
| 70 | block_size, |
| 71 | cp_rank, |
| 72 | CP_SIZE: tl.constexpr, |
| 73 | CP_INTERLEAVE: tl.constexpr, |
| 74 | PAD_ID: tl.constexpr, |
| 75 | ): |
| 76 | """Return rank-local KV slots, or PAD_ID for positions not owned by this rank.""" |
| 77 | block_offsets = positions % (block_size * CP_SIZE) |
| 78 | if CP_SIZE == 1: |
| 79 | return block_numbers * block_size + block_offsets |
| 80 | is_local = block_offsets // CP_INTERLEAVE % CP_SIZE == cp_rank |
| 81 | rounds = block_offsets // (CP_INTERLEAVE * CP_SIZE) |
| 82 | remainder = block_offsets % CP_INTERLEAVE |
| 83 | local_offsets = rounds * CP_INTERLEAVE + remainder |
| 84 | return tl.where(is_local, block_numbers * block_size + local_offsets, PAD_ID) |