Function _gumbel_noise

src/fused_mm_sampling/core.py:620–630

(seed, pid_v, pid_h, sample_idx, noise_offsets)

                )
                for peer_rank in tl.static_range(0, tp_world_size):
                    peer_base = tl.load(buffer_ptrs + peer_rank)
                    peer_maxs_ptr = peer_base.to(tl.pointer_type(tl.float32))
                    peer_maxs_idx_ptr = peer_base.to(tl.pointer_type(tl.int64))
                    tl.store(
                        peer_maxs_ptr + source_rank_base_offset,
                        gumbel_max,
                        mask=mask_h_out,
                    )
                    tl.store(
                        peer_maxs_idx_ptr + storage_offset_maxs_idx + source_rank_base_offset,
                        gumbel_max_idx_global,
                        mask=mask_h_out,
                    )
            else: