flashinfer.comm.PcieIpcAllReduceWorkspace

class flashinfer.comm.PcieIpcAllReduceWorkspace(group: ProcessGroup, max_numel: int, dtype: dtype = torch.bfloat16, max_blocks: int = 128, profile: str | None = None, tune_batches: Sequence[int] = (1, 2, 4, 8, 16, 32, 64, 128), tune_cache: str | None = None)

Shared workspace for the PCIe IPC all-reduce.

Allocates one slab per rank, shares it over CUDA IPC, and binds it to the kernels. The workspace is sized once and cannot grow, so max_numel must cover the largest collective that will be issued; anything larger must fall back to another backend.

This is a collective, and an unusually strict one. The kernels spin on peer flags with no timeout and no metadata exchange, so every rank must issue the same sequence of calls, with the same shape, dtype and launch configuration, in the same order. A rank that skips a call, reorders two, or passes a different explicit config does not get an error – the group hangs, or worse, one rank reads a neighbour’s partial sums as if they were finished. launch_config() is a pure function of shape, dtype and the workspace’s own immutable attributes precisely so that every rank derives the same answer without having to agree on one at runtime; passing config explicitly moves that obligation to the caller.

One workspace serves one CUDA stream. Its epoch and arrival counters assume the calls sharing it are totally ordered, which stream order gives and concurrent streams do not; the second stream is rejected. Build a separate workspace per stream.

Size max_numel to the real workload rather than to a round number. The epoch double buffer places its two halves world_size * max_numel elements apart, so an oversized workspace spreads them further than the payload needs and costs measurable time at small batch. The multiplier is the world size, not 2 – rounding max_numel up by 4x at 8 ranks moves the halves 32x the payload apart.

Parameters:

  • group (ProcessGroup) – Process group whose ranks share the workspace. Every rank must build the workspace with identical arguments.

  • max_numel (int) – Largest element count that will be all-reduced.

  • dtype (torch.dtype) – bfloat16 or float16. Only the element size is binding, so one workspace serves both.

  • max_blocks (int) – Upper bound on the block count any launch may request. Sizes the barrier and epoch slots.

  • profile (str, optional) – Force the interconnect label ("rootcplx" or "pcieswitch") instead of probing for it. The label does not pick a kernel; it partitions the tune cache so two topologies do not read each other’s measurements. Probing is collective and runs before any allocation.

  • tune_cache (str, optional) – Where tuned configurations are read from at construction and written by tune(). Defaults to FLASHINFER_AUTOTUNE_DIR (or the workspace directory). Give the same path to both, or a tuned result will not be found by the next process.

  • rather (Launch configurations start from a seed default that is workable)

:param than fast (see pcie_ipc_policy). Tune once to: :param replace it with measurements from this machine; the result is persisted and: :param later processes pick it up when the workspace is built. Tuning never changes: :param which shapes are supported: :param only which kernel a supported shape runs.:

Examples

>>> ws = PcieIpcAllReduceWorkspace(group=tp_group, max_numel=max_tokens * hidden)
>>> if ws.supports(x):
...     out = ws.all_reduce(x)
>>> ws.destroy()

Tuning, once per machine:

>>> ws.tune([hidden])  # collective; every rank calls it
__init__(group: ProcessGroup, max_numel: int, dtype: dtype = torch.bfloat16, max_blocks: int = 128, profile: str | None = None, tune_batches: Sequence[int] = (1, 2, 4, 8, 16, 32, 64, 128), tune_cache: str | None = None) → None

Methods

__init__(group, max_numel[, dtype, ...])

all_reduce(inp, *[, out, config, enable_pdl])

Out-of-place all-reduce.

destroy()

Release the handle and the shared slab.

launch_config(inp)

Seed launch configuration for inp, or None if unsupported.

prepare(shapes, *[, dtype])

Resolve the launch configuration for each shape now.

rebind_stream()

Allow the next call to come from a different stream.

supports(inp)

Whether the kernels can run inp at all.

tune(hiddens, *[, dtype, cache, tune_group, ...])

Measure the launch configuration for every tuned shape.

tuned_launch_config(inp)

Launch configuration for inp, measured if one has been persisted.

Attributes

handle