Method prepare

src/fused_mm_sampling/core.py:645–646

(self)

    # decimal multipliers.
    return -tl.log(-tl.log(tl.rand(seed + sample_idx, noise_offsets)))


@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"""