Skip to content

vllm.model_executor.layers.pooler.tokwise.heads

_mean_pool_chunks(data, chunk_size)

Mean-pool complete contextual states, accumulating in FP32.

Source code in vllm/model_executor/layers/pooler/tokwise/heads.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
def _mean_pool_chunks(data: torch.Tensor, chunk_size: int) -> torch.Tensor:
    """Mean-pool complete contextual states, accumulating in FP32."""
    chunk_size = min(chunk_size, len(data))
    num_full, remainder = divmod(len(data), chunk_size)
    full_end = num_full * chunk_size
    full_chunks = data[:full_end].reshape(num_full, chunk_size, data.shape[-1])
    if not remainder:
        return full_chunks.mean(dim=1, dtype=torch.float32)

    return torch.cat(
        [
            full_chunks.mean(dim=1, dtype=torch.float32),
            data[full_end:].mean(dim=0, keepdim=True, dtype=torch.float32),
        ]
    )