runtime/streamer.py

import torch


class GPUStreamer:
    """Manages direct GPU upload of token data.

    Uploads token tensors to the target device (XPU / CUDA / Metal)
    on a dedicated stream, overlapping transfer with decode compute.
    """
    def __init__(self, device: str = "cpu", stream: int = 0):
        self._device = torch.device(device)
        self._stream = stream

    def upload(self, tokens: torch.Tensor) -> torch.Tensor:
        return tokens.to(self._device, non_blocking=True)

    def upload_batch(self, tensors: list[torch.Tensor]) -> list[torch.Tensor]:
        return [t.to(self._device, non_blocking=True) for t in tensors]

    @property
    def device(self) -> torch.device:
        return self._device