Function alloc_on_cuda

src/fused_mm_sampling/core.py:640–641

(size: int, alignment: int, stream: int | None)

@triton.jit
def _gumbel_noise(seed, sample_idx, noise_offsets):
    # Use the seed as the sample-stream identifier and the offset as the
    # element identifier. Keeping these two dimensions separate avoids the
    # collisions caused by combining sample and tile IDs with arbitrary
    # decimal multipliers.
    return -tl.log(-tl.log(tl.rand(seed + sample_idx, noise_offsets)))