# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from __future__ import annotations

import numpy as np

from vllm.v1.worker.block_table import BlockTable


def fill_block_table_slot_mapping_np(
    *,
    block_table: BlockTable,
    req_indices: np.ndarray,  # [T] int64
    slot_positions_np: np.ndarray,  # [T] int64
    total_num_scheduled_tokens: int,
    block_size: int,
) -> None:
    """Fill `block_table.slot_mapping_np` for a packed batch of tokens."""
    block_table_indices = (req_indices * block_table.max_num_blocks_per_req +
                           slot_positions_np // int(block_size))
    block_table_cpu = block_table.get_cpu_tensor()
    block_numbers = block_table_cpu.flatten()[block_table_indices].numpy()
    block_offsets = slot_positions_np % int(block_size)
    np.add(
        block_numbers * int(block_size),
        block_offsets,
        out=block_table.slot_mapping_np[:total_num_scheduled_tokens],
    )

