Skip to content

vllm.models.beam.nvidia.allreduce

CUDA JIT MNNVL collective with Beam's FP32 sandwich-norm epilogue.

Classes:

Functions:

BeamAllReduceWorkspace

Reuse FlashInfer's multicast allocation and Lamport generation protocol.

Source code in vllm/models/beam/nvidia/allreduce.py 
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
class BeamAllReduceWorkspace:
    """Reuse FlashInfer's multicast allocation and Lamport generation protocol."""

    @classmethod
    def create(
        cls,
        group: torch.distributed.ProcessGroup,
        hidden_size: int,
        max_tokens: int,
        dtype: torch.dtype,
        *,
        device_group: torch.distributed.ProcessGroup | None = None,
    ) -> "BeamAllReduceWorkspace | None":
        # The CUDA epilogue handles one packed load per thread.
        if (
            group.size() not in (2, 4, 8)
            or dtype != torch.bfloat16
            or hidden_size % 1024
            or not 1024 <= hidden_size <= 8192
            or not has_flashinfer()
        ):
            return None
        from flashinfer.comm.mnnvl import (
            TorchDistBackend,
            all_ranks_agree,
            is_multicast_supported,
        )

        device = torch.accelerator.current_device_index()
        supported = torch.cuda.get_device_capability(device)[0] in (9, 10)
        supported = supported and is_multicast_supported(device)
        if not all_ranks_agree(TorchDistBackend(group=group), supported):
            return None
        return cls(group, hidden_size, max_tokens, device_group=device_group)

    def __init__(
        self,
        group: torch.distributed.ProcessGroup,
        hidden_size: int,
        max_tokens: int,
        *,
        device_group: torch.distributed.ProcessGroup | None = None,
    ) -> None:
        from flashinfer.comm import create_allreduce_fusion_workspace
        from flashinfer.comm.mnnvl import TorchDistBackend

        self.world_size = group.size()
        self.rank = torch.distributed.get_rank(group)
        self.hidden_size = hidden_size
        self.max_tokens = max_tokens
        self.device_group = device_group
        self.all_gather: Callable[[torch.Tensor], torch.Tensor] | None = None
        # Scale one-shot cutoffs by the BF16 payload staged across all ranks.
        token_bytes = hidden_size * 2
        rank_token_bytes = self.world_size * token_bytes
        self.small_batch_max_tokens = 3200 * 1024 // rank_token_bytes
        self.ep_finalize_oneshot_max_tokens = 2560 * 1024 // rank_token_bytes
        # GB300 TP2/4, H=5120, top-k=8: one-shot helps at one token;
        # two-shot wins through 131072 tokens. Keep the unfused cutoff tunable
        # for other geometries instead of imposing an unmeasured crossover.
        self.moe_gather_oneshot_max_tokens = 1
        self.moe_gather_separate_norm_min_tokens: int | None = None
        # Two-shot reserves two rows per token; one-shot needs one per rank.
        workspace_tokens = max(
            max_tokens,
            min(max_tokens, self.small_batch_max_tokens) * self.world_size // 2,
        )
        self.workspace = create_allreduce_fusion_workspace(
            backend="mnnvl",
            world_size=self.world_size,
            rank=self.rank,
            max_token_num=workspace_tokens,
            hidden_dim=hidden_size,
            dtype=torch.bfloat16,
            comm_backend=TorchDistBackend(group=group),
            group=group,
        )
        self.module = _module()
        self.allocation_bytes = (
            self.workspace.handle.mcast_device_memory.get_allocation_size()
        )
        self.separate_norm_min_tokens = cdiv(
            10 * 1024**2 * self.world_size, token_bytes
        )
        self.ep_finalize_max_tokens = 640 * 1024**2 // token_bytes

beam_finalize_allreduce_norm_residual_norm(x, residual, post_weight, post_eps, next_weight, next_eps, scales, workspace, routed_scale, *, oneshot_max_tokens=0)

Finalize local experts before AR, retaining the local BF16 rounding.

Source code in vllm/models/beam/nvidia/allreduce.py 
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
def beam_finalize_allreduce_norm_residual_norm(
    x: UnfinalizedMoEOutput,
    residual: torch.Tensor,
    post_weight: torch.Tensor | None,
    post_eps: float | None,
    next_weight: torch.Tensor | None,
    next_eps: float | None,
    scales: tuple[float, float] | None,
    workspace: BeamAllReduceWorkspace,
    routed_scale: float,
    *,
    oneshot_max_tokens: int = 0,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """Finalize local experts before AR, retaining the local BF16 rounding."""
    m, n = residual.shape
    assert 0 <= oneshot_max_tokens <= workspace.small_batch_max_tokens
    for weight in (post_weight, next_weight):
        if weight is not None:
            assert weight.dtype == torch.bfloat16 and weight.is_contiguous()
            assert weight.shape == (n,)
    assert x.gemm2_permuted.dtype == torch.bfloat16
    assert x.gemm2_permuted.is_contiguous() and x.gemm2_permuted.shape[1] == n
    assert x.expert_weights.dtype in (torch.float32, torch.bfloat16)
    assert x.expanded_idx_to_permuted_idx.dtype == torch.int32
    assert x.expert_weights.is_contiguous()
    assert x.expanded_idx_to_permuted_idx.is_contiguous()
    assert x.expert_weights.shape == x.expanded_idx_to_permuted_idx.shape
    assert x.expert_weights.shape[0] == m
    return _cuda_allreduce_norm_residual_norm(
        residual,
        residual,
        post_weight,
        post_eps,
        next_weight,
        next_eps,
        scales,
        workspace,
        x,
        routed_scale,
        oneshot_max_tokens,
    )