Function _local_reduce

src/fused_mm_sampling/core.py:344–354

Reduce across V-tiles (dim=1) on this rank and adjust to global vocab indices.

(
    maxs: torch.Tensor,  # [num_samples, n_tiles, H]
    maxs_idx: torch.Tensor,  # [num_samples, n_tiles, H]
    vocab_start_index: int,
)

@torch.compile(fullgraph=True)
def _local_reduce(
    maxs: torch.Tensor,  # [num_samples, n_tiles, H]
    maxs_idx: torch.Tensor,  # [num_samples, n_tiles, H]
    vocab_start_index: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Reduce across V-tiles (dim=1) on this rank and adjust to global vocab indices."""
    idxs = maxs.max(dim=1).indices  # [num_samples, H]
    samples = maxs_idx.gather(1, idxs.unsqueeze(1)).squeeze(1)  # [num_samples, H]
    max_values = maxs.gather(1, idxs.unsqueeze(1)).squeeze(1)  # [num_samples, H]
    samples += vocab_start_index
    return samples.T.contiguous(), max_values.T.contiguous()  # [H, num_samples]


def clip(low, high, x):