vllm.distributed.device_communicators.flashinfer_all_reduce

Classes:

• FlashInferAllReduce –

Functions:

• get_fi_ar_quant_workspace –

  Return the allreduce workspace for quant patterns, initializing if needed.

• get_fi_ar_workspace –

  Return the allreduce workspace for non-quant patterns, initializing if needed.

FlashInferAllReduce

Source code in vllm/distributed/device_communicators/flashinfer_all_reduce.py

class FlashInferAllReduce:
    def __init__(
        self,
        group: ProcessGroup,
        device: int | str | torch.device,
    ):
        self.disabled = True

        if not fi_ar_available:
            logger.info(
                "FlashInfer All Reduce is disabled because flashinfer is not available"
            )
            return

        if not current_platform.is_cuda():
            logger.info(
                "FlashInfer All Reduce is disabled because it requires CUDA platform"
            )
            return

        self.group = group
        self.world_size = dist.get_world_size(self.group)
        self.rank = dist.get_rank(self.group)
        self.device = device
        if self.world_size == 1:
            return

        default_max_size_mb = PassConfig.default_fi_allreduce_fusion_max_size_mb().get(
            self.world_size
        )
        if not default_max_size_mb:
            logger.warning(
                "FlashInfer All Reduce is disabled because it "
                "is not supported for world_size=%d.",
                self.world_size,
            )
            return
        backend, _ = _resolve_fi_ar_backend(self.group)
        if backend is None:
            logger.info(
                "FlashInfer All Reduce is disabled because no supported backend "
                "is available for this process group."
            )
            return
        tuned_max_size = _get_tuned_standalone_max_size(
            self.world_size,
            backend,
            self.group,
        )
        self.max_workspace_size = (
            tuned_max_size
            if tuned_max_size is not None
            else int(default_max_size_mb * MiB)
        )
        self.max_num_tokens = 0
        self.disabled = False

    def _ensure_workspace(self, hidden_dim: int, dtype: torch.dtype) -> bool:
        """Ensure the all reduce workspace is initialized."""
        if self.max_num_tokens == 0:
            element_size = torch.tensor([], dtype=dtype, device="cpu").element_size()
            self.max_num_tokens = self.max_workspace_size // (hidden_dim * element_size)
        workspace = get_fi_ar_workspace(
            world_size=self.world_size,
            rank=self.rank,
            max_token_num=self.max_num_tokens,
            hidden_dim=hidden_dim,
            dtype=dtype,
            group=self.group,
        )
        if workspace is None:
            self.disabled = True
            return False
        return True

    def should_use_fi_ar(self, input_tensor: torch.Tensor) -> bool:
        if self.disabled:
            return False

        if not input_tensor.is_cuda:
            return False

        if not input_tensor.is_contiguous():
            return False

        if len(input_tensor.shape) != 2:
            return False

        if input_tensor.dtype not in (
            torch.float16,
            torch.bfloat16,
            torch.float32,
        ):
            return False

        if input_tensor.nbytes > self.max_workspace_size:
            return False

        num_tokens, hidden_dim = input_tensor.shape
        if not self.max_num_tokens:
            element_size = torch.tensor([], dtype=input_tensor.dtype).element_size()
            self.max_num_tokens = self.max_workspace_size // (hidden_dim * element_size)

        if num_tokens > self.max_num_tokens:
            return False

        if not self._ensure_workspace(hidden_dim, input_tensor.dtype):
            return False

        workspace = get_fi_ar_workspace(
            world_size=self.world_size,
            rank=self.rank,
            max_token_num=self.max_num_tokens,
            hidden_dim=hidden_dim,
            dtype=input_tensor.dtype,
            group=self.group,
        )
        assert workspace is not None
        # The token bound above uses the full allocation budget, but mnnvl's
        # Lamport buffers rotate through three slots, so only ~1/3 is usable per
        # call. Asking the workspace directly to reject sizes the kernel can't fit.
        return workspace.is_buffer_size_sufficient(
            tp_size=self.world_size,
            num_tokens=num_tokens,
            hidden_dim=hidden_dim,
            dtype=input_tensor.dtype,
        )

    def all_reduce(self, input_tensor: torch.Tensor) -> torch.Tensor:
        num_tokens, hidden_dim = input_tensor.shape
        workspace = get_fi_ar_workspace(
            world_size=self.world_size,
            rank=self.rank,
            max_token_num=self.max_num_tokens,
            hidden_dim=hidden_dim,
            dtype=input_tensor.dtype,
            group=self.group,
        )
        return flashinfer_comm.allreduce_fusion(
            input=input_tensor,
            workspace=workspace,
            pattern=flashinfer_comm.AllReduceFusionPattern.kAllReduce,
            launch_with_pdl=True,
            trigger_completion_at_end=num_tokens > PDL_ADVANCE_LAUNCH_TOKENS,
        )

    def destroy(self):
        if not self.disabled:
            destroy_fi_ar_workspace()

_ensure_workspace(hidden_dim, dtype)

Ensure the all reduce workspace is initialized.

Source code in vllm/distributed/device_communicators/flashinfer_all_reduce.py

def _ensure_workspace(self, hidden_dim: int, dtype: torch.dtype) -> bool:
    """Ensure the all reduce workspace is initialized."""
    if self.max_num_tokens == 0:
        element_size = torch.tensor([], dtype=dtype, device="cpu").element_size()
        self.max_num_tokens = self.max_workspace_size // (hidden_dim * element_size)
    workspace = get_fi_ar_workspace(
        world_size=self.world_size,
        rank=self.rank,
        max_token_num=self.max_num_tokens,
        hidden_dim=hidden_dim,
        dtype=dtype,
        group=self.group,
    )
    if workspace is None:
        self.disabled = True
        return False
    return True

_create_workspace(backend, world_size, rank, max_token_num, hidden_dim, dtype, group)

Create a flashinfer allreduce workspace, returning None on failure.

Source code in vllm/distributed/device_communicators/flashinfer_all_reduce.py

def _create_workspace(
    backend: str,
    world_size: int,
    rank: int,
    max_token_num: int,
    hidden_dim: int,
    dtype: torch.dtype,
    group: ProcessGroup,
):
    """Create a flashinfer allreduce workspace, returning None on failure."""
    comm_backend = TorchDistBackend(group=group)
    rng_state = random.getstate()
    try:
        random.seed(int.from_bytes(os.urandom(16), byteorder="big"))
        # Creation may run lazily inside the first sync-checked forward (e.g.
        # with enforce_eager, which skips warmup).
        with gpu_sync_allowed(first_only=True):
            workspace = flashinfer_comm.create_allreduce_fusion_workspace(
                backend=backend,
                world_size=world_size,
                rank=rank,
                max_token_num=max_token_num,
                hidden_dim=hidden_dim,
                dtype=dtype,
                comm_backend=comm_backend,
                group=group,
            )
        if backend == "mnnvl" and not getattr(workspace, "mc_ptr", 0):
            workspace.destroy()
            logger.warning_once(
                "FlashInfer MNNVL multicast is unavailable on the current topology."
            )
            return None
    except Exception as e:
        if "multicast" in str(e).lower():
            logger.warning_once(
                "Failed to initialize FlashInfer All Reduce workspace: %s. "
                "This is expected on GPUs without NVSwitch (e.g., NVLink "
                "bridge-only or PCIe topologies).",
                e,
            )
        else:
            logger.warning_once(
                "Failed to initialize FlashInfer All Reduce workspace: %s.",
                e,
            )
        return None
    finally:
        random.setstate(rng_state)
    workspace_id = id(workspace)
    workspace_group = _fi_ar_workspace_groups.get(workspace_id)
    if workspace_group is not None and workspace_group is not group:
        raise RuntimeError(
            "FlashInfer returned an all-reduce workspace already associated "
            "with a different process group"
        )
    _fi_ar_workspace_groups[workspace_id] = group
    logger.debug(
        "Initialized FlashInfer All Reduce workspace: backend=%s, "
        "world_size=%d, rank=%d, max_token_num=%d, hidden_dim=%d, dtype=%s",
        backend,
        world_size,
        rank,
        max_token_num,
        hidden_dim,
        dtype,
    )
    return workspace

_resolve_fi_ar_backend(group)

Resolve the flashinfer allreduce backend for the current setup.

Returns:

• str | None –

  A (backend, allow_trtllm_fallback) tuple. backend is None

• bool –

  when automatic selection has no supported backend for the process

• tuple[str | None, bool] –

  group.

Source code in vllm/distributed/device_communicators/flashinfer_all_reduce.py

def _resolve_fi_ar_backend(group: ProcessGroup) -> tuple[str | None, bool]:
    """Resolve the flashinfer allreduce backend for the current setup.

    Returns:
        A ``(backend, allow_trtllm_fallback)`` tuple. ``backend`` is ``None``
        when automatic selection has no supported backend for the process
        group.

    """
    backend = envs.VLLM_FLASHINFER_ALLREDUCE_BACKEND
    if backend != "auto":
        logger.debug_once("Using flashinfer allreduce backend: %s", backend)
        return backend, False

    node_count = _node_count(group)
    if node_count == 1:
        # TRTLLM supports single-node allreduce on Hopper and Blackwell. Do not
        # create an MNNVL workspace when the process group is node-local.
        backend = "trtllm"
        allow_trtllm_fallback = False
    elif current_platform.has_device_capability(100):
        # MNNVL requires Blackwell-class GPUs and is the only supported
        # FlashInfer backend for a multi-node allreduce process group.
        backend = "mnnvl"
        allow_trtllm_fallback = False
    else:
        logger.debug_once(
            "FlashInfer allreduce fusion is disabled: multi-node process groups "
            "require an MNNVL-capable Blackwell system."
        )
        return None, False

    logger.debug_once("Auto-selected flashinfer allreduce backend: %s", backend)
    return backend, allow_trtllm_fallback

get_fi_ar_quant_workspace(world_size, rank, max_token_num, hidden_dim, dtype, group)

Return the allreduce workspace for quant patterns, initializing if needed.

Backend is controlled by VLLM_FLASHINFER_ALLREDUCE_BACKEND env var, matching non-quant fusion. With auto this prefers mnnvl and falls back to trtllm only on single-node topologies where mnnvl multicast is unavailable.

Source code in vllm/distributed/device_communicators/flashinfer_all_reduce.py

def get_fi_ar_quant_workspace(
    world_size: int,
    rank: int,
    max_token_num: int,
    hidden_dim: int,
    dtype: torch.dtype,
    group: ProcessGroup,
):
    """Return the allreduce workspace for quant patterns, initializing if needed.

    Backend is controlled by VLLM_FLASHINFER_ALLREDUCE_BACKEND env var, matching
    non-quant fusion. With ``auto`` this prefers mnnvl and falls back to trtllm
    only on single-node topologies where mnnvl multicast is unavailable.
    """
    global _fi_ar_quant_workspace, _fi_ar_quant_workspace_failed
    if _fi_ar_quant_workspace is not None or _fi_ar_quant_workspace_failed:
        return _fi_ar_quant_workspace

    backend, allow_trtllm_fallback = _resolve_fi_ar_backend(group)

    if backend is None:
        return None

    if backend == "trtllm" and _node_count(group) > 1:
        raise ValueError(
            "Flashinfer allreduce quantization fusion is not supported for "
            "multi-node allreduce with 'trtllm' backend. Please use 'mnnvl' "
            "backend instead."
        )

    # Reuse the non-quant workspace if it was already created with the same
    # backend.
    if _fi_ar_workspace is not None and _fi_ar_workspace.backend == backend:
        _fi_ar_quant_workspace = _fi_ar_workspace
        return _fi_ar_quant_workspace

    if (
        _fi_ar_workspace is not None
        and _fi_ar_workspace.backend == "trtllm"
        and allow_trtllm_fallback
        and backend != "trtllm"
    ):
        _fi_ar_quant_workspace = _fi_ar_workspace
        return _fi_ar_quant_workspace

    _fi_ar_quant_workspace = _create_workspace(
        backend, world_size, rank, max_token_num, hidden_dim, dtype, group
    )
    if _fi_ar_quant_workspace is None and allow_trtllm_fallback and backend != "trtllm":
        logger.warning_once(
            "FlashInfer mnnvl allreduce quantization fusion workspace unavailable "
            "(likely no NVSwitch multicast support); falling back to trtllm "
            "backend for single node."
        )
        backend = "trtllm"
        if _fi_ar_workspace is not None and _fi_ar_workspace.backend == backend:
            _fi_ar_quant_workspace = _fi_ar_workspace
        else:
            _fi_ar_quant_workspace = _create_workspace(
                backend, world_size, rank, max_token_num, hidden_dim, dtype, group
            )

    if _fi_ar_quant_workspace is not None:
        logger.info_once(
            "Initialized FlashInfer Allreduce norm quantization "
            f"fusion workspace with backend={backend}"
        )
    else:
        _fi_ar_quant_workspace_failed = True
        logger.error_once(
            "Failed to initialize FlashInfer Allreduce norm quantization "
            f"fusion workspace with backend={backend}"
        )

    return _fi_ar_quant_workspace

get_fi_ar_workspace(world_size, rank, max_token_num, hidden_dim, dtype, group)

Return the allreduce workspace for non-quant patterns, initializing if needed.

Used by AllReduceFusionPass (non-quant patterns) and FlashInferAllReduce for standalone allreduce. Backend is controlled by VLLM_FLASHINFER_ALLREDUCE_BACKEND env var.

Source code in vllm/distributed/device_communicators/flashinfer_all_reduce.py

def get_fi_ar_workspace(
    world_size: int,
    rank: int,
    max_token_num: int,
    hidden_dim: int,
    dtype: torch.dtype,
    group: ProcessGroup,
):
    """Return the allreduce workspace for non-quant patterns, initializing if needed.

    Used by AllReduceFusionPass (non-quant patterns) and FlashInferAllReduce
    for standalone allreduce. Backend is controlled by
    VLLM_FLASHINFER_ALLREDUCE_BACKEND env var.
    """
    global _fi_ar_workspace, _fi_ar_workspace_failed
    if _fi_ar_workspace is not None or _fi_ar_workspace_failed:
        return _fi_ar_workspace

    backend, allow_trtllm_fallback = _resolve_fi_ar_backend(group)

    if backend is None:
        return None

    if backend == "trtllm" and _node_count(group) > 1:
        raise ValueError(
            "Flashinfer allreduce is not supported for multi-node allreduce with "
            "'trtllm' backend. Please use 'mnnvl' backend instead."
        )

    if (
        envs.VLLM_ALLREDUCE_USE_FLASHINFER
        and (max_size := _get_tuned_standalone_max_size(world_size, backend, group))
        is not None
    ):
        element_size = torch.empty((), dtype=dtype, device="cpu").element_size()
        max_token_num = max(max_token_num, max_size // (hidden_dim * element_size))

    def _get_or_create(be: str):
        # Reuse the quant workspace if it was already created with the same backend
        if _fi_ar_quant_workspace is not None and _fi_ar_quant_workspace.backend == be:
            return _fi_ar_quant_workspace
        return _create_workspace(
            be, world_size, rank, max_token_num, hidden_dim, dtype, group
        )

    _fi_ar_workspace = _get_or_create(backend)
    if _fi_ar_workspace is None and allow_trtllm_fallback and backend != "trtllm":
        logger.warning_once(
            "FlashInfer mnnvl allreduce workspace unavailable (likely no NVSwitch "
            "multicast support); falling back to trtllm backend for single node."
        )
        backend = "trtllm"
        _fi_ar_workspace = _get_or_create(backend)

    if _fi_ar_workspace is not None:
        logger.info_once(
            "Initialized FlashInfer Allreduce norm fusion workspace "
            f"with backend={backend}"
        )
    else:
        _fi_ar_workspace_failed = True
        logger.error_once(
            "Failed to initialize FlashInfer Allreduce norm fusion workspace "
            f"with backend={backend}"
        )

    return _fi_ar_workspace