Function _allgather_logits

src/fused_mm_sampling/core.py:77–81

All-gather local logits along the vocab dimension to reconstruct [H, V_global].

(
    logits: torch.Tensor,  # [H, V_local]
)

def _allgather_logits(
    logits: torch.Tensor,  # [H, V_local]
) -> torch.Tensor:
    """All-gather local logits along the vocab dimension to reconstruct [H, V_global]."""
    return all_gather_tensor(logits, gather_dim=1, group=dist.group.WORLD)


def apply_top_k_top_p(