Harry MellorandClaude Opus 5
8c1557a79c
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Signed-off-by: Harry Mellor <hej.mellor@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
36 lines
1.3 KiB
Python
36 lines
1.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Shared helpers for DeepSeek-V4's CPU-ported kernels."""
|
|
|
|
import torch
|
|
|
|
|
|
def map_local_to_global_slots_cpu(
|
|
local_indices: torch.Tensor,
|
|
req_idx: torch.Tensor,
|
|
block_table: torch.Tensor,
|
|
block_size: int,
|
|
) -> torch.Tensor:
|
|
"""Map per-request-local (token or compressed-token) indices to global
|
|
paged-cache slot ids via a block table, mirroring
|
|
``compute_global_topk_indices_and_lens``.
|
|
|
|
Args:
|
|
local_indices: [N, K], -1 sentinel for invalid entries.
|
|
req_idx: [N], row into ``block_table`` for each of the N rows.
|
|
block_table: [num_reqs, max_blocks_per_seq].
|
|
block_size: tokens (or compressed tokens) per physical block.
|
|
|
|
Returns:
|
|
[N, K] int64 global slot ids, -1 where ``local_indices`` was invalid.
|
|
|
|
"""
|
|
valid = local_indices >= 0
|
|
safe_local = local_indices.clamp(min=0).to(torch.int64)
|
|
block_pos = torch.div(safe_local, block_size, rounding_mode="floor")
|
|
offset = safe_local % block_size
|
|
req_expand = req_idx.to(torch.int64).unsqueeze(-1).expand_as(local_indices)
|
|
block_num = block_table[req_expand, block_pos].to(torch.int64)
|
|
slot = block_num * block_size + offset
|
|
return torch.where(valid, slot, torch.full_like(slot, -1))
|