Method prepare

src/fused_mm_sampling/core.py:656–657

(self)

def alloc_on_cuda(size: int, alignment: int, stream: int | None):
    return torch.empty(size, device="cuda", dtype=torch.int8)


class Sampler(Protocol):
    def prepare(self) -> "Sampler":
        raise NotImplementedError()