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]
)

    H, V = probs.shape  # noqa: N806
    q = torch.empty(num_samples, H, V, device=probs.device, dtype=probs.dtype)
    q.exponential_()
    return probs.unsqueeze(0).div(q).argmax(dim=-1).T  # [H, num_samples]


def _allgather_logits(
    logits: torch.Tensor,  # [H, V_local]
) -> torch.Tensor:
    """All-gather local logits along the vocab dimension to reconstruct [H, V_global]."""