[feat]custom all reduce kernel

CLAUDE.md

@@ -580,7 +580,7 @@ Used by `flashinfer.trace` / `fi_trace`.
 | `FLASHINFER_AUTOTUNER_LOAD_FROM_FILE` | `0` | `flashinfer/autotuner/autotuner.py` | `1` loads previously serialized autotune results from disk instead of re-running the search. |
 | `FLASHINFER_DIST_AWARE_AUTOTUNE` | `0` | `flashinfer/fused_moe/da_config.py` | `1` enables experimental distribution-aware autotune and kernel dispatch (TRT-LLM MoE only). |
 | `FLASHINFER_DA_DISTRIBUTIONS` | built-in distribution catalog | `flashinfer/fused_moe/da_config.py` | Comma-separated training distributions used by the experimental TRT-LLM distribution-aware MoE autotuner. |
-| `FLASHINFER_AUTOTUNE_DIR` | unset | `flashinfer/mla/_sparse_mla_sm120.py` | Override the disk path for MLA AutoTuner cache files. Falls back to `FLASHINFER_WORKSPACE_DIR` when unset. |
+| `FLASHINFER_AUTOTUNE_DIR` | unset | `flashinfer/mla/_sparse_mla_sm120.py`, `flashinfer/comm/pcie_ipc_tuning.py` | Override the disk path for AutoTuner cache files (MLA, and the PCIe IPC all-reduce). Falls back to `FLASHINFER_WORKSPACE_DIR` when unset. |
 | `FLASHINFER_AUTOTUNE_TIMER` | unset (auto) | `flashinfer/autotuner/autotuner.py` | Selects the autotuner's per-tactic timer: `globaltimer` forces the GPU `%globaltimer` register, `cuda_event` forces `cudaEvent`, unset/anything-else auto-detects (uses `%globaltimer` only when Confidential Computing is detected). Under CC `cudaEventElapsedTime` is unreliable (can go negative), so the globaltimer path keeps tactic ranking stable. |
 | `FLASHINFER_CUTILE_AUTOTUNE_DISABLED` | `0` | `flashinfer/quantization/kernels/cutile/rope_quantize_fp8_cutile.py` | Non-zero skips exhaustive cuTile RoPE-FP8 tuning and uses the built-in token-count heuristic. |
 | `FLASHINFER_CONFIDENTIAL_COMPUTE` | unset | `flashinfer/utils.py` | Override NVIDIA Confidential Computing (CC) auto-detection used by `is_confidential_compute()` (which drives the autotuner timer above): `1` forces CC, `0` forces non-CC. Useful for CI or hosts without `pynvml`. |

csrc/pcie_ipc_all_reduce.cu

@@ -0,0 +1,206 @@
+/*
+ * Copyright (c) 2026 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <tvm/ffi/container/array.h>
+
+#include <cstdint>
+
+#include "flashinfer/comm/pcie_ipc_all_reduce.cuh"
+#include "tvm_ffi_utils.h"
+
+namespace fi = flashinfer::comm::pcie_ipc;
+
+using tvm::ffi::Array;
+
+// Opaque handle, matching the fptr_t convention used by the other custom
+// all-reduce bindings in this directory.
+using fptr_t = int64_t;
+static_assert(sizeof(void*) == sizeof(fptr_t));
+
+namespace {
+
+// Everything the launcher needs that does not change between calls. The
+// workspace itself is owned by the caller (see pcie_ipc_all_reduce.cuh).
+struct PcieIpcHandle {
+  fi::PeerViews views;
+  fi::WorkspaceLayout layout;
+  int rank;
+  int world_size;
+  int max_blocks;
+  int64_t max_numel;
+  int elem_size;
+};
+
+}  // namespace
+
+/*!
+ * \brief Bytes each rank must allocate and share over CUDA IPC.
+ *
+ * The caller passes the result to create_shared_buffer() and hands the
+ * resulting pointer array to pcie_ipc_init().
+ */
+int64_t pcie_ipc_workspace_size(int64_t world_size, int64_t max_numel, int64_t elem_size,
+                                int64_t max_blocks) {
+  TVM_FFI_ICHECK(world_size == 2 || world_size == 4 || world_size == 8)
+      << "pcie ipc all-reduce supports world_size 2, 4 or 8, got " << world_size;
+  TVM_FFI_ICHECK_GT(max_numel, 0) << "max_numel must be positive";
+  TVM_FFI_ICHECK_EQ(elem_size, 2)
+      << "only 2-byte dtypes (bfloat16, float16) are supported, got elem_size " << elem_size;
+  TVM_FFI_ICHECK_GT(max_blocks, 0) << "max_blocks must be positive";
+  return fi::workspace_size(static_cast<int>(world_size), max_numel, static_cast<int>(elem_size),
+                            static_cast<int>(max_blocks));
+}
+
+/*!
+ * \brief Bind an already-shared workspace and return an opaque handle.
+ *
+ * \param ipc_ptrs Peer pointers; entry i must address rank i's slab.
+ *
+ * The slab is zeroed here because the sentinel protocol reads +0.0 as "not yet
+ * written". The caller MUST barrier after this returns and before the first
+ * collective: a peer that starts pushing into this slab before we zero it
+ * would lose its payload.
+ */
+fptr_t pcie_ipc_init(Array<fptr_t> ipc_ptrs, int64_t rank, int64_t max_numel, int64_t elem_size,
+                     int64_t max_blocks) {
+  const int world_size = static_cast<int>(ipc_ptrs.size());
+  TVM_FFI_ICHECK(world_size == 2 || world_size == 4 || world_size == 8)
+      << "pcie ipc all-reduce supports world_size 2, 4 or 8, got " << world_size;
+  TVM_FFI_ICHECK(rank >= 0 && rank < world_size) << "rank " << rank << " out of range";
+  TVM_FFI_ICHECK_EQ(elem_size, 2)
+      << "only 2-byte dtypes (bfloat16, float16) are supported, got elem_size " << elem_size;
+  TVM_FFI_ICHECK_GT(max_blocks, 0) << "max_blocks must be positive";
+
+  int64_t ptrs[fi::kMaxWorldSize];
+  for (int i = 0; i < world_size; ++i) {
+    TVM_FFI_ICHECK_NE(ipc_ptrs[i], 0) << "ipc_ptrs[" << i << "] is null";
+    ptrs[i] = ipc_ptrs[i];
+  }
+
+  auto* handle = new PcieIpcHandle();
+  handle->layout = fi::compute_workspace_layout(world_size, max_numel, static_cast<int>(elem_size),
+                                                static_cast<int>(max_blocks));
+  handle->views = fi::make_peer_views(ptrs, world_size, static_cast<int>(rank), handle->layout);
+  handle->rank = static_cast<int>(rank);
+  handle->world_size = world_size;
+  handle->max_blocks = static_cast<int>(max_blocks);
+  handle->max_numel = max_numel;
+  handle->elem_size = static_cast<int>(elem_size);
+
+  cudaError_t err = cudaMemset(reinterpret_cast<void*>(ptrs[rank]), 0, handle->layout.total_bytes);
+  if (err != cudaSuccess) {
+    delete handle;
+    TVM_FFI_LOG_AND_THROW(RuntimeError)
+        << "failed to zero the pcie ipc workspace: " << cudaGetErrorString(err);
+  }
+  return reinterpret_cast<fptr_t>(handle);
+}
+
+void pcie_ipc_dispose(fptr_t handle) { delete reinterpret_cast<PcieIpcHandle*>(handle); }
+
+/*!
+ * \brief Out-of-place all-reduce over the shared workspace.
+ *
+ * \param blocks,threads,variant Launch configuration chosen by the caller;
+ *        \c variant is a fi::Variant and the (world_size, variant) pairs that
+ *        dispatch are listed in pcie_ipc_all_reduce.cuh.
+ */
+void pcie_ipc_all_reduce(fptr_t handle, TensorView inp, TensorView out, int64_t blocks,
+                         int64_t threads, int64_t variant, bool enable_pdl) {
+  auto* h = reinterpret_cast<PcieIpcHandle*>(handle);
+  ffi::CUDADeviceGuard device_guard(inp.device().device_id);
+  auto stream = get_stream(inp.device());
+
+  TVM_FFI_ICHECK(inp.IsContiguous() && out.IsContiguous()) << "input and output must be contiguous";
+  TVM_FFI_ICHECK_EQ(encode_dlpack_dtype(inp.dtype()), encode_dlpack_dtype(out.dtype()))
+      << "input and output dtype must match";
+  TVM_FFI_ICHECK_EQ(inp.numel(), out.numel()) << "input and output must have the same size";
+
+  const int64_t numel = inp.numel();
+  const int64_t elem_size = get_element_size(inp);
+  TVM_FFI_ICHECK_EQ(elem_size, h->elem_size)
+      << "dtype element size " << elem_size << " does not match the workspace's " << h->elem_size;
+  TVM_FFI_ICHECK_LE(static_cast<size_t>(numel * elem_size), h->layout.max_payload_bytes)
+      << "payload exceeds the workspace capacity";
+
+  const int64_t pack_elems = 16 / elem_size;
+  TVM_FFI_ICHECK_EQ(numel % pack_elems, 0)
+      << "numel must be divisible by the 16-byte pack width (" << pack_elems << ")";
+  TVM_FFI_ICHECK_EQ(h->max_numel % pack_elems, 0)
+      << "max_numel must be divisible by the 16-byte pack width";
+  TVM_FFI_ICHECK(blocks > 0 && blocks <= h->max_blocks)
+      << "blocks must be in (0, " << h->max_blocks << "], got " << blocks;
+  TVM_FFI_ICHECK(threads > 0 && threads <= 1024) << "threads must be in (0, 1024], got " << threads;
+  // Every barrier signals from threadIdx.x < world_size, so a narrower block
+  // leaves some peers with nobody to signal them and the collective hangs.
+  TVM_FFI_ICHECK_GE(threads, h->world_size)
+      << "threads must be at least world_size (" << h->world_size << "), got " << threads;
+  // Refused rather than silently wrong: ipc_topo_rsag8_block_param_kernel
+  // triggers launch completion before island_owner_ack and its barrier flag
+  // store, so a dependent kernel can start while this call's phase-4 state is
+  // still being written. Re-enabling needs that release moved past both stores,
+  // an audit of the other six, and an SM90 regression.
+  TVM_FFI_ICHECK(!enable_pdl)
+      << "enable_pdl is not supported yet: in the TP8 block kernel the launch-completion "
+         "trigger precedes the island ack and barrier flag stores";
+  TVM_FFI_ICHECK(variant >= 0 && variant < fi::kVariantCount)
+      << "variant must be in [0, " << fi::kVariantCount << "), got " << variant;
+  const auto algo = static_cast<fi::Variant>(variant);
+  // Reject rather than silently alias, so one configuration always names one
+  // kernel.
+  TVM_FFI_ICHECK(
+      !(h->world_size == 2 && algo != fi::Variant::kUnstaged && algo != fi::Variant::kStaged))
+      << "world_size 2 accepts only kUnstaged and kStaged, got variant " << variant;
+  TVM_FFI_ICHECK(!(algo == fi::Variant::kFlatStaged && h->world_size != 8))
+      << "kFlatStaged is world_size 8 only, got " << h->world_size;
+  // Only the block-partitioned TP8 kernel needs this: it derives its chunk
+  // from blockIdx.x & 3. Every other kernel uses flat grid-stride loops and
+  // accepts any block count.
+  if (h->world_size == 8 && algo == fi::Variant::kStaged) {
+    TVM_FFI_ICHECK_EQ(blocks % 4, 0)
+        << "the TP8 topology kernel requires blocks divisible by 4, got " << blocks;
+  }
+
+  cudaError_t err = cudaSuccess;
+  switch (encode_dlpack_dtype(out.dtype())) {
+    case bfloat16_code:
+      err = fi::all_reduce<nv_bfloat16>(static_cast<const nv_bfloat16*>(inp.data_ptr()),
+                                        static_cast<nv_bfloat16*>(out.data_ptr()), numel, h->views,
+                                        h->rank, h->world_size, h->max_blocks, h->max_numel,
+                                        static_cast<int>(blocks), static_cast<int>(threads), algo,
+                                        enable_pdl, stream);
+      break;
+    case float16_code:
+      err = fi::all_reduce<half>(
+          static_cast<const half*>(inp.data_ptr()), static_cast<half*>(out.data_ptr()), numel,
+          h->views, h->rank, h->world_size, h->max_blocks, h->max_numel, static_cast<int>(blocks),
+          static_cast<int>(threads), algo, enable_pdl, stream);
+      break;
+    default:
+      // The kernel templates carry a generic path, but only the two 2-byte
+      // dtypes are instantiated and measured.
+      TVM_FFI_LOG_AND_THROW(NotImplementedError)
+          << "pcie ipc all-reduce supports bfloat16 and float16 only";
+  }
+  if (err != cudaSuccess) {
+    TVM_FFI_LOG_AND_THROW(RuntimeError)
+        << "pcie ipc all-reduce launch failed: " << cudaGetErrorString(err);
+  }
+}
+
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(pcie_ipc_workspace_size, pcie_ipc_workspace_size);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(pcie_ipc_init, pcie_ipc_init);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(pcie_ipc_dispose, pcie_ipc_dispose);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(pcie_ipc_all_reduce, pcie_ipc_all_reduce);

docs/api/comm.rst

@@ -143,6 +143,55 @@ vLLM AllReduce
     vllm_get_graph_buffer_ipc_meta
     vllm_meta_size
 
+PCIe IPC AllReduce
+------------------
+
+Custom all-reduce for intra-node PCIe machines without NVLink. Admission is a
+capability check — world size, dtype, workspace capacity, and enough payload
+for every rank to own a share — and shapes it rejects fall back to the caller's
+own collective.
+
+.. code-block:: python
+
+    import flashinfer.comm as comm
+
+    # Collective: every rank builds the workspace with identical arguments.
+    # Size max_numel to the real workload -- an oversized one costs latency.
+    ws = comm.PcieIpcAllReduceWorkspace(group=group, max_numel=128 * 6144)
+
+    for x in activations:                     # same shapes, same order, all ranks
+        if ws.supports(x):
+            y = ws.all_reduce(x)              # out-of-place
+        else:
+            y = x.clone()
+            dist.all_reduce(y, group=group)   # unsupported shape: fall back
+
+    ws.destroy()                              # collective; all ranks together
+
+``supports()`` is a pure function of shape and dtype, so every rank reaches the
+same answer without agreeing on one at runtime. It says nothing about speed: a
+supported shape runs a seed launch configuration — one crossover keyed on the
+payload in bytes, no per-machine constants — and warns once per workspace until
+:meth:`~PcieIpcAllReduceWorkspace.tune` has measured the real one and persisted
+it, since the crossovers depend on the fabric.
+:func:`get_pcie_ipc_launch_config` exposes that seed for a
+``(world_size, numel, elem_size)`` triple, and returns ``None`` only for shapes
+the kernels cannot run.
+
+The kernels spin on peer flags with no timeout, so ranks that disagree on
+shape, dtype or call order hang rather than raise. One workspace serves one
+CUDA stream; use :meth:`~PcieIpcAllReduceWorkspace.rebind_stream` after
+ordering the two if a move is genuinely needed.
+
+.. autosummary::
+    :toctree: ../generated
+
+    PcieIpcAllReduceWorkspace
+    PcieIpcLaunchConfig
+    get_pcie_ipc_launch_config
+    probe_pcie_ipc_rank_topology
+    resolve_pcie_ipc_profile
+
 Ulysses Context-Parallel All-to-All
 -----------------------------------
 

flashinfer/aot.py

@@ -688,6 +688,7 @@ def gen_all_modules(
             gen_comm_alltoall_module,
             gen_dcp_alltoall_module,
             gen_moe_alltoall_module,
+            gen_pcie_ipc_comm_module,
             gen_trtllm_comm_module,
             gen_trtllm_mnnvl_comm_module,
             gen_vllm_comm_module,
@@ -713,6 +714,10 @@ def gen_all_modules(
             # SM90/SM12x users still get this via JIT.
             jit_specs.append(gen_dcp_alltoall_module())
         jit_specs.append(gen_vllm_comm_module())
+        # No architecture gate: the kernels use only plain PTX loads/stores
+        # and CUDA IPC, and target PCIe machines without NVLink, which is
+        # orthogonal to the SM version.
+        jit_specs.append(gen_pcie_ipc_comm_module())
 
     if add_misc:
         jit_specs += [
@@ -1163,7 +1168,7 @@ def main():
     parser.add_argument(
         "--add-comm",
         type=parse_bool,
-        help="Add communication kernels (trtllm_comm, vllm_comm)",
+        help="Add communication kernels (trtllm_comm, vllm_comm, pcie_ipc_comm)",
     )
     parser.add_argument(
         "--add-gemma",

flashinfer/comm/__init__.py

@@ -38,6 +38,24 @@
 from .vllm_ar import meta_size as vllm_meta_size
 from .vllm_ar import register_buffer as vllm_register_buffer
 from .vllm_ar import register_graph_buffers as vllm_register_graph_buffers
+from .pcie_ipc_ar import (
+    PcieIpcAllReduceWorkspace as PcieIpcAllReduceWorkspace,
+)
+from .pcie_ipc_ar import gen_pcie_ipc_comm_module as gen_pcie_ipc_comm_module
+from .pcie_ipc_ar import get_pcie_ipc_comm_module as get_pcie_ipc_comm_module
+from .pcie_ipc_policy import IpcLaunchConfig as PcieIpcLaunchConfig
+from .pcie_ipc_policy import IpcVariant as PcieIpcVariant
+from .pcie_ipc_tuning import PCIE_IPC_CUSTOM_OP as PCIE_IPC_CUSTOM_OP
+from .pcie_ipc_tuning import default_cache_path as pcie_ipc_default_cache_path
+from .pcie_ipc_policy import (
+    get_pcie_ipc_launch_config as get_pcie_ipc_launch_config,
+)
+from .pcie_ipc_topology import (
+    probe_pcie_ipc_rank_topology as probe_pcie_ipc_rank_topology,
+)
+from .pcie_ipc_topology import (
+    resolve_pcie_ipc_profile as resolve_pcie_ipc_profile,
+)
 from .ulysses import UlyssesCommunicator as UlyssesCommunicator
 from .ulysses import dispose_ulysses_a2a as dispose_ulysses_a2a
 from .ulysses import gen_ulysses_a2a_module as gen_ulysses_a2a_module