PCIe IPC AllGather and ReduceScatter schedules

The public entry points are the workspaces, launch configurations, variant enums, and policy getters exported from flashinfer.comm. Modules whose names start with _ contain package-internal implementation details. The prefix is a Python API convention, not a different import mechanism. Applications should use the public exports; internal workspace, topology, JIT, and tuning helpers can change without becoming additional public APIs.

Python package layout

flashinfer.comm.pcie_ipc_collectives groups the AG/RS Python implementation and shared PCIe collective helpers:

• all_gather.py and reduce_scatter.py provide the workspace APIs; all_gather_policy.py and reduce_scatter_policy.py define their launch configurations, variants, and seed policies.

• _ag_rs_workspace.py, _ag_rs_tuning.py, _ag_rs_module.py, and _ag_rs_topology.py contain AG/RS workspace management, exact-shape tuning, JIT loading, and operator-specific topology admission.

• _topology.py, _lifecycle.py, and _constants.py provide shared GPU identity/NVML evidence, stream and resource lifecycle, and constants. The existing AllReduce modules reuse these helpers but retain their current locations in flashinfer.comm and their own policy and admission rules.

Applications continue to import the public types from flashinfer.comm. The general-purpose cuda_ipc.py utilities remain outside this package. CUDA headers, TVM-FFI bindings, JIT generators, and tests retain FlashInfer’s existing include/, csrc/, flashinfer/jit/, and tests/ layout.

What a variant describes

The existing enum names identify complete implementations. They combine three choices: the communication algorithm, the transport (SM instructions or CUDA copy engines), and the amount of work between synchronizations. These are different dimensions; in particular, COPY_ENGINE is a recursive-doubling implementation, and the two ReduceScatter ONE_PACK implementations also visit peer destinations in cyclic order. Enum names and values remain stable.

Collective / variant	Admitted world sizes	Schedule and intended tradeoff
AG FLAT_PUSH	4	Load a local pack once, push it directly to each peer, synchronize once, then materialize peer shards from local staging. Low synchronization overhead is useful for small shards.
AG RECURSIVE_DOUBLING	2, 4, 8	Exchange accumulated shards with one XOR peer per pass, using SM load/store. More pass barriers, with less traffic across the slow boundary on ordered 4+4 TP8.
AG COPY_ENGINE	Ordered 4+4 TP8	The same XOR 4, 2, 1 exchange order using contiguous asynchronous CUDA copies. Bit-reversed staging makes each accumulated region contiguous. Copy submission and synchronization add fixed overhead; larger payloads can benefit from the copy-engine transport.
RS FLAT_CYCLIC	2, 4	Send every destination shard directly to its owner in cyclic destination order, synchronize once, then reduce the local staging contributions.
RS FLAT_ONE_PACK	2, 4	Use flat owner-push, but process one pack per thread in each grid-sized batch. Keep the local contribution in FP32 registers across publication, then reduce and store that batch.
RS TOPOLOGY_CYCLIC	Ordered 4+4 TP8	Reduce paired owner shards inside each four-GPU island, exchange one partial with the opposite-island owner, then finish the reduction. One island barrier and one owner-pair barrier follow the two sending phases.
RS TOPOLOGY_ONE_PACK	Ordered 4+4 TP8	The same hierarchical algorithm in grid-sized batches, retaining the local island partial in FP32 registers across the cross-island exchange. Two barriers are required per batch.

CYCLIC describes destination traversal, not a ring forwarding algorithm. ONE_PACK means one live pack per thread, not one pack for the entire GPU. A pack is 16 bytes: eight BF16/FP16 values or four FP32 values. With B blocks and T threads per block, each batch covers up to 16 * B * T shard bytes. For a shard of S bytes, the flat one-pack kernel has ceil(S / (16 * B * T)) publication barriers per CTA, while the topology one-pack kernel has twice that many rendezvous per CTA. These are per-CTA protocol counts; they are not a single GPU-wide barrier count.

The extra barriers explain a possible large-message cost of batching, but do not determine a universal crossover. The two schedules also differ in register lifetime, instruction order, memory access order, and the launch geometry that the tuner selects.

Why use two islands at TP8?

Let S be one rank-local shard in bytes: an AG input or an RS output. Count directed peer payload bytes across all eight ranks, excluding local copies, flags, protocol overhead, and retransmissions.

In a flat schedule every rank sends four shards across the island boundary: 8 * 4S = 32S total cross-island payload. With the hierarchical schedules each rank sends only one shard or partial across that boundary: 8 * S = 8S. Total peer payload remains 7S per rank; communication shifts from the cross-island links to the island-local links.

• AG starts by exchanging one input shard with XOR peer 4. The later XOR 2 and 1 passes relay accumulated shards within each island.

• RS first combines four contributions in each island, then each owner exchanges one shard-sized partial with XOR peer 4.

This is a traffic model, not a prediction of a fourfold latency improvement. Latency also depends on barriers, available links, contention, local memory traffic, launch geometry, and transport. On RS, each transmitted partial uses the output dtype. BF16/FP16 therefore introduce an intermediate rounding boundary before the final FP32 accumulation; results need not be bitwise equal to NCCL or to the flat reduction order.

Topology admission versus mathematical correctness

The current optimized TP8 admission contract requires logical ranks 0..3 and 4..7 to be the two islands. UUID-based NVML evidence must show unique physical GPUs on one host, non-SYSTEM links within each island, and SYSTEM links across islands. Both directions of every pair must agree. Missing or conflicting evidence leaves the optimized paths unavailable.

That fixed order matches the kernels’ rank arithmetic and the deployment whose performance motivated these schedules. SYSTEM is a physical topology classification, not a condition in the algebra of a sum or a copy. This check does not prove that an admitted configuration is fastest, and it does not replace CUDA IPC setup. Other placements would require a separately supported schedule or a rank mapping that preserves the caller’s rank-major tensor layout. The present API does not automatically reorder ranks.

AG recursive doubling remains available when this proof fails. TP8 RS reports unsupported; callers can select their own fallback. At TP2/TP4 the flat RS paths do not require the 4+4 proof. supports() describes this implementation’s admission contract, not a performance comparison with NCCL.

Defaults and measured selection

Without a matching tuning entry, the policy returns a deterministic seed:

RS world size	Seed variant by rank-local output bytes
2	FLAT_ONE_PACK below 4 MiB; FLAT_CYCLIC at or above 4 MiB
4	FLAT_ONE_PACK
8, ordered 4+4	TOPOLOGY_ONE_PACK below 128 KiB; TOPOLOGY_CYCLIC at or above 128 KiB

These boundaries are heuristic seeds. The historical RTX 6000D H6144/BF16 grid selected TP2 one-pack at 3 MiB and cyclic at 6 MiB, and TP8 one-pack at 96 KiB and cyclic at 192 KiB. The current defaults lie within those sampled intervals, but no exact-boundary calibration or derivation is recorded. Those observations do not establish the best variant at exactly 4 MiB or 128 KiB, for another dtype, or on another PCIe fabric.

The later BF16 campaign at commit 80f5ef1a5116ba777d40599724c38ae6610a8659 exercised H4096/H6144 and powers-of-two row counts from 1 through 32768 on RTX 6000D. Every variant above was selected for at least one shape. Some selected RS variants disagreed with the defaults, including topology one-pack at TP8/H4096/16 rows (128 KiB). A selected implementation beating NCCL does not by itself establish its margin over another custom variant. A robust deletion or threshold change needs repeated direct comparisons of the affected variants with their own useful launch geometries.

tune() enumerates every variant legal for the actual workspace and shape, including the fixed copy-engine configuration, and checks correctness before timing. It searches useful block/thread combinations within workspace limits; the seed is included. The tuner measures eager launches and minimizes the maximum elapsed time across ranks. Its exact-shape cache also records dtype, world size, placement fingerprint, and workspace limits. This is not an exhaustive graph-mode performance oracle. Tuning can choose a variant on either side of the seed thresholds.

Validation coverage

The GPU tests explicitly execute each admitted variant for BF16, FP16, and FP32, including a partial final pack batch. RS additionally compares every variant against an FP32 reference with dtype-appropriate tolerances.

Each core test also tunes one small BF16 shape, checks the default call against the reference, and repeats the check with a fresh workspace using the saved cache. This is a functional smoke test, not a performance assertion.

The graph tests queue mixed variants, changing grids and shapes, and then replay captured sequences with changed inputs and independent output snapshots. On TP8, the live topology gate controls whether CE/topology variants are exercised. Set FLASHINFER_TEST_PCIE_IPC_ORDERED_4PLUS4=1 on a designated 4+4 test node to require that coverage; arbitrary CI machines are not assumed to have that placement. Performance conclusions require separate measurements.