Function alloc_fn

src/fused_mm_sampling/tl_matmul.py:164–165

(size: int, alignment: int, stream: Optional[int])

    # TMA descriptors require a global memory allocation
    def alloc_fn(size: int, alignment: int, stream: Optional[int]):
        return torch.empty(size, device="cuda", dtype=torch.int8)

    triton.set_allocator(alloc_fn)