Repository: vllm-project/vllm
File: vllm/entrypoints/pooling/score/utils.py (Python)

Code before:

def _should_use_gpu_for_maxsim() -> bool:
    return envs.VLLM_USE_GPU_FOR_POOLING_SCORE and current_platform.is_cuda()


def compute_maxsim_scores(
    q_embs: Sequence[torch.Tensor],
    d_embs: Sequence[torch.Tensor],
    max_batch_size: int = 16,
    max_score_matrix_elements: int = 16_000_000,
) -> list[torch.Tensor]:
    """Compute ColBERT MaxSim scores in padded mini-batches."""
    if len(q_embs) != len(d_embs):
        raise ValueError("q_embs and d_embs must have the same length")

    num_pairs = len(q_embs)
    if num_pairs == 0:
        return []

    for q_emb, d_emb in zip(q_embs, d_embs):
        if q_emb.ndim != 2 or d_emb.ndim != 2:
            raise ValueError("Each embedding tensor must be 2-D")
        if q_emb.shape[1] != d_emb.shape[1]:
            raise ValueError("Query and document embeddings must have same dim")

    compute_device = torch.device("cuda" if _should_use_gpu_for_maxsim() else "cpu")
    scores: list[torch.Tensor] = []
    start = 0
    while start < num_pairs:
        end = min(start + max_batch_size, num_pairs)
        max_q = max(int(x.shape[0]) for x in q_embs[start:end])
        max_d = max(int(x.shape[0]) for x in d_embs[start:end])

        # keep score matrix bounded to avoid oversized allocations.
        while (
            end - start > 1
            and (end - start) * max_q * max_d > max_score_matrix_elements
        ):
            end -= 1
            max_q = max(int(x.shape[0]) for x in q_embs[start:end])
            max_d = max(int(x.shape[0]) for x in d_embs[start:end])

        batch_q = q_embs[start:end]
        batch_d = d_embs[start:end]
        batch_size = end - start
        dim = int(batch_q[0].shape[1])
        dtype = batch_q[0].dtype

        q_batch = torch.zeros(
            (batch_size, max_q, dim), dtype=dtype, device=compute_device
        )

Review comment (suggestion):
```suggestion
    compute_device = torch.device(current_platform.device_type if _should_use_gpu_for_maxsim() else "cpu")
```

I don't see why we can't enable this for other platforms as well

Code after:
def _should_use_gpu_for_maxsim(use_gpu_for_pooling_score: bool) -> bool:
    return use_gpu_for_pooling_score and not current_platform.is_cpu()


def compute_maxsim_scores(
    q_embs: Sequence[torch.Tensor],
    d_embs: Sequence[torch.Tensor],
    max_batch_size: int = 16,
    max_score_matrix_elements: int = 16_000_000,
    use_gpu_for_pooling_score: bool = False,
) -> list[torch.Tensor]:
    """Compute ColBERT MaxSim scores in padded mini-batches."""
    if len(q_embs) != len(d_embs):
        raise ValueError("q_embs and d_embs must have the same length")

    num_pairs = len(q_embs)
    if num_pairs == 0:
        return []

    for q_emb, d_emb in zip(q_embs, d_embs):
        if q_emb.ndim != 2 or d_emb.ndim != 2:
            raise ValueError("Each embedding tensor must be 2-D")
        if q_emb.shape[1] != d_emb.shape[1]:
            raise ValueError("Query and document embeddings must have same dim")

    compute_device = torch.device(
        current_platform.device_type
        if _should_use_gpu_for_maxsim(use_gpu_for_pooling_score)
        else "cpu"
    )
    scores: list[torch.Tensor] = []
    start = 0
    while start < num_pairs:
        end = min(start + max_batch_size, num_pairs)
        max_q = max(int(x.shape[0]) for x in q_embs[start:end])
        max_d = max(int(x.shape[0]) for x in d_embs[start:end])

        # keep score matrix bounded to avoid oversized allocations.
        while (
            end - start > 1
            and (end - start) * max_q * max_d > max_score_matrix_elements
        ):
            end -= 1
            max_q = max(int(x.shape[0]) for x in q_embs[start:end])
            max_d = max(int(x.shape[0]) for x in d_embs[start:end])

        batch_q = q_embs[start:end]
        batch_d = d_embs[start:end]
        batch_size = end - start
        dim = int(batch_q[0].shape[1])
        dtype = batch_q[0].dtype