| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Shared helpers for DeepSeek-V4's CPU-ported kernels.""" |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | |
| 8 | def map_local_to_global_slots_cpu( |
| 9 | local_indices: torch.Tensor, |
| 10 | req_idx: torch.Tensor, |
| 11 | block_table: torch.Tensor, |
| 12 | block_size: int, |
| 13 | ) -> torch.Tensor: |
| 14 | """Map per-request-local (token or compressed-token) indices to global |
| 15 | paged-cache slot ids via a block table, mirroring |
| 16 | ``compute_global_topk_indices_and_lens``. |
| 17 | |
| 18 | Args: |
| 19 | local_indices: [N, K], -1 sentinel for invalid entries. |
| 20 | req_idx: [N], row into ``block_table`` for each of the N rows. |
| 21 | block_table: [num_reqs, max_blocks_per_seq]. |
| 22 | block_size: tokens (or compressed tokens) per physical block. |
| 23 | |
| 24 | Returns: |
| 25 | [N, K] int64 global slot ids, -1 where ``local_indices`` was invalid. |
| 26 | |
| 27 | """ |
| 28 | valid = local_indices >= 0 |
| 29 | safe_local = local_indices.clamp(min=0).to(torch.int64) |
| 30 | block_pos = torch.div(safe_local, block_size, rounding_mode="floor") |
| 31 | offset = safe_local % block_size |
| 32 | req_expand = req_idx.to(torch.int64).unsqueeze(-1).expand_as(local_indices) |
| 33 | block_num = block_table[req_expand, block_pos].to(torch.int64) |
| 34 | slot = block_num * block_size + offset |
| 35 | return torch.where(valid, slot, torch.full_like(slot, -1)) |