Method sample

src/fused_mm_sampling/core.py:648–649

(self, **kwargs)

@lru_cache(maxsize=1)
def set_torch_allocator_for_tma_descriptors_cached():
    """From https://triton-lang.org/main/python-api/generated/triton.language.make_tensor_descriptor.html"""
    # TMA descriptors require a global memory allocation
    triton.set_allocator(alloc_on_cuda)